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 f14aa0e206..4f2668efee 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 @@ -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, + ) 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 f1cd16c959..4253289919 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 @@ -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. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 6e9f256b9b..47385df13a 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -1269,6 +1269,83 @@ def output_types(self) -> list[type[Any] | types.UnionType]: return list(output_types) + def get_last_checkpoint_id(self) -> str | None: + """Return the checkpoint id last persisted or restored by this workflow runner. + + Capture this value before ``run()`` when a host needs a run-scoped pause id: + after the run, only a different id indicates *this* run advanced the chain. + """ + checkpoint_id = self._runner._previous_checkpoint_id # pyright: ignore[reportPrivateUsage] + return str(checkpoint_id) if checkpoint_id is not None else None + + async def resolve_pause_checkpoint_id( + self, + request_ids: Collection[str], + *, + checkpoint_storage: CheckpointStorage | None = None, + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: str | None = None, + ) -> str | None: + """Resolve the persisted pause checkpoint that covers ``request_ids``. + + Prefers this runner's last-saved id when it advanced past ``baseline_checkpoint_id`` + (the id captured before ``run()``), over shared ``get_latest(workflow_name=...)``. + Storage precedence is the run argument, else the runner's effective + (runtime / builder) storage. When storage is available, candidates are accepted + only if their ``pending_request_info_events`` cover ``request_ids``. + + Args: + request_ids: Pending request_info ids that must be present on the checkpoint. + checkpoint_storage: Optional storage override for this lookup. + known_checkpoint_id: Fallback id (for example a cold-resume short-circuit). + baseline_checkpoint_id: Runner checkpoint id captured before the run that + produced these requests. When set (including ``None`` as “no prior id”), + the runner candidate is used only if ``get_last_checkpoint_id()`` differs. + + Returns: + A checkpoint id suitable for durable resume, or ``None`` when none is safe. + """ + ids = {str(request_id) for request_id in request_ids if request_id} + if not ids: + return None + + # Prefer explicit storage; otherwise load via public RunnerContext APIs + # (Protocol has no private `_get_effective_checkpoint_storage`). + storage = checkpoint_storage + use_context_storage = storage is None and self._runner.context.has_checkpointing() + + current = self.get_last_checkpoint_id() + # Run-scoped: do not advertise a pre-run leftover when this run did not persist. + runner_candidate: str | None = None + if current is not None and current != baseline_checkpoint_id: + runner_candidate = current + + candidates: list[str] = [] + if runner_candidate is not None: + candidates.append(runner_candidate) + if known_checkpoint_id is not None and known_checkpoint_id not in candidates: + candidates.append(str(known_checkpoint_id)) + + if storage is None and not use_context_storage: + # Without storage we cannot prove coverage; only advertise a run-scoped runner id. + return runner_candidate + + for candidate in candidates: + try: + if storage is not None: + checkpoint = await storage.load(candidate) + else: + checkpoint = await self._runner.context.load_checkpoint(candidate) + except Exception: # pragma: no cover - storage/type drift + logger.debug("Could not load pause checkpoint candidate %s", candidate, exc_info=True) + continue + if checkpoint is None: + continue + pending = checkpoint.pending_request_info_events or {} + if ids.issubset({str(key) for key in dict(pending)}): + return candidate + return None + async def cancel_pending_requests( self, request_ids: Collection[str], diff --git a/python/packages/core/tests/workflow/test_pause_checkpoint_resolve.py b/python/packages/core/tests/workflow/test_pause_checkpoint_resolve.py new file mode 100644 index 0000000000..5273bad577 --- /dev/null +++ b/python/packages/core/tests/workflow/test_pause_checkpoint_resolve.py @@ -0,0 +1,106 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for Workflow.resolve_pause_checkpoint_id / get_last_checkpoint_id.""" + +from typing import Any + +import pytest + +from agent_framework import ( + Content, + Executor, + InMemoryCheckpointStorage, + WorkflowBuilder, + WorkflowCheckpoint, + WorkflowContext, + handler, + response_handler, +) + + +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] + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_prefers_runner_over_shared_latest() -> None: + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + baseline = workflow.get_last_checkpoint_id() + + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + assert pause_id is not None + assert pause_id != baseline + + 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) + + resolved = await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + baseline_checkpoint_id=baseline, + ) + assert resolved == pause_id + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_requires_run_scoped_change_without_storage() -> None: + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor()).build() + workflow._runner._previous_checkpoint_id = "leftover" # pyright: ignore[reportPrivateUsage] + + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + baseline_checkpoint_id="leftover", + ) + is None + ) + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + baseline_checkpoint_id=None, + ) + == "leftover" + ) + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_uses_builder_storage() -> None: + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + baseline = workflow.get_last_checkpoint_id() + + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + resolved = await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + baseline_checkpoint_id=baseline, + ) + assert resolved == pause_id