From f2c5ee7c018d6a786f955036473ccb1efe82e869 Mon Sep 17 00:00:00 2001 From: vaibhav-patel Date: Sat, 20 Jun 2026 15:49:15 +0530 Subject: [PATCH 1/5] Python: add checkpointing support to AgentFrameworkWorkflow.run() in ag-ui The ag-ui AgentFrameworkWorkflow.run() previously accepted only a RunAgentInput payload and exposed no way to use the core workflow's checkpointing/state-persistence, unlike the core agent-framework workflow implementations. This left ag-ui workflows without resumable execution. Add optional checkpoint_storage and checkpoint_id keyword arguments to run(), threaded through run_workflow_stream() into the core Workflow.run(). This delegates to the existing core capability instead of reinventing it and keeps the public surface consistent with Workflow.run(): - checkpoint_storage enables checkpoint creation at each superstep boundary. - checkpoint_id resumes a run from a persisted checkpoint; incoming messages are forwarded only as request-info responses (never as a new start-executor message) to honor the core's message/checkpoint_id mutual exclusivity, and responses + checkpoint_id performs a restore-then-send in one call. Both can also be supplied via the input_data keys __ag_ui_checkpoint_storage and __ag_ui_checkpoint_id so the FastAPI endpoint (which calls run(input_data) positionally) can opt in without changing its call site; explicit keyword arguments take precedence. Checkpoint resume bypasses the AG-UI thread snapshot hydration early-returns so it always reaches the core restore path. Backward compatible: run(input_data) keeps working unchanged, and the non-checkpoint path still calls run_workflow_stream(input_data, workflow) with its original two-argument convention. Adds focused tests covering checkpoint creation, resume-from-checkpoint, input-data-keyed params, and the unchanged default path. Fixes #6632. --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 68 +++++++- .../agent_framework_ag_ui/_workflow_run.py | 58 ++++++- .../ag-ui/tests/ag_ui/test_workflow_agent.py | 147 +++++++++++++++++- 3 files changed, 259 insertions(+), 14 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index aa583856a65..7b41151afba 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -24,7 +24,7 @@ ToolCallResultEvent, ToolCallStartEvent, ) -from agent_framework import Workflow +from agent_framework import CheckpointStorage, Workflow from ._message_adapters import agui_messages_to_snapshot_format from ._run_common import ( @@ -45,6 +45,13 @@ WorkflowFactory = Callable[[str], Workflow] +# Input-payload keys used to surface workflow checkpointing through ``run(input_data)`` +# without changing the positional call convention used by the FastAPI endpoint. The +# corresponding ``run()`` keyword arguments take precedence over these when both are +# supplied. +_CHECKPOINT_ID_INPUT_KEY = "__ag_ui_checkpoint_id" +_CHECKPOINT_STORAGE_INPUT_KEY = "__ag_ui_checkpoint_storage" + def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]: """Convert AG-UI message event models to plain snapshot dictionaries.""" @@ -276,10 +283,38 @@ def clear_workflow_cache(self) -> None: """Drop all cached thread workflow instances.""" self._workflow_by_thread.clear() - async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: + async def run( + self, + input_data: dict[str, Any], + *, + checkpoint_storage: CheckpointStorage | None = None, + checkpoint_id: str | None = None, + ) -> AsyncGenerator[BaseEvent]: """Run the wrapped workflow and yield AG-UI events. + Args: + input_data: The AG-UI request payload (a ``RunAgentInput`` dump). + checkpoint_storage: Optional checkpoint storage to enable workflow + checkpointing for this run. When provided, the underlying core + workflow creates a checkpoint at the end of each superstep, matching + ``agent_framework.Workflow.run(checkpoint_storage=...)``. May also be + supplied via the ``input_data`` key ``__ag_ui_checkpoint_storage``; + the keyword argument takes precedence. + checkpoint_id: Optional checkpoint id to resume the workflow from. When + provided, execution restores the persisted workflow state instead of + starting a fresh turn, matching + ``agent_framework.Workflow.run(checkpoint_id=...)``. May also be + supplied via the ``input_data`` key ``__ag_ui_checkpoint_id``; the + keyword argument takes precedence. + Subclasses may override this to provide custom AG-UI streams. + + Note: + Checkpointing (the ``agent_framework`` workflow checkpoint mechanism) is + independent from AG-UI Thread Snapshot persistence (``snapshot_store``). + The two can be used together, but they persist different things: snapshots + capture replayable protocol output for a thread, while checkpoints capture + executor/runtime state for resumable execution. """ thread_id = self._thread_id_from_input(input_data) run_id = str(input_data.get("run_id") or input_data.get("runId") or uuid.uuid4()) @@ -288,7 +323,23 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: resume_payload = _extract_resume_payload(input_data) snapshot_store = self.snapshot_store - if snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None: + # Explicit keyword arguments win over values smuggled through input_data so the + # FastAPI endpoint (which calls ``run(input_data)`` positionally) can still opt + # into checkpointing without changing its call site. + if checkpoint_id is None: + checkpoint_id = cast(str | None, input_data.get(_CHECKPOINT_ID_INPUT_KEY)) + if checkpoint_storage is None: + checkpoint_storage = cast(CheckpointStorage | None, input_data.get(_CHECKPOINT_STORAGE_INPUT_KEY)) + + # A checkpoint resume legitimately carries no new messages; it must reach the + # core workflow's restore path rather than replaying a stored thread snapshot. + if ( + checkpoint_id is None + and snapshot_store is not None + and snapshot_scope is not None + and not raw_messages + and resume_payload is None + ): async for event in _hydrate_workflow_thread_snapshot( snapshot_store=snapshot_store, scope=snapshot_scope, @@ -346,8 +397,17 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: state_snapshot = make_json_safe(effective_state) if isinstance(state_snapshot, dict): snapshot_builder.state = cast(dict[str, Any], state_snapshot) + # Only forward checkpoint kwargs when checkpointing is requested so the + # non-checkpoint path keeps calling ``run_workflow_stream(input_data, workflow)`` + # exactly as before (preserves the established two-argument call convention). + stream_kwargs: dict[str, Any] = {} + if checkpoint_storage is not None: + stream_kwargs["checkpoint_storage"] = checkpoint_storage + if checkpoint_id is not None: + stream_kwargs["checkpoint_id"] = checkpoint_id + run_error_emitted = False - async for event in run_workflow_stream(input_data, workflow): + async for event in run_workflow_stream(input_data, workflow, **stream_kwargs): if snapshot_builder is not None: snapshot_builder.observe(event) if isinstance(event, RunErrorEvent): diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 44d571aef8d..1e3d4d0fff4 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -24,7 +24,15 @@ ToolCallEndEvent, ToolCallStartEvent, ) -from agent_framework import AgentResponse, AgentResponseUpdate, Content, Message, Workflow, WorkflowRunState +from agent_framework import ( + AgentResponse, + AgentResponseUpdate, + CheckpointStorage, + Content, + Message, + Workflow, + WorkflowRunState, +) from ._message_adapters import normalize_agui_input_messages from ._run_common import ( @@ -558,8 +566,24 @@ def _details_code(details: Any) -> str | None: async def run_workflow_stream( input_data: dict[str, Any], workflow: Workflow, + *, + checkpoint_storage: CheckpointStorage | None = None, + checkpoint_id: str | None = None, ) -> AsyncGenerator[BaseEvent]: - """Run a Workflow and emit AG-UI protocol events.""" + """Run a Workflow and emit AG-UI protocol events. + + Args: + input_data: Normalized AG-UI request payload (a ``RunAgentInput`` dump). + workflow: The core ``Workflow`` instance to execute. + checkpoint_storage: Optional checkpoint storage forwarded to the core + workflow. When provided, the workflow creates a checkpoint at the end + of each superstep, mirroring ``Workflow.run(checkpoint_storage=...)``. + checkpoint_id: Optional checkpoint id to resume from. When provided the run + restores the persisted workflow state instead of starting a fresh turn, + mirroring ``Workflow.run(checkpoint_id=...)``. Any incoming messages are + treated as request-info responses (or ignored) rather than a new + start-executor message, so resume stays consistent with the core API. + """ thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4()) run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4()) available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts") @@ -587,7 +611,11 @@ async def run_workflow_stream( if not responses and pending_before_run: responses.update(_single_pending_response_from_value(pending_before_run, _latest_user_text(messages))) - if not responses and pending_before_run: + # A checkpoint resume must always reach ``workflow.run(checkpoint_id=...)`` so the + # core restores persisted state and re-emits any pending requests from the + # checkpoint. ``pending_before_run`` reflects the live (pre-restore) instance, so + # short-circuiting on it here would skip the restore entirely. + if checkpoint_id is None and not responses and pending_before_run: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) for request_event in pending_before_run.values(): request_payload = _request_payload_from_request_event(request_event) @@ -604,7 +632,7 @@ async def run_workflow_stream( yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts) return - if not responses and not messages: + if checkpoint_id is None and not responses and not messages: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts) return @@ -640,11 +668,27 @@ def _drain_open_message() -> list[TextMessageEndEvent]: logger.debug("workflow.run() does not accept function_invocation_kwargs; dropping forwarded_props") fwd_kwargs = {} + # Forward checkpoint storage so the core workflow creates a checkpoint at the end + # of each superstep (parity with ``Workflow.run(checkpoint_storage=...)``). + checkpoint_kwargs: dict[str, Any] = {} + if checkpoint_storage is not None: + checkpoint_kwargs["checkpoint_storage"] = checkpoint_storage + try: - if responses: - event_stream = workflow.run(responses=responses, stream=True, **fwd_kwargs) + if checkpoint_id is not None: + # Resume from a checkpoint. ``message`` is mutually exclusive with + # ``checkpoint_id`` in the core API, so incoming messages are only + # forwarded as request-info responses (``responses``), never as a new + # start-executor message. ``responses`` + ``checkpoint_id`` performs a + # restore-then-send in a single call. + run_kwargs: dict[str, Any] = {"checkpoint_id": checkpoint_id, **checkpoint_kwargs} + if responses: + run_kwargs["responses"] = responses + event_stream = workflow.run(stream=True, **run_kwargs, **fwd_kwargs) + elif responses: + event_stream = workflow.run(responses=responses, stream=True, **checkpoint_kwargs, **fwd_kwargs) else: - event_stream = workflow.run(message=messages, stream=True, **fwd_kwargs) + event_stream = workflow.run(message=messages, stream=True, **checkpoint_kwargs, **fwd_kwargs) async for event in event_stream: event_type = getattr(event, "type", None) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 858d10370f0..7afbbd75efd 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -7,13 +7,25 @@ from typing import Any, cast import pytest -from agent_framework import Workflow, WorkflowBuilder, WorkflowContext, executor +from agent_framework import ( + InMemoryCheckpointStorage, + Workflow, + WorkflowBuilder, + WorkflowContext, + executor, + handler, +) +from agent_framework._workflows._executor import Executor from agent_framework_ag_ui import AgentFrameworkWorkflow -async def _run(agent: AgentFrameworkWorkflow, payload: dict[str, Any]) -> list[Any]: - return [event async for event in agent.run(payload)] +async def _run( + agent: AgentFrameworkWorkflow, + payload: dict[str, Any], + **run_kwargs: Any, +) -> list[Any]: + return [event async for event in agent.run(payload, **run_kwargs)] async def test_workflow_wrapper_rejects_workflow_and_factory_at_once() -> None: @@ -110,3 +122,132 @@ async def test_workflow_wrapper_factory_return_type_is_validated() -> None: with pytest.raises(TypeError, match="workflow_factory must return a Workflow instance"): _ = [event async for event in agent.run({"thread_id": "thread-a", "messages": []})] + + +# region checkpointing + + +class _StartExecutor(Executor): + @handler + async def run(self, message: Any, ctx: WorkflowContext[str]) -> None: + del message + await ctx.send_message("hello", target_id="middle") + + +class _MiddleExecutor(Executor): + @handler + async def process(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(f"{message}-processed", target_id="finish") + + +class _FinishExecutor(Executor): + @handler + async def finish(self, message: str, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(f"{message}-done") + + +def _build_multi_superstep_workflow(storage: InMemoryCheckpointStorage | None = None) -> Workflow: + """Build a start -> middle -> finish workflow that creates a checkpoint per superstep.""" + start = _StartExecutor(id="start") + middle = _MiddleExecutor(id="middle") + finish = _FinishExecutor(id="finish") + builder = WorkflowBuilder(max_iterations=10, start_executor=start) + if storage is not None: + builder = WorkflowBuilder(max_iterations=10, start_executor=start, checkpoint_storage=storage) + return builder.add_edge(start, middle).add_edge(middle, finish).build() + + +async def test_workflow_run_creates_checkpoints_via_storage_kwarg() -> None: + """Passing checkpoint_storage to run() should create workflow checkpoints (parity with core).""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow() + agent = AgentFrameworkWorkflow(workflow=workflow) + + events = await _run( + agent, + {"thread_id": "thread-cp", "messages": [{"role": "user", "content": "start"}]}, + checkpoint_storage=storage, + ) + + event_types = [event.type for event in events] + assert "RUN_STARTED" in event_types + assert "RUN_FINISHED" in event_types + assert "RUN_ERROR" not in event_types + + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + # One checkpoint per superstep boundary: at least the initial superstep plus follow-ups. + assert len(checkpoints) >= 2 + + +async def test_workflow_run_resumes_from_checkpoint_id() -> None: + """run(checkpoint_id=...) should restore persisted state and finish the workflow.""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow(storage) + agent = AgentFrameworkWorkflow(workflow=workflow) + + # First run: execute to completion while checkpoints are written. + first_events = await _run( + agent, + {"thread_id": "thread-cp", "messages": [{"role": "user", "content": "start"}]}, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the run to create at least one checkpoint" + # Resume from the earliest checkpoint so middle -> finish replays and re-produces output. + resume_checkpoint_id = checkpoints[0].checkpoint_id + + # Resume on the same thread (same underlying workflow instance) from the checkpoint. + resumed_events = await _run( + agent, + {"thread_id": "thread-cp", "messages": []}, + checkpoint_id=resume_checkpoint_id, + checkpoint_storage=storage, + ) + + resumed_types = [event.type for event in resumed_events] + assert "RUN_STARTED" in resumed_types + assert "RUN_FINISHED" in resumed_types + assert "RUN_ERROR" not in resumed_types + + # The resumed run should reproduce the final assistant output ("hello-processed-done"). + resumed_text = "".join( + getattr(event, "delta", "") for event in resumed_events if event.type == "TEXT_MESSAGE_CONTENT" + ) + assert "done" in resumed_text + + +async def test_workflow_run_reads_checkpoint_params_from_input_data() -> None: + """Checkpoint params smuggled through input_data should be honored (endpoint call convention).""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow() + agent = AgentFrameworkWorkflow(workflow=workflow) + + events = await _run( + agent, + { + "thread_id": "thread-cp-input", + "messages": [{"role": "user", "content": "start"}], + "__ag_ui_checkpoint_storage": storage, + }, + ) + + assert "RUN_ERROR" not in [event.type for event in events] + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + assert len(checkpoints) >= 1 + + +async def test_workflow_run_without_checkpointing_is_unchanged() -> None: + """Existing run(input_data) calls keep working unchanged when no checkpoint args are given.""" + workflow = _build_multi_superstep_workflow() + agent = AgentFrameworkWorkflow(workflow=workflow) + + events = await _run(agent, {"thread_id": "thread-plain", "messages": [{"role": "user", "content": "start"}]}) + + event_types = [event.type for event in events] + assert "RUN_STARTED" in event_types + assert "RUN_FINISHED" in event_types + assert "RUN_ERROR" not in event_types From b7f33842d040d9de4f2c21bb6dbcf145a32951b3 Mon Sep 17 00:00:00 2001 From: vaibhav-patel Date: Mon, 22 Jun 2026 14:01:35 +0530 Subject: [PATCH 2/5] Import Executor from the public agent_framework API in ag-ui workflow test --- python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 7afbbd75efd..9a66dea32a5 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -8,6 +8,7 @@ import pytest from agent_framework import ( + Executor, InMemoryCheckpointStorage, Workflow, WorkflowBuilder, @@ -15,8 +16,6 @@ executor, handler, ) -from agent_framework._workflows._executor import Executor - from agent_framework_ag_ui import AgentFrameworkWorkflow From 5b6dc4cea905b0c5ca98257ca9989908b09dcedc Mon Sep 17 00:00:00 2001 From: vaibhav-patel Date: Thu, 9 Jul 2026 12:04:41 +0400 Subject: [PATCH 3/5] Fix ag-ui checkpoint resume: preserve thread snapshot, coerce resume responses; fix CI lint/typing A checkpoint-only resume no longer clobbers the stored AG-UI thread snapshot: the snapshot builder is seeded with the prior stored history so the saved snapshot keeps the earlier replayable transcript plus the newly produced output. Resume responses are now coerced against the post-restore pending requests on a checkpoint restore, so a JSON function_approval_response resumes through AG-UI after a cold restore instead of failing with a response-type mismatch. Also update the test-double workflow run() overrides to match the new keyword-only parent signature and re-sort the workflow test imports so ruff and the typing checkers pass. --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 13 ++- .../agent_framework_ag_ui/_workflow_run.py | 41 +++++++++ .../ag-ui/tests/ag_ui/test_endpoint.py | 12 +-- .../ag-ui/tests/ag_ui/test_workflow_agent.py | 75 +++++++++++++++ .../ag-ui/tests/ag_ui/test_workflow_run.py | 91 +++++++++++++++++++ 5 files changed, 223 insertions(+), 9 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 4a529a41a67..fe777ce0e12 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -396,9 +396,16 @@ async def run( workflow = self._resolve_workflow(thread_id, snapshot_scope) builder_seed_messages = raw_messages - if resume_payload is not None and stored_snapshot is not None: - # Resume requests carry only the synthesized interrupt response, so seed - # the builder with stored history to avoid persisting a truncated thread. + # ``raw_messages`` lacks the prior transcript in two cases: a resume request + # carries only the synthesized interrupt response, and a checkpoint-only resume + # carries no new messages at all (so ``_reconstruct_messages_from_thread_snapshot`` + # returns the empty incoming list rather than folding in stored history). In + # both cases seed the builder with the stored history so the snapshot saved + # after the run preserves the earlier replayable messages instead of dropping + # them for just the newly produced output. + if stored_snapshot is not None and ( + resume_payload is not None or (checkpoint_id is not None and not raw_messages) + ): builder_seed_messages = [ copy.deepcopy(message) for message in stored_snapshot.messages ] + builder_seed_messages diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 304dad03ecd..cb37c637e48 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -150,6 +150,32 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: return {} +async def _restore_checkpoint_for_coercion( + workflow: Workflow, + checkpoint_id: str, + checkpoint_storage: CheckpointStorage | None, +) -> None: + """Restore a checkpoint into the live workflow so pending requests become visible. + + Resume responses are coerced against the workflow's *current* pending requests. + On a cold checkpoint restore those requests do not exist yet on the live instance, + so this best-effort restore hydrates them first. It is idempotent: the subsequent + ``workflow.run(checkpoint_id=...)`` restores again from the same checkpoint. Any + failure is swallowed here and surfaced by that core run instead. + """ + runner = getattr(workflow, "_runner", None) + restore_from_checkpoint = getattr(runner, "restore_from_checkpoint", None) + if restore_from_checkpoint is None: + return + try: + await restore_from_checkpoint(checkpoint_id, checkpoint_storage) + except Exception: # pragma: no cover - defensive; the core run re-raises the real error + logger.warning( + "Pre-restore for resume-response coercion failed; the core run will surface any error.", + exc_info=True, + ) + + def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None: """Build AG-UI interrupt payload from a workflow request_info event.""" request_payload = _request_payload_from_request_event(request_event) @@ -786,6 +812,21 @@ async def run_workflow_stream( last_assistant_text: str | None = None resume_payload = _extract_resume_payload(input_data) + + # A checkpoint resume that carries an explicit resume payload targets the requests + # that were pending when the checkpoint was written; those only reappear on the live + # instance once the checkpoint is restored. Restore up front so ``pending_before_run`` + # (and the resume contract + coercion below) see the post-restore pending set, exactly + # as the non-checkpoint path sees the live pending requests. Without this a raw JSON + # resume payload (e.g. a ``function_approval_response`` dict) would reach core + # uncoerced and be rejected with a response-type mismatch. Only pre-restore when a + # resume payload is present so a pure checkpoint restore still surfaces its pending + # interrupts instead of tripping the "resume required" contract. The + # ``workflow.run(checkpoint_id=...)`` call below restores again from the same + # checkpoint, so this pre-restore is idempotent. + if checkpoint_id is not None and resume_payload is not None: + await _restore_checkpoint_for_coercion(workflow, checkpoint_id, checkpoint_storage) + pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) resume_entries: list[dict[str, Any]] = [] diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index a7067594f14..48d43bca6d8 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -2243,8 +2243,8 @@ async def test_endpoint_streaming_error_emits_run_error_event(): """Streaming exceptions should emit RUN_ERROR instead of terminating silently.""" class FailingStreamWorkflow(AgentFrameworkWorkflow): - async def run(self, input_data: dict[str, Any]): - del input_data + async def run(self, input_data: dict[str, Any], **kwargs: Any): + del input_data, kwargs yield RunStartedEvent(run_id="run-1", thread_id="thread-1") raise RuntimeError("stream exploded") @@ -2997,8 +2997,8 @@ async def test_endpoint_encoding_failure_emits_run_error(): from unittest.mock import patch class SimpleWorkflow(AgentFrameworkWorkflow): - async def run(self, input_data: dict[str, Any]): - del input_data + async def run(self, input_data: dict[str, Any], **kwargs: Any): + del input_data, kwargs yield RunStartedEvent(run_id="run-1", thread_id="thread-1") app = FastAPI() @@ -3020,8 +3020,8 @@ async def test_endpoint_double_encoding_failure_terminates(): from unittest.mock import patch class SimpleWorkflow(AgentFrameworkWorkflow): - async def run(self, input_data: dict[str, Any]): - del input_data + async def run(self, input_data: dict[str, Any], **kwargs: Any): + del input_data, kwargs yield RunStartedEvent(run_id="run-1", thread_id="thread-1") app = FastAPI() diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 5dec793bcf0..fd8c8f2f750 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -16,6 +16,7 @@ executor, handler, ) + from agent_framework_ag_ui import AgentFrameworkWorkflow @@ -261,3 +262,77 @@ async def test_workflow_run_without_checkpointing_is_unchanged() -> None: assert "RUN_STARTED" in event_types assert "RUN_FINISHED" in event_types assert "RUN_ERROR" not in event_types + + +async def test_workflow_checkpoint_only_resume_preserves_thread_snapshot() -> None: + """A checkpoint-only resume must keep the prior stored thread snapshot, not truncate it. + + Regression test for a checkpoint-only resume (no new messages) silently replacing + the stored AG-UI Thread Snapshot with just the newly produced output, dropping the + earlier replayable transcript. + """ + from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore + from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot + + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow(storage) + store = InMemoryAGUIThreadSnapshotStore() + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store) + + # Prime the workflow so a checkpoint exists to resume from. + first_events = await _run( + agent, + { + "thread_id": "thread-cp-snap", + "run_id": "run-1", + "messages": [{"id": "user-1", "role": "user", "content": "First question"}], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the primed run to create at least one checkpoint" + # Resume from the earliest checkpoint so middle -> finish replays and re-produces output. + resume_checkpoint_id = checkpoints[0].checkpoint_id + + # Stand in for a richer stored transcript: two prior replayable messages that a + # checkpoint-only resume must preserve alongside the resumed output. + await store.save( + scope="tenant-a", + thread_id="thread-cp-snap", + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "First question"}, + {"id": "assistant-1", "role": "assistant", "content": "Earlier reply"}, + ], + state=None, + interrupt=None, + ), + ) + + # Checkpoint-only resume: no new messages, resume from the checkpoint. + resumed_events = await _run( + agent, + { + "thread_id": "thread-cp-snap", + "run_id": "run-2", + "messages": [], + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + checkpoint_id=resume_checkpoint_id, + checkpoint_storage=storage, + ) + assert "RUN_ERROR" not in [event.type for event in resumed_events] + + snapshot = await store.get(scope="tenant-a", thread_id="thread-cp-snap") + assert snapshot is not None + contents = [message.get("content") for message in snapshot.messages] + # Prior transcript preserved... + assert "First question" in contents + assert "Earlier reply" in contents + # ...plus the newly produced output from the resumed run. + assert any(isinstance(content, str) and "done" in content for content in contents) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index 0eb50938169..c6d59fe9458 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -13,6 +13,7 @@ AgentResponseUpdate, Content, Executor, + InMemoryCheckpointStorage, Message, WorkflowBuilder, WorkflowContext, @@ -325,6 +326,96 @@ async def handle_approval(self, original_request: Content, response: Content, ct assert any("approved" in delta for delta in text_deltas) +async def test_workflow_run_resume_content_response_after_checkpoint_restore() -> None: + """A JSON function_approval_response resumes correctly through a cold checkpoint restore. + + Regression test: on a checkpoint restore the pending approval request only reappears + after the checkpoint is restored, so resume responses must be coerced against the + post-restore pending set. Without that, the raw JSON payload reaches core uncoerced + and is rejected with "Response type mismatch ... expected Content, got dict" -- the + same payload that already resumes cleanly on the non-checkpoint path. + """ + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": "$89.99"}, + ) + approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call) + await ctx.request_info(approval_request, Content, request_id="approval-1") + + @response_handler + async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None: + del original_request + status = "approved" if bool(response.approved) else "rejected" + await ctx.yield_output(f"Refund tool call {status}.") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + + # First run: hit the approval interrupt and let core checkpoint the pending state. + first_events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "go"}]}, + workflow, + checkpoint_storage=storage, + ) + ] + assert "RUN_ERROR" not in [event.type for event in first_events] + first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0] + interrupt_payload = _interrupts_from_run_finished(first_finished) + interrupt_value = _interrupt_metadata_value(interrupt_payload[0]) + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the interrupted run to create a checkpoint" + resume_checkpoint_id = checkpoints[-1].checkpoint_id + + # Resume on a FRESH workflow instance so no pending requests exist in memory until + # the checkpoint is restored -- a cold restore, as after a process restart. + resumed_workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + resumed_events: list[Any] = [ + event + async for event in run_workflow_stream( + { + "messages": [], + "resume": { + "interrupts": [ + { + "id": "approval-1", + "value": { + "type": "function_approval_response", + "approved": True, + "id": interrupt_value.get("id", "approval-1"), + "function_call": interrupt_value.get("function_call"), + }, + } + ] + }, + }, + resumed_workflow, + checkpoint_id=resume_checkpoint_id, + checkpoint_storage=storage, + ) + ] + + resumed_types = [event.type for event in resumed_events] + assert "RUN_ERROR" not in resumed_types + assert "TEXT_MESSAGE_CONTENT" in resumed_types + text_deltas = [event.delta for event in resumed_events if event.type == "TEXT_MESSAGE_CONTENT"] + assert any("approved" in delta for delta in text_deltas) + + async def test_workflow_run_resume_message_list_from_json_payload() -> None: """Resume payloads should coerce AG-UI message dictionaries into list[Message] responses.""" From 4c3c3af9b722bad3095b38464bd487568e5f5186 Mon Sep 17 00:00:00 2001 From: vaibhav-patel Date: Mon, 13 Jul 2026 14:58:45 +0400 Subject: [PATCH 4/5] Coerce ag-ui resume responses without a second checkpoint restore Reading pending request_info events for resume-response coercion previously restored the checkpoint into the live workflow, which invoked every executor's on_checkpoint_restore hook. workflow.run(checkpoint_id=...) then restored again, running those hooks a second time. Custom restore hooks are not required to be idempotent, so this could duplicate restoration work or break workflows that expect exactly one restore per resume. Load the persisted WorkflowCheckpoint directly from storage (runtime override or the workflow's build-time context storage) and read its pending_request_info_events instead. This exposes the same post-restore pending set for the resume contract and response coercion without mutating workflow state or running any restore hook, leaving workflow.run(checkpoint_id=...) as the single restore per resume. Add a regression test asserting on_checkpoint_restore runs exactly once on a checkpointed ag-ui resume. --- .../agent_framework_ag_ui/_workflow_run.py | 83 ++++++++++------ .../ag-ui/tests/ag_ui/test_workflow_run.py | 94 +++++++++++++++++++ 2 files changed, 150 insertions(+), 27 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 4bea5a7d2d2..0e4bf2d984a 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -31,6 +31,7 @@ Content, Message, Workflow, + WorkflowCheckpoint, WorkflowRunState, ) @@ -154,30 +155,57 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: return {} -async def _restore_checkpoint_for_coercion( +async def _load_checkpoint( workflow: Workflow, checkpoint_id: str, checkpoint_storage: CheckpointStorage | None, -) -> None: - """Restore a checkpoint into the live workflow so pending requests become visible. - - Resume responses are coerced against the workflow's *current* pending requests. - On a cold checkpoint restore those requests do not exist yet on the live instance, - so this best-effort restore hydrates them first. It is idempotent: the subsequent - ``workflow.run(checkpoint_id=...)`` restores again from the same checkpoint. Any - failure is swallowed here and surfaced by that core run instead. +) -> WorkflowCheckpoint | None: + """Best-effort load of a persisted checkpoint *without* restoring it. + + Storage resolution mirrors ``Workflow.run(checkpoint_id=..., checkpoint_storage=...)``: + an explicitly supplied ``checkpoint_storage`` (which the resume run installs as its + runtime override) takes precedence; otherwise the workflow's build-time checkpoint + storage, exposed through its runner context, is used. Loading only reads the persisted + ``WorkflowCheckpoint`` -- it never mutates workflow state and never invokes executor + ``on_checkpoint_restore`` hooks. Any failure is swallowed here and surfaced by the + core ``workflow.run(checkpoint_id=...)`` call instead. """ - runner = getattr(workflow, "_runner", None) - restore_from_checkpoint = getattr(runner, "restore_from_checkpoint", None) - if restore_from_checkpoint is None: - return try: - await restore_from_checkpoint(checkpoint_id, checkpoint_storage) + if checkpoint_storage is not None: + return await checkpoint_storage.load(checkpoint_id) + runner_context = getattr(workflow, "_runner_context", None) + has_checkpointing = getattr(runner_context, "has_checkpointing", None) + load_checkpoint = getattr(runner_context, "load_checkpoint", None) + if callable(has_checkpointing) and has_checkpointing() and load_checkpoint is not None: + return cast(WorkflowCheckpoint, await load_checkpoint(checkpoint_id)) except Exception: # pragma: no cover - defensive; the core run re-raises the real error logger.warning( - "Pre-restore for resume-response coercion failed; the core run will surface any error.", + "Could not load checkpoint for resume-response coercion; the core run will surface any error.", exc_info=True, ) + return None + + +async def _pending_request_events_from_checkpoint( + workflow: Workflow, + checkpoint_id: str, + checkpoint_storage: CheckpointStorage | None, +) -> dict[str, Any]: + """Read pending request_info events from a persisted checkpoint without restoring it. + + Resume responses are coerced against the requests that were pending when the + checkpoint was written. On a cold checkpoint resume those requests are not yet live + on the workflow instance, so the coercion cannot see them. Reading + ``pending_request_info_events`` straight from the persisted ``WorkflowCheckpoint`` + exposes them without running any executor ``on_checkpoint_restore`` hook. The single + ``workflow.run(checkpoint_id=...)`` below then performs the one real restore, so the + restore -- and every custom restore hook -- runs exactly once per resume. + """ + checkpoint = await _load_checkpoint(workflow, checkpoint_id, checkpoint_storage) + pending = getattr(checkpoint, "pending_request_info_events", None) + if isinstance(pending, dict): + return dict(cast(dict[str, Any], pending)) + return {} def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None: @@ -819,19 +847,20 @@ async def run_workflow_stream( # A checkpoint resume that carries an explicit resume payload targets the requests # that were pending when the checkpoint was written; those only reappear on the live - # instance once the checkpoint is restored. Restore up front so ``pending_before_run`` - # (and the resume contract + coercion below) see the post-restore pending set, exactly - # as the non-checkpoint path sees the live pending requests. Without this a raw JSON - # resume payload (e.g. a ``function_approval_response`` dict) would reach core - # uncoerced and be rejected with a response-type mismatch. Only pre-restore when a - # resume payload is present so a pure checkpoint restore still surfaces its pending - # interrupts instead of tripping the "resume required" contract. The - # ``workflow.run(checkpoint_id=...)`` call below restores again from the same - # checkpoint, so this pre-restore is idempotent. + # instance once the checkpoint is restored. Rather than restore up front (which would + # invoke every executor's ``on_checkpoint_restore`` hook a second time, since + # ``workflow.run(checkpoint_id=...)`` restores again below), read the pending request + # events straight from the persisted checkpoint. That gives ``pending_before_run`` + # (and the resume contract + coercion below) the same post-restore pending set the + # non-checkpoint path sees from live requests, without any extra restore. Without it a + # raw JSON resume payload (e.g. a ``function_approval_response`` dict) would reach core + # uncoerced and be rejected with a response-type mismatch. Only load the checkpoint's + # pending set when a resume payload is present so a pure checkpoint restore still + # surfaces its pending interrupts instead of tripping the "resume required" contract. if checkpoint_id is not None and resume_payload is not None: - await _restore_checkpoint_for_coercion(workflow, checkpoint_id, checkpoint_storage) - - pending_before_run = await _pending_request_events(workflow) + pending_before_run = await _pending_request_events_from_checkpoint(workflow, checkpoint_id, checkpoint_storage) + else: + pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) resume_entries: list[dict[str, Any]] = [] if pending_interrupt_ids: diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index 8b98893c158..36eca2e70b5 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -421,6 +421,100 @@ async def handle_approval(self, original_request: Content, response: Content, ct assert any("approved" in delta for delta in text_deltas) +async def test_workflow_run_resume_restores_checkpoint_exactly_once() -> None: + """A checkpointed AG-UI resume must restore -- and run on_checkpoint_restore -- once. + + Custom ``on_checkpoint_restore`` hooks are not required to be idempotent. Coercing the + resume responses against the checkpoint's pending requests must therefore read the + persisted pending set without a second restore. This asserts the resumed executor's + restore hook fires exactly once (it would fire twice under a pre-restore-then-run + approach that hydrates pending requests by restoring the whole checkpoint first). + """ + + class CountingApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + self.restore_count = 0 + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345", "amount": "$89.99"}, + ) + approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call) + await ctx.request_info(approval_request, Content, request_id="approval-1") + + @response_handler + async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None: + del original_request + status = "approved" if bool(response.approved) else "rejected" + await ctx.yield_output(f"Refund tool call {status}.") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.restore_count += 1 + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=CountingApprovalExecutor(), checkpoint_storage=storage).build() + + # First run: hit the approval interrupt and let core checkpoint the pending state. + first_events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "go"}]}, + workflow, + checkpoint_storage=storage, + ) + ] + assert "RUN_ERROR" not in [event.type for event in first_events] + first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0] + interrupt_payload = _interrupts_from_run_finished(first_finished) + interrupt_value = _interrupt_metadata_value(interrupt_payload[0]) + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the interrupted run to create a checkpoint" + resume_checkpoint_id = checkpoints[-1].checkpoint_id + + # Resume on a FRESH workflow instance (cold restore, as after a process restart) so the + # restore hook count starts at zero and reflects only restores performed by this resume. + resumed_executor = CountingApprovalExecutor() + resumed_workflow = WorkflowBuilder(start_executor=resumed_executor).build() + resumed_events: list[Any] = [ + event + async for event in run_workflow_stream( + { + "messages": [], + "resume": { + "interrupts": [ + { + "id": "approval-1", + "value": { + "type": "function_approval_response", + "approved": True, + "id": interrupt_value.get("id", "approval-1"), + "function_call": interrupt_value.get("function_call"), + }, + } + ] + }, + }, + resumed_workflow, + checkpoint_id=resume_checkpoint_id, + checkpoint_storage=storage, + ) + ] + + assert "RUN_ERROR" not in [event.type for event in resumed_events] + assert resumed_executor.restore_count == 1, ( + f"expected exactly one checkpoint restore per resume, got {resumed_executor.restore_count}" + ) + + async def test_workflow_run_resume_message_list_from_json_payload() -> None: """Resume payloads should coerce AG-UI message dictionaries into list[Message] responses.""" From d0c27e5ccef5f5e6fd3f484b42a680e4b4996ebe Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 3 Aug 2026 15:52:00 +0900 Subject: [PATCH 5/5] Python: rework AG-UI workflow checkpointing onto public configuration surfaces Checkpoint storage is now configured on AgentFrameworkWorkflow (or the FastAPI endpoint) instead of being smuggled through input_data keys, and a run resumes by supplying its checkpoint id in the AG-UI forwarded props. With storage always in hand, resume-response coercion reads the pending request set straight from the persisted checkpoint via the public CheckpointStorage.load(), replacing the private runner-context fallback, and the core run call forwards checkpoint arguments directly, relying on core validation for conflicting parameters. Requesting a resume without configured storage now fails with a clear error. --- .../ag-ui/agent_framework_ag_ui/_endpoint.py | 19 +++- .../ag-ui/agent_framework_ag_ui/_workflow.py | 79 +++++++-------- .../agent_framework_ag_ui/_workflow_run.py | 97 ++++++------------- .../ag-ui/tests/ag_ui/test_endpoint.py | 70 +++++++++++-- .../ag-ui/tests/ag_ui/test_workflow_agent.py | 81 ++++++++++------ 5 files changed, 194 insertions(+), 152 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py index 5486802d647..06a3031f61a 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -12,7 +12,7 @@ from ag_ui.core import RunErrorEvent from ag_ui.encoder import EventEncoder -from agent_framework import SupportsAgentRun, Workflow +from agent_framework import CheckpointStorage, SupportsAgentRun, Workflow from fastapi import FastAPI, HTTPException from fastapi.params import Depends from fastapi.responses import Response, StreamingResponse @@ -90,6 +90,7 @@ def add_agent_framework_fastapi_endpoint( dependencies: Sequence[Depends] | None = None, snapshot_store: AGUIThreadSnapshotStore | None = None, snapshot_scope_resolver: SnapshotScopeResolver | None = None, + checkpoint_storage: CheckpointStorage | None = None, keepalive_seconds: float | None = 15, ) -> None: """Add an AG-UI endpoint to a FastAPI app. @@ -113,6 +114,10 @@ def add_agent_framework_fastapi_endpoint( snapshot_scope_resolver: Optional resolver for the application-defined Snapshot Scope. Required whenever a snapshot store is configured because an AG-UI Thread id is not an authorization boundary. Also scopes in-memory workflow_factory instances when provided without a snapshot store. + checkpoint_storage: Optional workflow checkpoint storage, applied when the endpoint exposes a workflow. + When provided, each run creates a checkpoint at the end of every superstep, and a run may resume from + a persisted checkpoint by supplying its id in the AG-UI forwarded props + (``forwarded_props: {"checkpoint_id": ...}``). keepalive_seconds: Endpoint SSE keepalive interval in seconds. Defaults to 15. Positive values emit fixed SSE comments while the stream is open. None disables keepalive and preserves the non-keepalive response path. Keepalive comments are transport traffic and do not change AG-UI events. @@ -125,7 +130,7 @@ def add_agent_framework_fastapi_endpoint( elif isinstance(agent, AgentFrameworkAgent): protocol_runner = agent elif isinstance(agent, Workflow): - protocol_runner = AgentFrameworkWorkflow(workflow=agent) + protocol_runner = AgentFrameworkWorkflow(workflow=agent, checkpoint_storage=checkpoint_storage) elif isinstance(agent, SupportsAgentRun): protocol_runner = AgentFrameworkAgent( agent=agent, @@ -136,6 +141,16 @@ def add_agent_framework_fastapi_endpoint( else: raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.") + if checkpoint_storage is not None: + if not isinstance(protocol_runner, AgentFrameworkWorkflow): + raise ValueError("checkpoint_storage is only supported when the endpoint exposes a workflow.") + if ( + protocol_runner.checkpoint_storage is not None + and protocol_runner.checkpoint_storage is not checkpoint_storage + ): + raise ValueError("checkpoint_storage is already configured on the AG-UI workflow runner.") + protocol_runner.checkpoint_storage = checkpoint_storage + _configure_snapshot_persistence( protocol_runner, snapshot_store=snapshot_store, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 9cdda94c526..b499c250aa7 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -49,12 +49,16 @@ WorkflowFactory = Callable[[str], Workflow] -# Input-payload keys used to surface workflow checkpointing through ``run(input_data)`` -# without changing the positional call convention used by the FastAPI endpoint. The -# corresponding ``run()`` keyword arguments take precedence over these when both are -# supplied. -_CHECKPOINT_ID_INPUT_KEY = "__ag_ui_checkpoint_id" -_CHECKPOINT_STORAGE_INPUT_KEY = "__ag_ui_checkpoint_storage" + +def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: + """Read an optional checkpoint id to resume from out of the AG-UI forwarded props.""" + forwarded_props = input_data.get("forwarded_props") or input_data.get("forwardedProps") + if not isinstance(forwarded_props, dict): + return None + checkpoint_id = forwarded_props.get("checkpoint_id") or forwarded_props.get("checkpointId") + if checkpoint_id is None: + return None + return str(checkpoint_id) def _cancelled_resume_interrupt_ids(resume_payload: Any) -> set[str]: @@ -240,6 +244,7 @@ def __init__( name: str | None = None, description: str | None = None, snapshot_store: AGUIThreadSnapshotStore | None = None, + checkpoint_storage: CheckpointStorage | None = None, ) -> None: """Initialize the AG-UI workflow wrapper. @@ -250,6 +255,11 @@ def __init__( description: Optional workflow description. snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless endpoint setup also provides an explicit Snapshot Scope resolver. + checkpoint_storage: Optional workflow checkpoint storage. When provided, each run + creates a checkpoint at the end of every superstep (matching + ``agent_framework.Workflow.run(checkpoint_storage=...)``), and a run may resume + from a persisted checkpoint by supplying its id in the AG-UI forwarded props + (``forwarded_props: {"checkpoint_id": ...}``). Required for checkpoint resume. """ if workflow is not None and workflow_factory is not None: raise ValueError("Pass either workflow= or workflow_factory=, not both.") @@ -264,6 +274,7 @@ def __init__( self.name = name if name is not None else getattr(workflow, "name", "workflow") self.description = description if description is not None else getattr(workflow, "description", "") self.snapshot_store = snapshot_store + self.checkpoint_storage = checkpoint_storage @staticmethod def _thread_id_from_input(input_data: dict[str, Any]) -> str: @@ -302,32 +313,17 @@ def clear_workflow_cache(self) -> None: """Drop all cached thread workflow instances.""" self._workflow_by_thread.clear() - async def run( - self, - input_data: dict[str, Any], - *, - checkpoint_storage: CheckpointStorage | None = None, - checkpoint_id: str | None = None, - ) -> AsyncGenerator[BaseEvent]: + async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: """Run the wrapped workflow and yield AG-UI events. - Args: - input_data: The AG-UI request payload (a ``RunAgentInput`` dump). - checkpoint_storage: Optional checkpoint storage to enable workflow - checkpointing for this run. When provided, the underlying core - workflow creates a checkpoint at the end of each superstep, matching - ``agent_framework.Workflow.run(checkpoint_storage=...)``. May also be - supplied via the ``input_data`` key ``__ag_ui_checkpoint_storage``; - the keyword argument takes precedence. - checkpoint_id: Optional checkpoint id to resume the workflow from. When - provided, execution restores the persisted workflow state instead of - starting a fresh turn, matching - ``agent_framework.Workflow.run(checkpoint_id=...)``. May also be - supplied via the ``input_data`` key ``__ag_ui_checkpoint_id``; the - keyword argument takes precedence. - Subclasses may override this to provide custom AG-UI streams. + When ``checkpoint_storage`` is configured on this wrapper, the underlying core + workflow creates a checkpoint at the end of each superstep, and a run may resume + from a persisted checkpoint by supplying its id in the AG-UI forwarded props + (``forwarded_props: {"checkpoint_id": ...}``), which restores the persisted + workflow state instead of starting a fresh turn. + Note: Checkpointing (the ``agent_framework`` workflow checkpoint mechanism) is independent from AG-UI Thread Snapshot persistence (``snapshot_store``). @@ -343,13 +339,13 @@ async def run( resume_payload = _extract_resume_payload(input_data) snapshot_store = self.snapshot_store - # Explicit keyword arguments win over values smuggled through input_data so the - # FastAPI endpoint (which calls ``run(input_data)`` positionally) can still opt - # into checkpointing without changing its call site. - if checkpoint_id is None: - checkpoint_id = cast(str | None, input_data.get(_CHECKPOINT_ID_INPUT_KEY)) - if checkpoint_storage is None: - checkpoint_storage = cast(CheckpointStorage | None, input_data.get(_CHECKPOINT_STORAGE_INPUT_KEY)) + checkpoint_storage = self.checkpoint_storage + checkpoint_id = _checkpoint_id_from_input(input_data) + if checkpoint_id is not None and checkpoint_storage is None: + raise ValueError( + "Resuming from a checkpoint requires checkpoint_storage to be configured on " + "AgentFrameworkWorkflow (or the AG-UI endpoint)." + ) # A checkpoint resume legitimately carries no new messages; it must reach the # core workflow's restore path rather than replaying a stored thread snapshot. @@ -424,17 +420,10 @@ async def run( state_snapshot = make_json_safe(effective_state) if isinstance(state_snapshot, dict): snapshot_builder.state = cast(dict[str, Any], state_snapshot) - # Only forward checkpoint kwargs when checkpointing is requested so the - # non-checkpoint path keeps calling ``run_workflow_stream(input_data, workflow)`` - # exactly as before (preserves the established two-argument call convention). - stream_kwargs: dict[str, Any] = {} - if checkpoint_storage is not None: - stream_kwargs["checkpoint_storage"] = checkpoint_storage - if checkpoint_id is not None: - stream_kwargs["checkpoint_id"] = checkpoint_id - run_error_emitted = False - async for event in run_workflow_stream(input_data, workflow, **stream_kwargs): + async for event in run_workflow_stream( + input_data, workflow, checkpoint_storage=checkpoint_storage, checkpoint_id=checkpoint_id + ): if snapshot_builder is not None: snapshot_builder.observe(event) if isinstance(event, RunErrorEvent): diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 5e5f4b0270d..86d9e369974 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -31,7 +31,6 @@ Content, Message, Workflow, - WorkflowCheckpoint, WorkflowRunState, ) @@ -155,41 +154,9 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: return {} -async def _load_checkpoint( - workflow: Workflow, - checkpoint_id: str, - checkpoint_storage: CheckpointStorage | None, -) -> WorkflowCheckpoint | None: - """Best-effort load of a persisted checkpoint *without* restoring it. - - Storage resolution mirrors ``Workflow.run(checkpoint_id=..., checkpoint_storage=...)``: - an explicitly supplied ``checkpoint_storage`` (which the resume run installs as its - runtime override) takes precedence; otherwise the workflow's build-time checkpoint - storage, exposed through its runner context, is used. Loading only reads the persisted - ``WorkflowCheckpoint`` -- it never mutates workflow state and never invokes executor - ``on_checkpoint_restore`` hooks. Any failure is swallowed here and surfaced by the - core ``workflow.run(checkpoint_id=...)`` call instead. - """ - try: - if checkpoint_storage is not None: - return await checkpoint_storage.load(checkpoint_id) - runner_context = getattr(workflow, "_runner_context", None) - has_checkpointing = getattr(runner_context, "has_checkpointing", None) - load_checkpoint = getattr(runner_context, "load_checkpoint", None) - if callable(has_checkpointing) and has_checkpointing() and load_checkpoint is not None: - return cast(WorkflowCheckpoint, await load_checkpoint(checkpoint_id)) - except Exception: # pragma: no cover - defensive; the core run re-raises the real error - logger.warning( - "Could not load checkpoint for resume-response coercion; the core run will surface any error.", - exc_info=True, - ) - return None - - async def _pending_request_events_from_checkpoint( - workflow: Workflow, checkpoint_id: str, - checkpoint_storage: CheckpointStorage | None, + checkpoint_storage: CheckpointStorage, ) -> dict[str, Any]: """Read pending request_info events from a persisted checkpoint without restoring it. @@ -197,15 +164,19 @@ async def _pending_request_events_from_checkpoint( checkpoint was written. On a cold checkpoint resume those requests are not yet live on the workflow instance, so the coercion cannot see them. Reading ``pending_request_info_events`` straight from the persisted ``WorkflowCheckpoint`` - exposes them without running any executor ``on_checkpoint_restore`` hook. The single - ``workflow.run(checkpoint_id=...)`` below then performs the one real restore, so the + exposes them without running any executor ``on_checkpoint_restore`` hook; the single + ``workflow.run(checkpoint_id=...)`` then performs the one real restore, so the restore -- and every custom restore hook -- runs exactly once per resume. """ - checkpoint = await _load_checkpoint(workflow, checkpoint_id, checkpoint_storage) - pending = getattr(checkpoint, "pending_request_info_events", None) - if isinstance(pending, dict): - return dict(cast(dict[str, Any], pending)) - return {} + try: + checkpoint = await checkpoint_storage.load(checkpoint_id) + except Exception: + logger.warning( + "Could not load checkpoint for resume-response coercion; the core run will surface any error.", + exc_info=True, + ) + return {} + return dict(checkpoint.pending_request_info_events) def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None: @@ -875,18 +846,14 @@ async def run_workflow_stream( # A checkpoint resume that carries an explicit resume payload targets the requests # that were pending when the checkpoint was written; those only reappear on the live - # instance once the checkpoint is restored. Rather than restore up front (which would - # invoke every executor's ``on_checkpoint_restore`` hook a second time, since - # ``workflow.run(checkpoint_id=...)`` restores again below), read the pending request - # events straight from the persisted checkpoint. That gives ``pending_before_run`` - # (and the resume contract + coercion below) the same post-restore pending set the - # non-checkpoint path sees from live requests, without any extra restore. Without it a - # raw JSON resume payload (e.g. a ``function_approval_response`` dict) would reach core - # uncoerced and be rejected with a response-type mismatch. Only load the checkpoint's - # pending set when a resume payload is present so a pure checkpoint restore still - # surfaces its pending interrupts instead of tripping the "resume required" contract. + # instance once the checkpoint is restored, so coerce against the checkpoint's + # persisted pending set instead. Only do so when a resume payload is present, so a + # pure checkpoint restore still surfaces its pending interrupts instead of tripping + # the "resume required" contract. if checkpoint_id is not None and resume_payload is not None: - pending_before_run = await _pending_request_events_from_checkpoint(workflow, checkpoint_id, checkpoint_storage) + if checkpoint_storage is None: + raise ValueError("Resuming a checkpoint with an AG-UI resume payload requires checkpoint_storage.") + pending_before_run = await _pending_request_events_from_checkpoint(checkpoint_id, checkpoint_storage) else: pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) @@ -987,25 +954,19 @@ def _drain_open_message() -> list[TextMessageEndEvent]: logger.debug("workflow.run() does not accept function_invocation_kwargs; dropping forwarded_props") fwd_kwargs = {} - # Forward checkpoint storage so the core workflow creates a checkpoint at the end - # of each superstep (parity with ``Workflow.run(checkpoint_storage=...)``). + # When checkpointing is not in play, keep the exact legacy call shape so duck-typed + # workflows with narrower ``run`` signatures keep working. Otherwise forward the + # checkpoint arguments as-is (``None`` included) and let core validate conflicts. checkpoint_kwargs: dict[str, Any] = {} - if checkpoint_storage is not None: - checkpoint_kwargs["checkpoint_storage"] = checkpoint_storage + if checkpoint_storage is not None or checkpoint_id is not None: + checkpoint_kwargs = {"checkpoint_storage": checkpoint_storage, "checkpoint_id": checkpoint_id} try: - if checkpoint_id is not None: - # Resume from a checkpoint. ``message`` is mutually exclusive with - # ``checkpoint_id`` in the core API, so incoming messages are only - # forwarded as request-info responses (``responses``), never as a new - # start-executor message. ``responses`` + ``checkpoint_id`` performs a - # restore-then-send in a single call. - run_kwargs: dict[str, Any] = {"checkpoint_id": checkpoint_id, **checkpoint_kwargs} - if responses: - run_kwargs["responses"] = responses - event_stream = workflow.run(stream=True, **run_kwargs, **fwd_kwargs) - elif responses: - event_stream = workflow.run(responses=responses, stream=True, **checkpoint_kwargs, **fwd_kwargs) + if responses or checkpoint_id is not None: + # ``message`` is mutually exclusive with both ``responses`` and + # ``checkpoint_id`` in the core API; ``responses`` + ``checkpoint_id`` + # restores the checkpoint and delivers the responses in a single call. + event_stream = workflow.run(stream=True, responses=responses or None, **checkpoint_kwargs, **fwd_kwargs) else: event_stream = workflow.run(message=messages, stream=True, **checkpoint_kwargs, **fwd_kwargs) diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index e5960c51891..71fc1849614 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -22,6 +22,7 @@ ContextProvider, Executor, FunctionTool, + InMemoryCheckpointStorage, InMemoryHistoryProvider, Message, SessionContext, @@ -150,6 +151,59 @@ async def start(message: Any, ctx: WorkflowContext[Any, Any]) -> None: assert "RUN_FINISHED" in event_types +async def test_add_endpoint_workflow_checkpointing_over_the_wire(): + """Endpoint checkpoint_storage creates checkpoints, and forwardedProps.checkpoint_id resumes.""" + + @executor(id="start") + async def start(message: Any, ctx: WorkflowContext[str]) -> None: + del message + await ctx.send_message("hello", target_id="finish") + + @executor(id="finish") + async def finish(message: str, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(f"{message}-done") + + def event_types_of(response: Any) -> list[str]: + lines = [line for line in response.content.decode("utf-8").split("\n") if line.startswith("data: ")] + return [json.loads(line[6:]).get("type") for line in lines] + + app = FastAPI() + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=start).add_edge(start, finish).build() + + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow", checkpoint_storage=storage) + + client = TestClient(app) + response = client.post( + "/workflow", + json={"threadId": "thread-cp", "messages": [{"role": "user", "content": "go"}]}, + ) + assert response.status_code == 200 + assert "RUN_ERROR" not in event_types_of(response) + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the run to create at least one checkpoint" + + resume_response = client.post( + "/workflow", + json={ + "threadId": "thread-cp", + "messages": [], + "forwardedProps": {"checkpoint_id": checkpoints[0].checkpoint_id}, + }, + ) + assert resume_response.status_code == 200 + resumed_types = event_types_of(resume_response) + assert "RUN_FINISHED" in resumed_types + assert "RUN_ERROR" not in resumed_types + # The restored run must replay the remaining superstep and re-produce the final + # output; a run that silently ignored the checkpoint id would finish with no text. + assert "TEXT_MESSAGE_CONTENT" in resumed_types + + async def test_add_endpoint_accepts_keepalive_option_for_supported_runners(build_chat_client): """Keepalive configuration is accepted at the endpoint seam for every supported runner shape.""" @@ -3375,8 +3429,8 @@ async def test_endpoint_streaming_error_emits_run_error_event(): """Streaming exceptions should emit RUN_ERROR instead of terminating silently.""" class FailingStreamWorkflow(AgentFrameworkWorkflow): - async def run(self, input_data: dict[str, Any], **kwargs: Any): - del input_data, kwargs + async def run(self, input_data: dict[str, Any]): + del input_data yield RunStartedEvent(run_id="run-1", thread_id="thread-1") raise RuntimeError("stream exploded") @@ -4147,8 +4201,8 @@ async def test_endpoint_encoding_failure_emits_run_error(): from unittest.mock import patch class SimpleWorkflow(AgentFrameworkWorkflow): - async def run(self, input_data: dict[str, Any], **kwargs: Any): - del input_data, kwargs + async def run(self, input_data: dict[str, Any]): + del input_data yield RunStartedEvent(run_id="run-1", thread_id="thread-1") app = FastAPI() @@ -4170,8 +4224,8 @@ async def test_endpoint_double_encoding_failure_terminates(): from unittest.mock import patch class SimpleWorkflow(AgentFrameworkWorkflow): - async def run(self, input_data: dict[str, Any], **kwargs: Any): - del input_data, kwargs + async def run(self, input_data: dict[str, Any]): + del input_data yield RunStartedEvent(run_id="run-1", thread_id="thread-1") app = FastAPI() @@ -4780,8 +4834,8 @@ async def test_workflow_resume_preserves_persisted_history(monkeypatch): ), ) - async def fake_run_workflow_stream(input_data: Any, workflow: Any): - del input_data, workflow + async def fake_run_workflow_stream(input_data: Any, workflow: Any, **kwargs: Any): + del input_data, workflow, kwargs yield RunStartedEvent(run_id="run-2", thread_id="workflow-thread") yield TextMessageStartEvent(message_id="resume-msg", role="assistant") yield TextMessageContentEvent(message_id="resume-msg", delta="Resumed reply") diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index fd8c8f2f750..4a4ec411605 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -20,12 +20,8 @@ from agent_framework_ag_ui import AgentFrameworkWorkflow -async def _run( - agent: AgentFrameworkWorkflow, - payload: dict[str, Any], - **run_kwargs: Any, -) -> list[Any]: - return [event async for event in agent.run(payload, **run_kwargs)] +async def _run(agent: AgentFrameworkWorkflow, payload: dict[str, Any]) -> list[Any]: + return [event async for event in agent.run(payload)] def _interrupts_from_finished(event: Any) -> list[dict[str, Any]]: @@ -168,16 +164,15 @@ def _build_multi_superstep_workflow(storage: InMemoryCheckpointStorage | None = return builder.add_edge(start, middle).add_edge(middle, finish).build() -async def test_workflow_run_creates_checkpoints_via_storage_kwarg() -> None: - """Passing checkpoint_storage to run() should create workflow checkpoints (parity with core).""" +async def test_workflow_run_creates_checkpoints_via_constructor_storage() -> None: + """Configuring checkpoint_storage on the wrapper should create workflow checkpoints (parity with core).""" storage = InMemoryCheckpointStorage() workflow = _build_multi_superstep_workflow() - agent = AgentFrameworkWorkflow(workflow=workflow) + agent = AgentFrameworkWorkflow(workflow=workflow, checkpoint_storage=storage) events = await _run( agent, {"thread_id": "thread-cp", "messages": [{"role": "user", "content": "start"}]}, - checkpoint_storage=storage, ) event_types = [event.type for event in events] @@ -191,10 +186,10 @@ async def test_workflow_run_creates_checkpoints_via_storage_kwarg() -> None: async def test_workflow_run_resumes_from_checkpoint_id() -> None: - """run(checkpoint_id=...) should restore persisted state and finish the workflow.""" + """A checkpoint_id in the forwarded props should restore persisted state and finish the workflow.""" storage = InMemoryCheckpointStorage() workflow = _build_multi_superstep_workflow(storage) - agent = AgentFrameworkWorkflow(workflow=workflow) + agent = AgentFrameworkWorkflow(workflow=workflow, checkpoint_storage=storage) # First run: execute to completion while checkpoints are written. first_events = await _run( @@ -214,9 +209,11 @@ async def test_workflow_run_resumes_from_checkpoint_id() -> None: # Resume on the same thread (same underlying workflow instance) from the checkpoint. resumed_events = await _run( agent, - {"thread_id": "thread-cp", "messages": []}, - checkpoint_id=resume_checkpoint_id, - checkpoint_storage=storage, + { + "thread_id": "thread-cp", + "messages": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, + }, ) resumed_types = [event.type for event in resumed_events] @@ -231,24 +228,51 @@ async def test_workflow_run_resumes_from_checkpoint_id() -> None: assert "done" in resumed_text -async def test_workflow_run_reads_checkpoint_params_from_input_data() -> None: - """Checkpoint params smuggled through input_data should be honored (endpoint call convention).""" +async def test_workflow_run_reads_checkpoint_id_from_camelcase_forwarded_props() -> None: + """A camelCase ``forwardedProps.checkpointId`` payload (wire format) should also resume.""" storage = InMemoryCheckpointStorage() - workflow = _build_multi_superstep_workflow() - agent = AgentFrameworkWorkflow(workflow=workflow) + workflow = _build_multi_superstep_workflow(storage) + agent = AgentFrameworkWorkflow(workflow=workflow, checkpoint_storage=storage) - events = await _run( + first_events = await _run( + agent, + {"thread_id": "thread-cp-camel", "messages": [{"role": "user", "content": "start"}]}, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the run to create at least one checkpoint" + + resumed_events = await _run( agent, { - "thread_id": "thread-cp-input", - "messages": [{"role": "user", "content": "start"}], - "__ag_ui_checkpoint_storage": storage, + "thread_id": "thread-cp-camel", + "messages": [], + "forwardedProps": {"checkpointId": checkpoints[0].checkpoint_id}, }, ) + resumed_types = [event.type for event in resumed_events] + assert "RUN_FINISHED" in resumed_types + assert "RUN_ERROR" not in resumed_types - assert "RUN_ERROR" not in [event.type for event in events] - checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) - assert len(checkpoints) >= 1 + +async def test_workflow_resume_without_checkpoint_storage_raises() -> None: + """Requesting a checkpoint resume without configured storage should fail loudly.""" + workflow = _build_multi_superstep_workflow() + agent = AgentFrameworkWorkflow(workflow=workflow) + + with pytest.raises(ValueError, match="requires checkpoint_storage"): + await _run( + agent, + { + "thread_id": "thread-cp-nostorage", + "messages": [], + "forwarded_props": {"checkpoint_id": "some-checkpoint"}, + }, + ) async def test_workflow_run_without_checkpointing_is_unchanged() -> None: @@ -277,7 +301,7 @@ async def test_workflow_checkpoint_only_resume_preserves_thread_snapshot() -> No storage = InMemoryCheckpointStorage() workflow = _build_multi_superstep_workflow(storage) store = InMemoryAGUIThreadSnapshotStore() - agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store) + agent = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store, checkpoint_storage=storage) # Prime the workflow so a checkpoint exists to resume from. first_events = await _run( @@ -321,10 +345,9 @@ async def test_workflow_checkpoint_only_resume_preserves_thread_snapshot() -> No "thread_id": "thread-cp-snap", "run_id": "run-2", "messages": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", }, - checkpoint_id=resume_checkpoint_id, - checkpoint_storage=storage, ) assert "RUN_ERROR" not in [event.type for event in resumed_events]