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 6eeca4b3bbc..becb21d7a6e 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 @@ -266,7 +266,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. @@ -277,16 +279,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: @@ -1153,9 +1173,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) 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 3f163cda361..e419b2113ec 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 @@ -228,6 +228,51 @@ async def handle_approval(self, original_request: Content, response: Content, ct 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] + + 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."""