agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BREAKING] Python: Checkpoint refactor: encode/decode, checkpoint format, etc (#3744)

Browse Source 
* WIP: Checkpoint refactor: encode/decode, checkpoint format, etc

* WIP: Remove workflow ID in checkpoints

* Refactor checkpointing

* Add get_latest tests

* Increase test coverage

* Fix formatting

* Fix unit tests

* Fix samples

* fix unit tests

* fix pipeline

* Copilot comments

* Fix tests

* Fix more tests

* Address comments part 1

* Address comments part 2

* Comments
This commit is contained in:
Tao Chen
2026-02-11 20:57:15 +00:00
committed by GitHub
parent a2a672b687
commit 7db6c4ab4e
43 changed files with 3335 additions and 2075 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/orchestrations/tests/test_magentic.py
+16 -14
View File
@@ -362,7 +362,7 @@ async def test_magentic_checkpoint_resume_round_trip():
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = checkpoints[-1]
@@ -605,8 +605,9 @@ async def test_agent_executor_invoke_with_assistants_client_messages():
async def _collect_checkpoints(
storage: InMemoryCheckpointStorage,
workflow_name: str,
) -> list[WorkflowCheckpoint]:
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=workflow_name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
return checkpoints
@@ -619,12 +620,13 @@ async def test_magentic_checkpoint_resume_inner_loop_superstep():
participants=[StubThreadAgent()], checkpoint_storage=storage, manager=InvokeOnceManager()
).build()
async for event in workflow.run("inner-loop task", stream=True):
if event.type == "output":
break
async for _ in workflow.run("inner-loop task", stream=True):
continue
checkpoints = await _collect_checkpoints(storage)
inner_loop_checkpoint = next(cp for cp in checkpoints if cp.metadata.get("superstep") == 1) # type: ignore[reportUnknownMemberType]
checkpoints = await _collect_checkpoints(storage, workflow.name)
# The first checkpoint is after the manager has run.
# The second checkpoint is after the participant has run.
inner_loop_checkpoint = checkpoints[1]
resumed = MagenticBuilder(
participants=[StubThreadAgent()], checkpoint_storage=storage, manager=InvokeOnceManager()
@@ -651,7 +653,7 @@ async def test_magentic_checkpoint_resume_from_saved_state():
if event.type == "output":
break
checkpoints = await _collect_checkpoints(storage)
checkpoints = await _collect_checkpoints(storage, workflow.name)
# Verify we can resume from the last saved checkpoint
resumed_state = checkpoints[-1] # Use the last checkpoint
@@ -688,7 +690,7 @@ async def test_magentic_checkpoint_resume_rejects_participant_renames():
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
checkpoints = await _collect_checkpoints(storage)
checkpoints = await _collect_checkpoints(storage, workflow.name)
target_checkpoint = checkpoints[-1]
renamed_workflow = MagenticBuilder(
@@ -772,7 +774,7 @@ async def test_magentic_checkpoint_runtime_only() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert len(checkpoints) > 0, "Runtime-only checkpointing should have created checkpoints"
@@ -806,8 +808,8 @@ async def test_magentic_checkpoint_runtime_overrides_buildtime() -> None:
assert baseline_output is not None
buildtime_checkpoints = await buildtime_storage.list_checkpoints()
runtime_checkpoints = await runtime_storage.list_checkpoints()
buildtime_checkpoints = await buildtime_storage.list_checkpoints(workflow_name=wf.name)
runtime_checkpoints = await runtime_storage.list_checkpoints(workflow_name=wf.name)
assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints"
assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden"
@@ -856,13 +858,13 @@ async def test_magentic_checkpoint_restore_no_duplicate_history():
break
# Get checkpoint
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert len(checkpoints) > 0, "Should have created checkpoints"
latest_checkpoint = checkpoints[-1]
# Load checkpoint and verify no duplicates in state
checkpoint_data = await storage.load_checkpoint(latest_checkpoint.checkpoint_id)
checkpoint_data = await storage.load(latest_checkpoint.checkpoint_id)
assert checkpoint_data is not None
# Check the magentic_context in the checkpoint