Python: add checkpointing support to AgentFrameworkWorkflow.run() in agent-framework-ag-ui

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,

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):

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)