[BREAKING] Python: Refactor Checkpointing for runner and runner context (#1645)

* Refactor Checkpointing for runner and runner context

* exception

* Fix formatting

* Comments

* rename

* Add detailed doc string

python/packages/core/tests/workflow/test_checkpoint.py

@@ -20,9 +20,7 @@ def test_workflow_checkpoint_default_values():
     assert checkpoint.timestamp != ""
     assert checkpoint.messages == {}
     assert checkpoint.shared_state == {}
-    assert checkpoint.executor_states == {}
     assert checkpoint.iteration_count == 0
-    assert checkpoint.max_iterations == 100
     assert checkpoint.metadata == {}
     assert checkpoint.version == "1.0"
 
@@ -35,9 +33,7 @@ def test_workflow_checkpoint_custom_values():
         timestamp=custom_timestamp,
         messages={"executor1": [{"data": "test"}]},
         shared_state={"key": "value"},
-        executor_states={"executor1": {"state": "active"}},
         iteration_count=5,
-        max_iterations=50,
         metadata={"test": True},
         version="2.0",
     )
@@ -47,9 +43,7 @@ def test_workflow_checkpoint_custom_values():
     assert checkpoint.timestamp == custom_timestamp
     assert checkpoint.messages == {"executor1": [{"data": "test"}]}
     assert checkpoint.shared_state == {"key": "value"}
-    assert checkpoint.executor_states == {"executor1": {"state": "active"}}
     assert checkpoint.iteration_count == 5
-    assert checkpoint.max_iterations == 50
     assert checkpoint.metadata == {"test": True}
     assert checkpoint.version == "2.0"
 
@@ -290,7 +284,6 @@ async def test_file_checkpoint_storage_json_serialization():
             workflow_id="complex-workflow",
             messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
             shared_state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
-            executor_states={"executor1": {"state": "active", "config": {"timeout": 30, "retries": 3}}},
         )
 
         # Save and load
@@ -300,7 +293,6 @@ async def test_file_checkpoint_storage_json_serialization():
         assert loaded is not None
         assert loaded.messages == checkpoint.messages
         assert loaded.shared_state == checkpoint.shared_state
-        assert loaded.executor_states == checkpoint.executor_states
 
         # Verify the JSON file is properly formatted
         file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"

python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py

@@ -8,6 +8,7 @@ from typing import Any
 from agent_framework._workflows._checkpoint import WorkflowCheckpoint
 from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
 from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary
+from agent_framework._workflows._const import EXECUTOR_STATE_KEY
 from agent_framework._workflows._events import RequestInfoEvent, WorkflowEvent
 from agent_framework._workflows._request_info_executor import (
     PendingRequestDetails,
@@ -16,10 +17,7 @@ from agent_framework._workflows._request_info_executor import (
     RequestInfoMessage,
     RequestResponse,
 )
-from agent_framework._workflows._runner_context import (
-    Message,
-    WorkflowState,
-)
+from agent_framework._workflows._runner_context import Message
 from agent_framework._workflows._shared_state import SharedState
 from agent_framework._workflows._workflow_context import WorkflowContext
 
@@ -29,9 +27,6 @@ PENDING_STATE_KEY = RequestInfoExecutor._PENDING_SHARED_STATE_KEY  # pyright: ig
 class _StubRunnerContext:
     """Minimal runner context stub for exercising WorkflowContext helpers."""
 
-    def __init__(self, stored_state: dict[str, Any] | None = None) -> None:
-        self._state = stored_state or {}
-
     async def send_message(self, message: Message) -> None:  # pragma: no cover - unused in tests
         return None
 
@@ -53,31 +48,27 @@ class _StubRunnerContext:
     async def next_event(self) -> WorkflowEvent:  # pragma: no cover - unused
         raise RuntimeError("Not implemented in stub context")
 
-    async def get_executor_state(self, executor_id: str) -> dict[str, Any] | None:  # pragma: no cover - trivial
-        return self._state
-
-    async def set_executor_state(self, executor_id: str, state: dict[str, Any]) -> None:  # pragma: no cover - unused
-        self._state = state
-
     def has_checkpointing(self) -> bool:  # pragma: no cover - unused
         return False
 
     def set_workflow_id(self, workflow_id: str) -> None:  # pragma: no cover - unused
         pass
 
-    def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:  # pragma: no cover - unused
+    def reset_for_new_run(self) -> None:  # pragma: no cover - unused
         pass
 
-    async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:  # pragma: no cover - unused
+    async def create_checkpoint(
+        self,
+        shared_state: SharedState,
+        iteration_count: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> str:  # pragma: no cover - unused
         raise RuntimeError("Checkpointing not supported in stub context")
 
     async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:  # pragma: no cover - unused
         return None
 
-    async def get_workflow_state(self) -> WorkflowState:  # pragma: no cover - unused
-        return {}  # type: ignore[return-value]
-
-    async def set_workflow_state(self, state: WorkflowState) -> None:  # pragma: no cover - unused
+    async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:  # pragma: no cover - unused
         pass
 
     def set_streaming(self, streaming: bool) -> None:  # pragma: no cover - unused
@@ -120,8 +111,8 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None:
         },
     )
 
-    runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx)
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
+    await ctx.set_executor_state({PENDING_STATE_KEY: {request_id: snapshot}})
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -143,8 +134,8 @@ async def test_has_pending_request_detects_snapshot() -> None:
         },
     )
 
-    runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx)
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
+    await ctx.set_executor_state({PENDING_STATE_KEY: {request_id: snapshot}})
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -152,9 +143,8 @@ async def test_has_pending_request_detects_snapshot() -> None:
 
 
 async def test_has_pending_request_false_when_snapshot_absent() -> None:
-    shared_state = SharedState()
-    runner_ctx = _StubRunnerContext({"pending_requests": {}})
-    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], shared_state, runner_ctx)
+    ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
+    await ctx.set_executor_state({PENDING_STATE_KEY: {}})
 
     executor = RequestInfoExecutor(id="request_info")
 
@@ -196,7 +186,6 @@ def test_pending_requests_from_checkpoint_and_summary() -> None:
                 }
             }
         },
-        executor_states={},
         iteration_count=1,
     )
 
@@ -284,10 +273,13 @@ async def test_run_persists_pending_requests_in_runner_state() -> None:
     await executor.execute(approval, ctx.source_executor_ids, shared_state, runner_ctx)
 
     # Runner state should include both pending snapshot and serialized request events
-    assert PENDING_STATE_KEY in runner_ctx._state  # pyright: ignore[reportPrivateUsage]
-    assert approval.request_id in runner_ctx._state[PENDING_STATE_KEY]  # pyright: ignore[reportPrivateUsage]
+    assert await shared_state.has(EXECUTOR_STATE_KEY)
+    executor_state = await shared_state.get(EXECUTOR_STATE_KEY)
+    assert executor.id in executor_state
+    assert PENDING_STATE_KEY in executor_state[executor.id]
+    assert approval.request_id in executor_state[executor.id][PENDING_STATE_KEY]
 
     response_ctx: WorkflowContext[None] = WorkflowContext("request_info", ["source"], shared_state, runner_ctx)
     await executor.handle_response("approved", approval.request_id, response_ctx)  # type: ignore
 
-    assert runner_ctx._state[PENDING_STATE_KEY] == {}  # pyright: ignore[reportPrivateUsage]
+    assert executor_state[executor.id][PENDING_STATE_KEY] == {}

python/packages/core/tests/workflow/test_workflow.py

@@ -414,9 +414,7 @@ async def test_workflow_run_stream_from_checkpoint_with_external_storage(simple_
             workflow_id="test-workflow",
             messages={},
             shared_state={},
-            executor_states={},
             iteration_count=0,
-            max_iterations=100,
         )
         checkpoint_id = await storage.save_checkpoint(test_checkpoint)
 
@@ -451,9 +449,7 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
             workflow_id="test-workflow",
             messages={},
             shared_state={},
-            executor_states={},
             iteration_count=0,
-            max_iterations=100,
         )
         checkpoint_id = await storage.save_checkpoint(test_checkpoint)
 
@@ -484,9 +480,7 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executo
             workflow_id="test-workflow",
             messages={},
             shared_state={},
-            executor_states={},
             iteration_count=0,
-            max_iterations=100,
         )
         checkpoint_id = await storage.save_checkpoint(test_checkpoint)
 
@@ -525,7 +519,7 @@ class StateTrackingExecutor(Executor):
     """An executor that tracks state in shared state to test context reset behavior."""
 
     @handler
-    async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any, list]) -> None:
+    async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any, list[Any]]) -> None:
         """Handle the message and track it in shared state."""
         # Get existing messages from shared state
         try:

python/packages/core/tests/workflow/test_workflow_observability.py

@@ -6,7 +6,7 @@ import pytest
 from opentelemetry import trace
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
-from agent_framework import WorkflowBuilder
+from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder
 from agent_framework._workflows._executor import Executor, handler
 from agent_framework._workflows._runner_context import InProcRunnerContext, Message
 from agent_framework._workflows._shared_state import SharedState
@@ -426,7 +426,7 @@ async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExp
 @pytest.mark.parametrize("enable_otel", [False], indirect=True)
 async def test_message_trace_context_serialization(span_exporter: InMemorySpanExporter) -> None:
     """Test that message trace context is properly serialized/deserialized."""
-    ctx = InProcRunnerContext()
+    ctx = InProcRunnerContext(InMemoryCheckpointStorage())
 
     # Create message with trace context
     message = Message(
@@ -439,16 +439,18 @@ async def test_message_trace_context_serialization(span_exporter: InMemorySpanEx
 
     await ctx.send_message(message)
 
-    # Get context state (which serializes messages)
-    state = await ctx.get_workflow_state()
+    # Create a checkpoint that includes the message
+    checkpoint_id = await ctx.create_checkpoint(SharedState(), 0)
+    checkpoint = await ctx.load_checkpoint(checkpoint_id)
+    assert checkpoint is not None
 
     # Check serialized message includes trace context
-    serialized_msg = state["messages"]["source"][0]
+    serialized_msg = checkpoint.messages["source"][0]
     assert serialized_msg["trace_contexts"] == [{"traceparent": "00-trace-span-01"}]
     assert serialized_msg["source_span_ids"] == ["span123"]
 
     # Test deserialization
-    await ctx.set_workflow_state(state)
+    await ctx.apply_checkpoint(checkpoint)
     restored_messages = await ctx.drain_messages()
 
     restored_msg = list(restored_messages.values())[0][0]