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 539ddd21855..33d3bede30e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -525,6 +525,12 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: (``forwarded_props: {"checkpoint_id": ...}``), which restores the persisted workflow state instead of starting a fresh turn. + Resume also works when storage lives only on the resolved workflow + (``WorkflowBuilder(checkpoint_storage=...)``) and this wrapper / endpoint omits + ``checkpoint_storage``, so emitted pause ids remain round-trippable for hosts that + expose the workflow through ``AgentFrameworkWorkflow`` or + ``add_agent_framework_fastapi_endpoint`` without duplicating storage configuration. + Note: Checkpointing (the ``agent_framework`` workflow checkpoint mechanism) is independent from AG-UI Thread Snapshot persistence (``snapshot_store``). @@ -547,11 +553,6 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: 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)." - ) supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") request_owner = (snapshot_scope, str(supplied_thread_id) if supplied_thread_id is not None else None) @@ -561,6 +562,15 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: if interrupt.get("id") is not None } workflow = self._resolve_workflow(thread_id, snapshot_scope) + # Prefer wrapper/endpoint storage; otherwise allow builder/runtime storage on the + # resolved workflow so pause ids emitted without AG-UI storage still resume. + if checkpoint_id is not None and checkpoint_storage is None: + if not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] + raise ValueError( + "Resuming from a checkpoint requires checkpoint_storage to be configured on " + "AgentFrameworkWorkflow (or the AG-UI endpoint), or WorkflowBuilder " + "checkpoint storage on the workflow instance." + ) live_pending_events = await _pending_request_events(self.workflow) if self.workflow is not None else {} if self.workflow is not None and checkpoint_id is None: for request_event in live_pending_events.values(): @@ -572,9 +582,14 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: code="WORKFLOW_RESUME_NOT_FOUND", ) return - if checkpoint_id is not None and checkpoint_storage is not None: + if checkpoint_id is not None: try: - checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + else: + checkpoint = await workflow._runner.context.load_checkpoint(checkpoint_id) # pyright: ignore[reportPrivateUsage] + if checkpoint is None: + raise LookupError(f"checkpoint '{checkpoint_id}' was not found") except Exception as exc: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield 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 9ee70075f2b..9a5126cc813 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 @@ -61,6 +61,8 @@ logger = logging.getLogger(__name__) +_BASELINE_OMITTED = object() + _PUBLIC_WORKFLOW_ERROR_MESSAGE = "Workflow execution failed." @@ -130,7 +132,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) @@ -148,6 +153,99 @@ 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``. + + 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: Any = _BASELINE_OMITTED, +) -> str | None: + """Resolve the pause checkpoint for *this* run's interrupts via core. + + When ``baseline_checkpoint_id`` is omitted, core uses the baseline captured at + ``workflow.run()`` start. Pass ``None`` explicitly for short-circuit paths that + did not call ``run()`` and should advertise the current runner id when present. + """ + 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, + ) + kwargs: dict[str, Any] = { + "checkpoint_storage": checkpoint_storage, + "known_checkpoint_id": known_checkpoint_id, + } + if baseline_checkpoint_id is not _BASELINE_OMITTED: + kwargs["baseline_checkpoint_id"] = baseline_checkpoint_id + return await resolve_fn(_interrupt_request_ids(interrupts), **kwargs) + + +async def _interrupts_with_pause_checkpoint( + *, + interrupts: list[dict[str, Any]], + workflow: Workflow, + checkpoint_storage: CheckpointStorage | None, + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: Any = _BASELINE_OMITTED, +) -> list[dict[str, Any]]: + """Attach a run-scoped pause checkpoint id to interrupts when available.""" + if not interrupts: + return 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 _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) @@ -171,7 +269,9 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: async def _pending_request_events_from_checkpoint( checkpoint_id: str, - checkpoint_storage: CheckpointStorage, + checkpoint_storage: CheckpointStorage | None = None, + *, + workflow: Any | None = None, ) -> dict[str, Any]: """Read pending request_info events from a persisted checkpoint without restoring it. @@ -182,16 +282,34 @@ async def _pending_request_events_from_checkpoint( exposes them without running any executor ``on_checkpoint_restore`` hook; the single ``workflow.run(checkpoint_id=...)`` then performs the one real restore, so the restore -- and every custom restore hook -- runs exactly once per resume. + + ``checkpoint_storage`` is preferred when provided. Otherwise the workflow's + effective builder/runtime storage is used when ``has_checkpointing()`` is true, + so AG-UI can round-trip pause IDs emitted from ``WorkflowBuilder(checkpoint_storage=...)`` + without requiring a duplicate AG-UI storage argument. """ try: - checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + else: + context = getattr(getattr(workflow, "_runner", None), "context", None) + if context is None or not context.has_checkpointing(): + raise ValueError( + "Resuming a checkpoint with an AG-UI resume payload requires checkpoint_storage " + "(or WorkflowBuilder checkpoint storage on the workflow instance)." + ) + checkpoint = await context.load_checkpoint(checkpoint_id) + except ValueError: + raise 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) + if checkpoint is None: + return {} + return dict(checkpoint.pending_request_info_events or {}) def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None: @@ -1062,9 +1180,15 @@ async def run_workflow_stream( # 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: + # Prefer the explicit AG-UI storage argument; otherwise allow the workflow's + # builder/runtime storage so builder-emitted pause IDs remain round-trippable. + if checkpoint_storage is None and not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] 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) + pending_before_run = await _pending_request_events_from_checkpoint( + checkpoint_id, + checkpoint_storage, + workflow=workflow, + ) else: pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) @@ -1140,12 +1264,30 @@ 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) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=pending_interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + baseline_checkpoint_id=None, + ), + ) 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) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=pending_interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + baseline_checkpoint_id=None, + ), + ) return def _drain_open_message() -> list[TextMessageEndEvent]: @@ -1281,7 +1423,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 _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + ), + ) terminal_emitted = True elif state_value not in _TERMINAL_STATES: yield CustomEvent(name="status", value={"state": state_value}) @@ -1447,4 +1597,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 _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + ), + ) 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 8702adcd047..615d3ce4f1a 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 @@ -260,7 +260,7 @@ async def test_workflow_run_reads_checkpoint_id_from_camelcase_forwarded_props() async def test_workflow_resume_without_checkpoint_storage_raises() -> None: - """Requesting a checkpoint resume without configured storage should fail loudly.""" + """Requesting a checkpoint resume without any storage should fail loudly.""" workflow = _build_multi_superstep_workflow() agent = AgentFrameworkWorkflow(workflow=workflow) @@ -275,6 +275,39 @@ async def test_workflow_resume_without_checkpoint_storage_raises() -> None: ) +async def test_workflow_wrapper_resumes_builder_storage_without_agui_storage() -> None: + """Builder-owned storage must round-trip through AgentFrameworkWorkflow.run() without wrapper storage.""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow(storage) + # Host configures storage only on the builder; AG-UI wrapper / endpoint omit it. + agent = AgentFrameworkWorkflow(workflow=workflow) + + first_events = await _run( + agent, + {"thread_id": "thread-builder-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 builder-storage run to create at least one checkpoint" + resume_checkpoint_id = checkpoints[0].checkpoint_id + + resume_events = await _run( + agent, + { + "thread_id": "thread-builder-cp", + "messages": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, + }, + ) + resumed_types = [event.type for event in resume_events] + assert "RUN_FINISHED" in resumed_types + assert "RUN_ERROR" not in resumed_types + + 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() 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 d9cb44fe01c..fc79a1ae07f 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 @@ -98,6 +98,229 @@ 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._run_common import _build_run_finished_event + from agent_framework_ag_ui._workflow_run import ( + _interrupts_with_pause_checkpoint, + _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 + + # Workflow-owned baseline from the pause run should already allow advertising pause_id. + resolved = await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=storage, + interrupts=interrupt_payload, + ) + assert resolved == pause_id + + rebuilt = _build_run_finished_event( + "run-1", + "thread-1", + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupt_payload, + workflow=workflow, + checkpoint_storage=storage, + ), + ) + 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_builder_checkpoint_storage_resume_round_trips_without_agui_storage() -> None: + """Pause IDs from builder storage must resume with resume payload even without AG-UI storage.""" + + 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 + await ctx.yield_output(f"got:{response}") # 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() + + pause_events = [ + event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow) + ] + finished = [event for event in pause_events if event.type == "RUN_FINISHED"][0] + pause_id = _interrupts_from_run_finished(finished)[0]["metadata"]["agent_framework"]["checkpoint_id"] + + # Cold resume: omit AG-UI checkpoint_storage; rely on builder storage only. + resume_events = [ + event + async for event in run_workflow_stream( + { + "messages": [], + "resume": {"interrupts": [{"id": "req-1", "value": "ok"}]}, + "forwarded_props": {"checkpoint_id": pause_id}, + }, + workflow, + ) + ] + assert "RUN_ERROR" not in [event.type for event in resume_events] + assert "RUN_FINISHED" in [event.type for event in resume_events] + text = "".join(getattr(event, "delta", "") for event in resume_events if event.type == "TEXT_MESSAGE_CONTENT") + assert "got:ok" in text + + +@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"}] + workflow._run_baseline_checkpoint_id = "stale-from-prior-run" # pyright: ignore[reportPrivateUsage] + assert ( + await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=None, + interrupts=interrupts, + ) + 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.""" @@ -843,6 +1066,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 afbf45af11b..a24a2fade38 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -410,6 +410,12 @@ def __init__( # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). + # Captured at the start of each ``_run_core`` so ``resolve_pause_checkpoint_id`` + # can tell a newly persisted pause from a leftover / restored id. + self._run_baseline_checkpoint_id: str | None = None + self._restored_checkpoint_id: str | None = None + @property def status(self) -> WorkflowRunState: """Return the current run-level status of this workflow instance. @@ -907,6 +913,12 @@ async def _run_core( if checkpoint_storage is not None: self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + # Own the run boundary for pause-checkpoint resolution: capture the runner id + # before this run advances it, and remember an incoming restore so it is not + # re-advertised if a later pause save fails. + self._run_baseline_checkpoint_id = self.get_last_checkpoint_id() + self._restored_checkpoint_id = str(checkpoint_id) if checkpoint_id is not None else None + try: # Async validation: a fresh-message run is only allowed when the # runner context has fully drained from any prior run. If it still @@ -1308,6 +1320,93 @@ 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.""" + 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 | object = _MISSING, + ) -> str | None: + """Resolve the persisted pause checkpoint that covers ``request_ids``. + + Prefers this runner's last-saved id when it advanced past the run baseline + (captured automatically at ``run()`` start), 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``. + + An incoming restored checkpoint id is excluded so a failed pause save after + resume cannot re-advertise pre-response state. + + 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: Optional override for the pre-run runner id. When omitted, + uses the baseline captured by the most recent ``run()``. Pass ``None`` + explicitly to treat any current runner id as newly advanced. + + 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 + + if baseline_checkpoint_id is _MISSING: + baseline_checkpoint_id = self._run_baseline_checkpoint_id + + # 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() + excluded: set[str] = set() + if self._restored_checkpoint_id is not None: + excluded.add(self._restored_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 and current not in excluded: + 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 + and known_checkpoint_id not in excluded + ): + 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_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index bf5487ce0f8..fa96ee9d60f 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -26,6 +26,7 @@ Message, ResponseStream, WorkflowBuilder, + WorkflowCheckpoint, WorkflowCheckpointException, WorkflowContext, WorkflowConvergenceException, @@ -1769,3 +1770,108 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None # endregion + + +# --------------------------------------------------------------------------- +# Pause checkpoint resolution (AG-UI interrupt metadata) +# --------------------------------------------------------------------------- + +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() + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + assert pause_id is not None + + 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) + + # Workflow owns the run baseline captured at run() start. + resolved = await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + ) + 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] + workflow._run_baseline_checkpoint_id = "leftover" # pyright: ignore[reportPrivateUsage] + + assert await workflow.resolve_pause_checkpoint_id({"approval-1"}) 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_excludes_restored_checkpoint() -> None: + """After resume, the restored id must not be re-advertised if a new pause save fails.""" + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + + async for _ in workflow.run("go", stream=True): + pass + restored = workflow.get_last_checkpoint_id() + assert restored is not None + + # Simulate a resume that restored `restored` but did not persist a newer pause. + workflow._run_baseline_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + workflow._restored_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + workflow._runner._previous_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + known_checkpoint_id=restored, + ) + is None + ) + + +@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() + 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"}) + assert resolved == pause_id