Python: include checkpoint_id on AG-UI interrupt metadata (#8150)

python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py

@@ -8,7 +8,7 @@
 import json
 import logging
 import uuid
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from functools import partial
 from types import UnionType
 from typing import Any, Union, cast, get_args, get_origin, get_type_hints
@@ -127,7 +127,10 @@ def _workflow_interrupt_value(request_data: Any) -> Any:
     return {"data": safe_request_data}
 
 
-def _workflow_interrupt_metadata(request_payload: dict[str, Any], value: Any) -> dict[str, Any]:
+def _workflow_interrupt_metadata(
+    request_payload: dict[str, Any],
+    value: Any,
+) -> dict[str, Any]:
     """Build Agent Framework metadata for workflow request_info interrupts."""
     agent_framework_metadata = {
         key: make_json_safe(value)
@@ -145,6 +148,104 @@ def _workflow_interrupt_metadata(request_payload: dict[str, Any], value: Any) ->
     return {"agent_framework": agent_framework_metadata}
 
 
+def _attach_checkpoint_id_to_interrupts(
+    interrupts: list[dict[str, Any]],
+    checkpoint_id: str | None,
+) -> list[dict[str, Any]]:
+    """Attach ``checkpoint_id`` to each interrupt's ``metadata.agent_framework`` (issue #8150).
+
+    Multi-worker hosts need the pause checkpoint on the wire so the next resume can pass
+    ``forwardedProps.checkpoint_id`` without a side-channel lookup. No-op when checkpointing
+    is inactive or the id is already present.
+    """
+    if not checkpoint_id or not interrupts:
+        return interrupts
+
+    attached: list[dict[str, Any]] = []
+    for interrupt in interrupts:
+        entry = dict(interrupt)
+        metadata = entry.get("metadata")
+        if isinstance(metadata, dict):
+            metadata = dict(metadata)
+        else:
+            metadata = {}
+        agent_framework = metadata.get("agent_framework")
+        if isinstance(agent_framework, dict):
+            agent_framework = dict(agent_framework)
+        else:
+            agent_framework = {}
+        agent_framework.setdefault("checkpoint_id", checkpoint_id)
+        metadata["agent_framework"] = agent_framework
+        entry["metadata"] = metadata
+        attached.append(entry)
+    return attached
+
+
+def _interrupt_request_ids(interrupts: list[dict[str, Any]]) -> set[str]:
+    return {str(item["id"]) for item in interrupts if item.get("id") is not None}
+
+
+async def _pause_checkpoint_id_for_interrupts(
+    *,
+    workflow: Workflow,
+    checkpoint_storage: CheckpointStorage | None,
+    interrupts: list[dict[str, Any]],
+    known_checkpoint_id: str | None = None,
+    baseline_checkpoint_id: str | None = None,
+) -> str | None:
+    """Resolve the pause checkpoint for *this* run's interrupts via core.
+
+    ``baseline_checkpoint_id`` must be the runner id captured before ``workflow.run()`` so
+    a leftover id from a prior run is not advertised when this run did not persist.
+    """
+    if not interrupts:
+        return None
+
+    resolve = getattr(workflow, "resolve_pause_checkpoint_id", None)
+    if not callable(resolve):
+        return None
+
+    # getattr returns a plain object to the type checker; cast to an awaitable callable.
+    resolve_fn = cast(
+        Callable[..., Awaitable[str | None]],
+        resolve,
+    )
+    return await resolve_fn(
+        _interrupt_request_ids(interrupts),
+        checkpoint_storage=checkpoint_storage,
+        known_checkpoint_id=known_checkpoint_id,
+        baseline_checkpoint_id=baseline_checkpoint_id,
+    )
+
+
+async def _build_run_finished_with_checkpointed_interrupts(
+    *,
+    run_id: str,
+    thread_id: str,
+    interrupts: list[dict[str, Any]],
+    workflow: Workflow,
+    checkpoint_storage: CheckpointStorage | None,
+    known_checkpoint_id: str | None = None,
+    baseline_checkpoint_id: str | None = None,
+) -> Any:
+    """Build RUN_FINISHED, attaching the pause checkpoint id to interrupt metadata when available."""
+    if not interrupts:
+        return _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts)
+
+    pause_checkpoint_id = await _pause_checkpoint_id_for_interrupts(
+        workflow=workflow,
+        checkpoint_storage=checkpoint_storage,
+        interrupts=interrupts,
+        known_checkpoint_id=known_checkpoint_id,
+        baseline_checkpoint_id=baseline_checkpoint_id,
+    )
+    return _build_run_finished_event(
+        run_id=run_id,
+        thread_id=thread_id,
+        interrupts=_attach_checkpoint_id_to_interrupts(interrupts, pause_checkpoint_id),
+    )
+
+
 async def _pending_request_events(workflow: Workflow) -> dict[str, Any]:
     """Best-effort retrieval of pending request_info events from workflow context."""
     runner_context = getattr(workflow, "_runner_context", None)
@@ -1133,12 +1234,34 @@ async def run_workflow_stream(
             interrupt_event_value = _workflow_interrupt_event_value(request_payload)
             if interrupt_event_value is not None:
                 yield CustomEvent(name=_INTERRUPT_CARD_EVENT_NAME, value=interrupt_event_value)
-        yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts)
+        baseline_checkpoint_id = (
+            workflow.get_last_checkpoint_id() if hasattr(workflow, "get_last_checkpoint_id") else None
+        )
+        yield await _build_run_finished_with_checkpointed_interrupts(
+            run_id=run_id,
+            thread_id=thread_id,
+            interrupts=pending_interrupts,
+            workflow=workflow,
+            checkpoint_storage=checkpoint_storage,
+            known_checkpoint_id=checkpoint_id,
+            baseline_checkpoint_id=baseline_checkpoint_id,
+        )
         return
 
     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)
+        baseline_checkpoint_id = (
+            workflow.get_last_checkpoint_id() if hasattr(workflow, "get_last_checkpoint_id") else None
+        )
+        yield await _build_run_finished_with_checkpointed_interrupts(
+            run_id=run_id,
+            thread_id=thread_id,
+            interrupts=pending_interrupts,
+            workflow=workflow,
+            checkpoint_storage=checkpoint_storage,
+            known_checkpoint_id=checkpoint_id,
+            baseline_checkpoint_id=baseline_checkpoint_id,
+        )
         return
 
     def _drain_open_message() -> list[TextMessageEndEvent]:
@@ -1195,6 +1318,8 @@ def _drain_open_blocks() -> list[BaseEvent]:
     if checkpoint_storage is not None or checkpoint_id is not None:
         checkpoint_kwargs = {"checkpoint_storage": checkpoint_storage, "checkpoint_id": checkpoint_id}
 
+    baseline_checkpoint_id = workflow.get_last_checkpoint_id() if hasattr(workflow, "get_last_checkpoint_id") else None
+
     try:
         telemetry_conversation_id = str(supplied_thread_id) if supplied_thread_id is not None else None
         telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id)
@@ -1244,7 +1369,15 @@ def _drain_open_blocks() -> list[BaseEvent]:
                         yield end_event
                     if not interrupts:
                         interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow)))
-                    yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts)
+                    yield await _build_run_finished_with_checkpointed_interrupts(
+                        run_id=run_id,
+                        thread_id=thread_id,
+                        interrupts=interrupts,
+                        workflow=workflow,
+                        checkpoint_storage=checkpoint_storage,
+                        known_checkpoint_id=checkpoint_id,
+                        baseline_checkpoint_id=baseline_checkpoint_id,
+                    )
                     terminal_emitted = True
                 elif state_value not in _TERMINAL_STATES:
                     yield CustomEvent(name="status", value={"state": state_value})
@@ -1391,4 +1524,12 @@ def _drain_open_blocks() -> list[BaseEvent]:
     if not terminal_emitted and not run_error_emitted:
         if not interrupts:
             interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow)))
-        yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts)
+        yield await _build_run_finished_with_checkpointed_interrupts(
+            run_id=run_id,
+            thread_id=thread_id,
+            interrupts=interrupts,
+            workflow=workflow,
+            checkpoint_storage=checkpoint_storage,
+            known_checkpoint_id=checkpoint_id,
+            baseline_checkpoint_id=baseline_checkpoint_id,
+        )

python/packages/ag-ui/tests/ag_ui/test_workflow_run.py

@@ -95,6 +95,183 @@ def _interrupt_metadata_value(interrupt: dict[str, Any]) -> dict[str, Any]:
     return cast(dict[str, Any], value)
 
 
+def test_attach_checkpoint_id_to_interrupts_setdefault() -> None:
+    """Issue #8150: attach pause checkpoint_id without overwriting an existing one."""
+    from agent_framework_ag_ui._workflow_run import _attach_checkpoint_id_to_interrupts
+
+    empty = _attach_checkpoint_id_to_interrupts([{"id": "r1"}], None)
+    assert empty == [{"id": "r1"}]
+
+    attached = _attach_checkpoint_id_to_interrupts(
+        [{"id": "r1", "metadata": {"agent_framework": {"type": "workflow_request_info"}}}],
+        "cp-123",
+    )
+    assert attached[0]["metadata"]["agent_framework"]["checkpoint_id"] == "cp-123"
+    assert attached[0]["metadata"]["agent_framework"]["type"] == "workflow_request_info"
+
+    preserved = _attach_checkpoint_id_to_interrupts(attached, "cp-other")
+    assert preserved[0]["metadata"]["agent_framework"]["checkpoint_id"] == "cp-123"
+
+
+@pytest.mark.asyncio
+async def test_pause_checkpoint_id_ignores_competing_shared_latest() -> None:
+    """Prefer this runner's pause checkpoint over a newer shared get_latest() winner."""
+    from agent_framework import WorkflowCheckpoint
+
+    from agent_framework_ag_ui._workflow_run import (
+        _build_run_finished_with_checkpointed_interrupts,
+        _pause_checkpoint_id_for_interrupts,
+    )
+
+    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"},
+            )
+            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, response
+            await ctx.yield_output("done")  # 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_events = [
+        event
+        async for event in run_workflow_stream(
+            {"messages": [{"role": "user", "content": "go"}]},
+            workflow,
+            checkpoint_storage=storage,
+        )
+    ]
+    first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0]
+    interrupt_payload = _interrupts_from_run_finished(first_finished)
+    pause_id = interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"]
+
+    # Inject a newer shared checkpoint that does not cover this interrupt (other owner / stale).
+    competing = WorkflowCheckpoint(
+        workflow_name=workflow.name,
+        graph_signature_hash="competing",
+        pending_request_info_events={},
+        timestamp="9999-01-01T00:00:00+00:00",
+    )
+    await storage.save(competing)
+    latest = await storage.get_latest(workflow_name=workflow.name)
+    assert latest is not None
+    assert latest.checkpoint_id == competing.checkpoint_id
+
+    # baseline=None: treat as first advertisement after the pause run already advanced the runner id.
+    resolved = await _pause_checkpoint_id_for_interrupts(
+        workflow=workflow,
+        checkpoint_storage=storage,
+        interrupts=interrupt_payload,
+        baseline_checkpoint_id=None,
+    )
+    assert resolved == pause_id
+
+    rebuilt = await _build_run_finished_with_checkpointed_interrupts(
+        run_id="run-1",
+        thread_id="thread-1",
+        interrupts=interrupt_payload,
+        workflow=workflow,
+        checkpoint_storage=storage,
+        baseline_checkpoint_id=None,
+    )
+    rebuilt_interrupts = _interrupts_from_run_finished(rebuilt)
+    assert rebuilt_interrupts[0]["metadata"]["agent_framework"]["checkpoint_id"] == pause_id
+
+
+@pytest.mark.asyncio
+async def test_builder_checkpoint_storage_attaches_id_without_run_arg() -> None:
+    """WorkflowBuilder(checkpoint_storage=...) alone must still advertise pause checkpoint_id."""
+
+    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"},
+            )
+            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, response
+            await ctx.yield_output("done")  # 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()
+    # Deliberately omit checkpoint_storage= on the AG-UI entrypoint (builder path only).
+    events = [
+        event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)
+    ]
+    finished = [event for event in events if event.type == "RUN_FINISHED"][0]
+    interrupt_payload = _interrupts_from_run_finished(finished)
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+    assert checkpoints
+    assert interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] == checkpoints[-1].checkpoint_id
+
+
+@pytest.mark.asyncio
+async def test_resolve_pause_checkpoint_id_is_run_scoped_without_storage() -> None:
+    """Stale runner ids must not be advertised when baseline shows this run did not persist."""
+
+    class ApprovalExecutor(Executor):
+        def __init__(self) -> None:
+            super().__init__(id="approval_executor")
+
+        @handler
+        async def start(self, message: Any, ctx: WorkflowContext) -> None:
+            del message
+            await ctx.request_info("need-input", str, request_id="req-1")
+
+        @response_handler
+        async def handle(self, original_request: str, response: str, ctx: WorkflowContext) -> None:
+            del original_request, response
+            await ctx.yield_output("done")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+    workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build()
+    # Simulate a leftover id from a prior run on a workflow with no checkpoint storage.
+    workflow._runner._previous_checkpoint_id = "stale-from-prior-run"  # pyright: ignore[reportPrivateUsage]
+
+    from agent_framework_ag_ui._workflow_run import _pause_checkpoint_id_for_interrupts
+
+    interrupts = [{"id": "req-1", "value": "need-input"}]
+    assert (
+        await _pause_checkpoint_id_for_interrupts(
+            workflow=workflow,
+            checkpoint_storage=None,
+            interrupts=interrupts,
+            baseline_checkpoint_id="stale-from-prior-run",
+        )
+        is None
+    )
+    assert (
+        await _pause_checkpoint_id_for_interrupts(
+            workflow=workflow,
+            checkpoint_storage=None,
+            interrupts=interrupts,
+            baseline_checkpoint_id=None,
+        )
+        == "stale-from-prior-run"
+    )
+
+
 async def test_workflow_run_maps_custom_and_text_events():
     """Custom workflow events and yielded text are mapped to AG-UI events."""
 
@@ -682,6 +859,8 @@ async def handle_approval(self, original_request: Content, response: Content, ct
     )
     assert checkpoints, "expected the interrupted run to create a checkpoint"
     resume_checkpoint_id = checkpoints[-1].checkpoint_id
+    # Issue #8150: interrupt metadata must carry the pause checkpoint for multi-worker resume.
+    assert interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] == resume_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.