Python: Add checkpoint save and restore hooks to executor (#2097)

* Add checkpoint hooks

* Deprecate get_executor_state and set_executor_state

* Fix tests and samples

* Add doc strings

* Add sample

* Fix import

* Address comments and fix tests

* Address comments

* conditional import

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

@@ -158,8 +158,8 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     assert thread_messages[1].text == "Initial response 1"
 
 
-async def test_agent_executor_snapshot_and_restore_state_directly() -> None:
-    """Test AgentExecutor's snapshot_state and restore_state methods directly."""
+async def test_agent_executor_save_and_restore_state_directly() -> None:
+    """Test AgentExecutor's on_checkpoint_save and on_checkpoint_restore methods directly."""
     # Create agent with thread containing messages
     agent = _CountingAgent(id="direct_test_agent", name="DirectTestAgent")
     thread = AgentThread(message_store=ChatMessageStore())
@@ -182,7 +182,7 @@ async def test_agent_executor_snapshot_and_restore_state_directly() -> None:
     executor._cache = list(cache_messages)  # type: ignore[reportPrivateUsage]
 
     # Snapshot the state
-    state = await executor.snapshot_state()  # type: ignore[reportUnknownMemberType]
+    state = await executor.on_checkpoint_save()
 
     # Verify snapshot contains both cache and thread
     assert "cache" in state
@@ -206,7 +206,7 @@ async def test_agent_executor_snapshot_and_restore_state_directly() -> None:
     assert len(initial_thread_msgs) == 0
 
     # Restore state
-    await new_executor.restore_state(state)  # type: ignore[reportUnknownMemberType]
+    await new_executor.on_checkpoint_restore(state)
 
     # Verify cache is restored
     restored_cache = new_executor._cache  # type: ignore[reportPrivateUsage]

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

@@ -288,57 +288,6 @@ def test_build_fails_without_participants():
         HandoffBuilder().build()
 
 
-async def test_multiple_runs_dont_leak_conversation():
-    """Verify that running the same workflow multiple times doesn't leak conversation history."""
-    triage = _RecordingAgent(name="triage", handoff_to="specialist")
-    specialist = _RecordingAgent(name="specialist")
-
-    workflow = (
-        HandoffBuilder(participants=[triage, specialist])
-        .set_coordinator("triage")
-        .with_termination_condition(lambda conv: sum(1 for m in conv if m.role == Role.USER) >= 2)
-        .build()
-    )
-
-    # First run
-    events = await _drain(workflow.run_stream("First run message"))
-    requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
-    assert requests
-    events = await _drain(workflow.send_responses_streaming({requests[-1].request_id: "Second message"}))
-    outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
-    assert outputs, "First run should emit output"
-
-    first_run_conversation = outputs[-1].data
-    assert isinstance(first_run_conversation, list)
-    first_run_conv_list = cast(list[ChatMessage], first_run_conversation)
-    first_run_user_messages = [msg for msg in first_run_conv_list if msg.role == Role.USER]
-    assert len(first_run_user_messages) == 2
-    assert any("First run message" in msg.text for msg in first_run_user_messages if msg.text)
-
-    # Second run - should start fresh, not include first run's messages
-    triage.calls.clear()
-    specialist.calls.clear()
-
-    events = await _drain(workflow.run_stream("Second run different message"))
-    requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
-    assert requests
-    events = await _drain(workflow.send_responses_streaming({requests[-1].request_id: "Another message"}))
-    outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
-    assert outputs, "Second run should emit output"
-
-    second_run_conversation = outputs[-1].data
-    assert isinstance(second_run_conversation, list)
-    second_run_conv_list = cast(list[ChatMessage], second_run_conversation)
-    second_run_user_messages = [msg for msg in second_run_conv_list if msg.role == Role.USER]
-    assert len(second_run_user_messages) == 2, (
-        "Second run should have exactly 2 user messages, not accumulate first run"
-    )
-    assert any("Second run different message" in msg.text for msg in second_run_user_messages if msg.text)
-    assert not any("First run message" in msg.text for msg in second_run_user_messages if msg.text), (
-        "Second run should NOT contain first run's messages"
-    )
-
-
 async def test_handoff_async_termination_condition() -> None:
     """Test that async termination conditions work correctly."""
     termination_call_count = 0
@@ -585,7 +534,7 @@ async def test_return_to_previous_state_serialization():
     coordinator._current_agent_id = "specialist_a"  # type: ignore[reportPrivateUsage]
 
     # Snapshot the state
-    state = coordinator.snapshot_state()
+    state = await coordinator.on_checkpoint_save()
 
     # Verify pattern metadata includes current_agent_id
     assert "metadata" in state
@@ -603,7 +552,7 @@ async def test_return_to_previous_state_serialization():
     )
 
     # Restore state
-    coordinator2.restore_state(state)
+    await coordinator2.on_checkpoint_restore(state)
 
     # Verify current_agent_id was restored
     assert coordinator2._current_agent_id == "specialist_a", "Current agent should be restored from checkpoint"  # type: ignore[reportPrivateUsage]

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

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import sys
 from collections.abc import AsyncIterable
 from dataclasses import dataclass
 from typing import Any, cast
@@ -42,6 +43,11 @@ from agent_framework._workflows._magentic import (  # type: ignore[reportPrivate
     _MagenticStartMessage,  # type: ignore
 )
 
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
 
 def test_magentic_start_message_from_string():
     msg = _MagenticStartMessage.from_string("Do the thing")
@@ -101,8 +107,9 @@ class FakeManager(MagenticManagerBase):
     next_speaker_name: str = "agentA"
     instruction_text: str = "Proceed with step 1"
 
-    def snapshot_state(self) -> dict[str, Any]:
-        state = super().snapshot_state()
+    @override
+    def on_checkpoint_save(self) -> dict[str, Any]:
+        state = super().on_checkpoint_save()
         if self.task_ledger is not None:
             state = dict(state)
             state["task_ledger"] = {
@@ -111,8 +118,9 @@ class FakeManager(MagenticManagerBase):
             }
         return state
 
-    def restore_state(self, state: dict[str, Any]) -> None:
-        super().restore_state(state)
+    @override
+    def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
+        super().on_checkpoint_restore(state)
         ledger_state = state.get("task_ledger")
         if isinstance(ledger_state, dict):
             ledger_dict = cast(dict[str, Any], ledger_state)
@@ -185,7 +193,6 @@ async def test_standard_manager_progress_ledger_and_fallback():
     assert ledger2.is_request_satisfied.answer is False
 
 
-@pytest.mark.skip(reason="Response handling refactored - responses no longer passed to run_stream()")
 async def test_magentic_workflow_plan_review_approval_to_completion():
     manager = FakeManager(max_round_count=10)
     wf = (
@@ -204,7 +211,7 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
 
     completed = False
     output: ChatMessage | None = None
-    async for ev in wf.run_stream(
+    async for ev in wf.send_responses_streaming(
         responses={req_event.request_id: MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE)}
     ):
         if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
@@ -218,7 +225,6 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
     assert isinstance(output, ChatMessage)
 
 
-@pytest.mark.skip(reason="Response handling refactored - responses no longer passed to run_stream()")
 async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds():
     class CountingManager(FakeManager):
         # Declare as a model field so assignment is allowed under Pydantic
@@ -250,7 +256,7 @@ async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds()
     # Reply APPROVE with comments (no edited text). Expect one replan and no second review round.
     saw_second_review = False
     completed = False
-    async for ev in wf.run_stream(
+    async for ev in wf.send_responses_streaming(
         responses={
             req_event.request_id: MagenticPlanReviewReply(
                 decision=MagenticPlanReviewDecision.APPROVE,
@@ -298,7 +304,6 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result():
     assert data.role == Role.ASSISTANT
 
 
-@pytest.mark.skip(reason="Response handling refactored - send_responses_streaming no longer exists")
 async def test_magentic_checkpoint_resume_round_trip():
     storage = InMemoryCheckpointStorage()
 
@@ -369,7 +374,7 @@ class _DummyExec(Executor):
         pass
 
 
-def test_magentic_agent_executor_snapshot_roundtrip():
+async def test_magentic_agent_executor_on_checkpoint_save_and_restore_roundtrip():
     backing_executor = _DummyExec("backing")
     agent_exec = MagenticAgentExecutor(backing_executor, "agentA")
     agent_exec._chat_history.extend([  # type: ignore[reportPrivateUsage]
@@ -377,10 +382,10 @@ def test_magentic_agent_executor_snapshot_roundtrip():
         ChatMessage(role=Role.ASSISTANT, text="world", author_name="agentA"),
     ])
 
-    state = agent_exec.snapshot_state()
+    state = await agent_exec.on_checkpoint_save()
 
     restored_executor = MagenticAgentExecutor(_DummyExec("backing2"), "agentA")
-    restored_executor.restore_state(state)
+    await restored_executor.on_checkpoint_restore(state)
 
     assert len(restored_executor._chat_history) == 2  # type: ignore[reportPrivateUsage]
     assert restored_executor._chat_history[0].text == "hello"  # type: ignore[reportPrivateUsage]

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

@@ -199,7 +199,10 @@ async def test_fan_out():
 
     # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent
     # executor_b will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore)
-    assert len(events) == 7
+    # Each superstep will emit also emit a WorkflowStartedEvent and WorkflowCompletedEvent
+    # This workflow will converge in 2 supersteps because executor_c will send one more message
+    # after executor_b completes
+    assert len(events) == 11
 
     assert events.get_final_state() == WorkflowRunState.IDLE
     outputs = events.get_outputs()
@@ -220,7 +223,9 @@ async def test_fan_out_multiple_completed_events():
 
     # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent
     # executor_b and executor_c will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore)
-    assert len(events) == 8
+    # Each superstep will emit also emit a WorkflowStartedEvent and WorkflowCompletedEvent
+    # This workflow will converge in 1 superstep because executor_a and executor_b will not send further messages
+    assert len(events) == 10
 
     # Multiple outputs are expected from both executors
     outputs = events.get_outputs()
@@ -246,7 +251,8 @@ async def test_fan_in():
 
     # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent
     # aggregator will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore)
-    assert len(events) == 9
+    # Each superstep will emit also emit a WorkflowStartedEvent and WorkflowCompletedEvent
+    assert len(events) == 13
 
     assert events.get_final_state() == WorkflowRunState.IDLE
     outputs = events.get_outputs()