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 5486802d64..06a3031f61 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 20f44361c9..b499c250aa 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 agent_framework._telemetry import mark_feature_used from ._feature_usage import FeatureIndex @@ -50,6 +50,17 @@ WorkflowFactory = Callable[[str], Workflow] +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]: """Return cancelled interrupt ids from a resume payload.""" return { @@ -233,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. @@ -243,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.") @@ -257,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: @@ -299,6 +317,19 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: """Run the wrapped workflow and yield AG-UI events. 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``). + 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. """ mark_feature_used(FeatureIndex.AG_UI) thread_id = self._thread_id_from_input(input_data) @@ -308,7 +339,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: + 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. + 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, @@ -349,9 +396,16 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: 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 @@ -367,7 +421,9 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: if isinstance(state_snapshot, dict): snapshot_builder.state = cast(dict[str, Any], state_snapshot) run_error_emitted = False - async for event in run_workflow_stream(input_data, workflow): + 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 3f6e76e5fd..86d9e36997 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 ( @@ -146,6 +154,31 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: return {} +async def _pending_request_events_from_checkpoint( + checkpoint_id: str, + checkpoint_storage: CheckpointStorage, +) -> 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=...)`` then performs the one real restore, so the + restore -- and every custom restore hook -- runs exactly once per resume. + """ + 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: """Build AG-UI interrupt payload from a workflow request_info event.""" request_payload = _request_payload_from_request_event(request_event) @@ -775,8 +808,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") @@ -794,7 +843,19 @@ async def run_workflow_stream( last_assistant_text: str | None = None resume_payload = _extract_resume_payload(input_data) - pending_before_run = await _pending_request_events(workflow) + + # 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, 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: + 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) resume_entries: list[dict[str, Any]] = [] if pending_interrupt_ids: @@ -836,7 +897,11 @@ async def run_workflow_stream( return pending_interrupts = _interrupts_from_pending_requests(pending_before_run) - 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) @@ -853,7 +918,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 @@ -889,11 +954,21 @@ def _drain_open_message() -> list[TextMessageEndEvent]: logger.debug("workflow.run() does not accept function_invocation_kwargs; dropping forwarded_props") fwd_kwargs = {} + # 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 or checkpoint_id is not None: + checkpoint_kwargs = {"checkpoint_storage": checkpoint_storage, "checkpoint_id": checkpoint_id} + try: - if responses: - event_stream = workflow.run(responses=responses, stream=True, **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, **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_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 0408de841d..71fc184961 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.""" @@ -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 6fa23ac2cc..4a4ec41160 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,7 +7,15 @@ from typing import Any, cast import pytest -from agent_framework import Workflow, WorkflowBuilder, WorkflowContext, executor +from agent_framework import ( + Executor, + InMemoryCheckpointStorage, + Workflow, + WorkflowBuilder, + WorkflowContext, + executor, + handler, +) from agent_framework_ag_ui import AgentFrameworkWorkflow @@ -121,3 +129,233 @@ 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_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, checkpoint_storage=storage) + + events = await _run( + agent, + {"thread_id": "thread-cp", "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 + + 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: + """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, checkpoint_storage=storage) + + # 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": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, + }, + ) + + 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_id_from_camelcase_forwarded_props() -> None: + """A camelCase ``forwardedProps.checkpointId`` payload (wire format) should also resume.""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow(storage) + agent = AgentFrameworkWorkflow(workflow=workflow, checkpoint_storage=storage) + + 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-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 + + +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: + """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 + + +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, checkpoint_storage=storage) + + # 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": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, + _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", + }, + ) + 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 8af9a33be2..8ab02a6c85 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 @@ -16,6 +16,7 @@ ChatResponseUpdate, Content, Executor, + InMemoryCheckpointStorage, Message, WorkflowBuilder, WorkflowContext, @@ -330,6 +331,190 @@ 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_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."""