From 3866974e66a65880a5ad7ec361a8ddeac10dcc40 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Fri, 14 Aug 2026 13:38:06 +0900 Subject: [PATCH 1/3] Python: Align AG-UI run continuity Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aadad05d-2646-405f-8c62-a7a223abfc92 --- .../specs/004-python-function-calling-loop.md | 2 + python/packages/ag-ui/AGENTS.md | 3 + python/packages/ag-ui/README.md | 16 + .../ag-ui/agent_framework_ag_ui/_agent_run.py | 3 +- .../ag-ui/agent_framework_ag_ui/_tooling.py | 21 - .../ag-ui/agent_framework_ag_ui/_workflow.py | 123 ++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 501 +++++++++++++++++- .../ag-ui/tests/ag_ui/test_tooling.py | 39 -- 8 files changed, 641 insertions(+), 67 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 4027d9a8b1..eeb4180f7e 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -468,6 +468,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | | AG-UI standard approval payload | Agent and workflow tool approvals emit canonical `tool_call` interrupts. `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. Hosted approvals remain decision-only. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments`, `test_workflow_endpoint_emits_canonical_tool_approval_interrupt`, `test_workflow_endpoint_accepts_canonical_tool_approval_resume`, `test_workflow_endpoint_applies_canonical_approval_edited_args`, `test_workflow_endpoint_accepts_legacy_partial_approval_edits`, `test_workflow_endpoint_hosted_approval_rejects_argument_edits` | | AG-UI cancellation | A cancelled interrupt executes zero times and completes normally, including an identical retry during retained cancellation state; resolved siblings in the same complete resume still execute once. Workflow cancellation clears both runner correlation and the owning agent executor's pending request so later approvals remain resumable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_replayed_cancellation_completes_idempotently`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally`, `test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval` | +| AG-UI shared workflow interrupt ownership | A direct shared `Workflow` request-info interrupt can only be resolved or cancelled by the Snapshot Scope and AG-UI thread that created it. Ownership is recorded before the interrupt is exposed and survives checkpoint restoration; explicitly threaded callers fail closed when ownership is unavailable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_workflow_request_info_rejects_resume_from_different_thread`, `test_endpoint_workflow_request_info_rejects_resume_from_different_scope`, `test_endpoint_workflow_request_info_rejects_cancellation_from_different_thread`, `test_endpoint_workflow_request_info_remains_owned_after_client_disconnect`, `test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt`, `test_endpoint_workflow_checkpoint_resume_rejects_different_thread_after_restart` | | AG-UI approval retention and capacity | Pending authority expires automatically, indeterminate outcomes remain non-retryable until their safety window permits reclamation, and one trusted scope cannot consume another scope's occurrence quota. | `packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py::test_abandoned_pending_occurrence_expires_and_releases_capacity`, `test_indeterminate_occurrence_is_reclaimed_after_its_safety_window`, `test_capacity_is_enforced_per_trusted_scope` | | AG-UI local executor unavailable on resume | A claimed local occurrence whose executor disappeared releases its unstarted claim, reports temporary unavailability, and remains safely retryable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable` | | AG-UI forwarded execution interruption | A provider failure, cancellation, or stream close after forwarding an approval recovers the open occurrence as indeterminate when no idempotency key proves retry safety. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_hosted_approval_becomes_indeterminate_when_provider_stream_fails` | @@ -511,6 +512,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | AG-UI approval-time follow-up | The full grouped user-input pause remains in message history and emits no synthetic `TOOL_CALL_RESULT`. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_follow_up_group_remains_in_history_without_live_tool_result` | | AG-UI approval execution failure | A grouped executor failure becomes one deterministic terminal error result for the approved call. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_approval_execution_failure_emits_one_terminal_error_result` | | AG-UI no-approval path | Ordinary tool results do not gain an extra approval result event. | `packages/ag-ui/tests/ag_ui/test_approval_result_event.py::test_no_approval_path_emits_no_approval_specific_duplicate_result` | +| AG-UI client-tool request isolation | Client tool declarations are validated before use and remain request-scoped; a rejected collision or earlier successful request cannot change a later request's server-tool execution. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_failed_client_tool_collision_does_not_affect_next_request`, `test_endpoint_client_tools_do_not_persist_into_next_request` | | AG-UI `confirm_changes` snapshot | An accepted synthetic confirmation is replaced only when its original function call has a real result; rejection is cleaned explicitly, and missing accepted results remain inert. | `packages/ag-ui/tests/ag_ui/test_confirm_changes_snapshot.py` | | AG-UI malformed `confirm_changes` metadata | Non-list tool-call metadata and malformed argument JSON are ignored without guessing a target call. | `test_confirm_changes_target_ignores_non_list_tool_calls`, `test_confirm_changes_target_rejects_malformed_arguments_json` | | Compaction pair integrity | Adjacent and non-adjacent pairs, including assistant-embedded results and completed reused-id occurrences, remain atomic without pairing ambiguous or out-of-order ids. | `packages/core/tests/core/test_compaction.py::test_group_annotations_keep_tool_call_and_tool_result_atomic`, `test_group_annotations_include_reasoning_in_tool_call_group`, `test_group_annotations_pair_nonadjacent_function_result_by_call_id`, `test_group_annotations_pair_multiple_nonadjacent_results_with_declaration`, `test_group_annotations_pair_completed_reused_call_id_occurrences`, `test_group_annotations_close_assistant_embedded_result_before_reused_call_id`, `test_sliding_window_does_not_retain_orphan_result_after_assistant_embedded_result`, `test_sliding_window_keeps_reused_call_id_occurrences_atomic`, `test_group_annotations_do_not_pair_ambiguous_duplicate_call_ids` | diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index c1b6c32896..389f5038fc 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -40,6 +40,9 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - `_approval_lifecycle.py` is the sole owner of approval occurrence registration, trusted aliases, authority validation, claims, terminal outcomes, and retry deduplication. Runner code normalizes AG-UI protocol values and projects lifecycle outcomes but must not maintain a parallel pending-approval registry. +- Default stateless conversation history is client-controlled, including historical tool calls and results. Never + document conversational tool results as authorization or policy evidence; use deterministic server-side checks, + server-validated approvals, or scoped authoritative snapshots. - `confirm_changes` snapshot cleanup resolves the synthetic confirmation back to its original `function_call_id`; it must never concatenate unrelated tool results or record accepted changes without a matching real result. - SSE keepalive is endpoint-owned transport behavior configured through diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index a7688e59c4..7b93c64a5c 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -313,6 +313,22 @@ The `dependencies` parameter accepts any FastAPI dependency, enabling integratio For a complete authentication example, see [getting_started/server.py](getting_started/server.py). +### Conversation and Tool Result Trust + +In the default stateless mode, the AG-UI client sends the conversation history for each run. Treat that history as +untrusted input, including client-supplied `assistant` tool calls and `tool` results. A historical tool result is not +proof that the server emitted the matching call or executed the named backend tool. + +Do not use conversation history, tool results, or the model's decision to call a tool as an authorization, +entitlement, approval, or policy signal. Enforce security decisions deterministically in authenticated server code, +such as endpoint dependencies, tool middleware, or the server-validated human-in-the-loop approval flow. Tool +implementations must also authorize the current principal before accessing protected data or performing sensitive +actions. + +For applications that need server-authoritative thread history, configure scoped AG-UI Thread Snapshots. Snapshot +mode only accepts user turns and results for backend-issued tool calls when extending stored history. It complements +endpoint authentication and authorization; it does not replace them. + ## AG-UI Thread Snapshots AG-UI Thread Snapshot persistence is opt-in and disabled by default. Existing endpoints keep their current behavior diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index a5d461a932..703d75635a 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -72,7 +72,7 @@ from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id from ._message_adapters import normalize_agui_input_messages from ._predictive_state import PredictiveStateHandler -from ._tooling import collect_server_tools, merge_tools, register_additional_client_tools +from ._tooling import collect_server_tools, merge_tools from ._run_common import ( FlowState, _approval_interrupt_for_function_call, # type: ignore @@ -2303,7 +2303,6 @@ async def run_agent_stream( approval_snapshot_reconciliations: list[ApprovalSnapshotReconciliation] = [] client_tools = convert_agui_tools_to_agent_framework(input_data.get("tools")) server_tools = collect_server_tools(agent) - register_additional_client_tools(agent, client_tools) tools = merge_tools(server_tools, client_tools) approval_resume_messages, handled_resume_ids, cancelled_resume_ids, resume_error = ( _canonical_approval_resume_messages( diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_tooling.py b/python/packages/ag-ui/agent_framework_ag_ui/_tooling.py index 585bcb5c3e..9a4cdf1437 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_tooling.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_tooling.py @@ -7,7 +7,6 @@ import logging from typing import TYPE_CHECKING, Any -from agent_framework import BaseChatClient from agent_framework._tools import _append_unique_tools # pyright: ignore[reportPrivateUsage] if TYPE_CHECKING: @@ -71,26 +70,6 @@ def collect_server_tools(agent: SupportsAgentRun) -> list[Any]: return server_tools -def register_additional_client_tools(agent: SupportsAgentRun, client_tools: list[Any] | None) -> None: - """Register client tools as additional declaration-only tools to avoid server execution. - - Args: - agent: Agent instance to register tools on. Works with Agent - or any agent with a client attribute. - client_tools: List of client tools to register. - """ - if not client_tools: - return - - client = getattr(agent, "client", None) - if client is None: - return - - if isinstance(client, BaseChatClient) and client.function_invocation_configuration is not None: # type: ignore[attr-defined] - client.function_invocation_configuration["additional_tools"] = client_tools # type: ignore[attr-defined] - logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)") - - def _has_approval_tools(tools: list[Any]) -> bool: """Check if any tools require approval.""" return any(getattr(tool, "approval_mode", None) == "always_require" for tool in tools) 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 34bcc8b699..6f1eb6a1ed 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -14,6 +14,7 @@ MessagesSnapshotEvent, RunErrorEvent, RunFinishedEvent, + RunStartedEvent, StateSnapshotEvent, TextMessageContentEvent, TextMessageEndEvent, @@ -30,6 +31,7 @@ from ._run_common import ( _cancelled_resume_interrupt_ids, _extract_resume_payload, + _normalize_resume_interrupts, _reconstruct_messages_from_thread_snapshot, ) from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts @@ -45,6 +47,9 @@ logger = logging.getLogger(__name__) WorkflowFactory = Callable[[str], Workflow] +WorkflowRequestOwner = tuple[str | None, str | None] + +_CHECKPOINT_REQUEST_OWNER_KEY = "ag_ui_workflow_request_owner" def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: @@ -58,6 +63,42 @@ def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: return str(checkpoint_id) +def _checkpoint_request_owner(metadata: dict[str, Any]) -> WorkflowRequestOwner | None: + """Read AG-UI workflow request ownership from checkpoint metadata.""" + raw_owner = metadata.get(_CHECKPOINT_REQUEST_OWNER_KEY) + if not isinstance(raw_owner, dict): + return None + snapshot_scope = raw_owner.get("snapshot_scope") + thread_id = raw_owner.get("thread_id") + if snapshot_scope is not None and not isinstance(snapshot_scope, str): + return None + if thread_id is not None and not isinstance(thread_id, str): + return None + return snapshot_scope, thread_id + + +async def _persist_checkpoint_request_owner( + *, + checkpoint_storage: CheckpointStorage, + workflow: Workflow, + pending_interrupt_ids: set[str], + request_owner: WorkflowRequestOwner, +) -> None: + """Persist the owner of pending workflow requests on the latest checkpoint.""" + checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow.name) + if checkpoint is None: + return + checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} + if not pending_interrupt_ids.issubset(checkpoint_pending_ids): + return + checkpoint.metadata = dict(checkpoint.metadata) + checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = { + "snapshot_scope": request_owner[0], + "thread_id": request_owner[1], + } + await checkpoint_storage.save(checkpoint) + + class _WorkflowSnapshotBuilder: """Capture replayable workflow protocol output without retaining raw events.""" @@ -225,6 +266,7 @@ def __init__( self.workflow = workflow self._workflow_factory = workflow_factory + self._shared_workflow_request_owners: dict[str, WorkflowRequestOwner] = {} # Cache keyed by (snapshot_scope, thread_id): the Snapshot Scope is the # authorization boundary for both snapshots and in-memory workflow_factory # instances, so the same thread id under different scopes must never share @@ -301,6 +343,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: scope=snapshot_scope, thread_id=thread_id, ) + stored_snapshot = snapshot_session.stored checkpoint_storage = self.checkpoint_storage checkpoint_id = _checkpoint_id_from_input(input_data) @@ -310,6 +353,50 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: "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) + resume_interrupt_ids = { + str(interrupt["id"]) + for interrupt in _normalize_resume_interrupts(resume_payload) + if interrupt.get("id") is not None + } + if self.workflow is not None and stored_snapshot is not None: + for interrupt in stored_snapshot.interrupt or []: + interrupt_id = interrupt.get("id") + if interrupt_id is not None: + self._shared_workflow_request_owners[str(interrupt_id)] = request_owner + if self.workflow is not None and checkpoint_id is not None and checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} + checkpoint_owner = _checkpoint_request_owner(checkpoint.metadata) + if checkpoint_pending_ids and checkpoint_owner is None and request_owner != (None, None): + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", + code="WORKFLOW_RESUME_NOT_FOUND", + ) + return + if checkpoint_owner is not None and checkpoint_owner != request_owner: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", + code="WORKFLOW_RESUME_NOT_FOUND", + ) + return + if checkpoint_owner is not None: + for interrupt_id in checkpoint_pending_ids: + self._shared_workflow_request_owners[interrupt_id] = checkpoint_owner + if self.workflow is not None: + for interrupt_id in resume_interrupt_ids: + owner = self._shared_workflow_request_owners.get(interrupt_id) + if owner != request_owner and (owner is not None or request_owner != (None, None)): + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"No pending interrupt found for resume interruptId '{interrupt_id}'.", + code="WORKFLOW_RESUME_NOT_FOUND", + ) + return + # 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_session.enabled and not raw_messages and resume_payload is None: @@ -319,7 +406,6 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: # Seed follow-up turns so the workflow runs with the full persisted thread # history instead of just the latest request messages. - stored_snapshot = snapshot_session.stored if stored_snapshot is not None and resume_payload is None: raw_messages = _reconstruct_messages_from_thread_snapshot( stored_messages=stored_snapshot.messages, @@ -350,11 +436,46 @@ 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 + pending_interrupt_ids: set[str] = set() 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 ( + self.workflow is not None + and isinstance(event, ToolCallStartEvent) + and event.tool_call_name == "request_info" + ): + interrupt_id = str(event.tool_call_id) + self._shared_workflow_request_owners[interrupt_id] = request_owner + if checkpoint_storage is not None: + await _persist_checkpoint_request_owner( + checkpoint_storage=checkpoint_storage, + workflow=workflow, + pending_interrupt_ids={interrupt_id}, + request_owner=request_owner, + ) + if isinstance(event, RunFinishedEvent): + outcome = getattr(event, "outcome", None) + if getattr(outcome, "type", None) == "interrupt": + pending_interrupt_ids = { + str(interrupt.id) + for interrupt in getattr(outcome, "interrupts", []) + if getattr(interrupt, "id", None) + } + if self.workflow is not None: + for interrupt_id in resume_interrupt_ids: + self._shared_workflow_request_owners.pop(interrupt_id, None) + for interrupt_id in pending_interrupt_ids: + self._shared_workflow_request_owners[interrupt_id] = request_owner + if checkpoint_storage is not None and pending_interrupt_ids: + await _persist_checkpoint_request_owner( + checkpoint_storage=checkpoint_storage, + workflow=workflow, + pending_interrupt_ids=pending_interrupt_ids, + request_owner=request_owner, + ) if isinstance(event, RunErrorEvent): run_error_emitted = True if getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED": diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index da84d6d2eb..09c88b0172 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -7,7 +7,7 @@ import subprocess import sys from collections import Counter -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable from inspect import signature from typing import Any, cast @@ -58,6 +58,50 @@ def _decode_sse_events(response: Any) -> list[dict[str, Any]]: return [json.loads(line[6:]) for line in content.splitlines() if line.startswith("data: ")] +async def _post_until_sse_event_then_disconnect( + app: FastAPI, + path: str, + payload: dict[str, Any], + *, + event_type: str, +) -> None: + """Run one ASGI request until an SSE event is sent, then disconnect the client.""" + request_sent = False + disconnect = asyncio.Event() + body = json.dumps(payload).encode() + + async def receive() -> dict[str, Any]: + nonlocal request_sent + if not request_sent: + request_sent = True + return {"type": "http.request", "body": body, "more_body": False} + await disconnect.wait() + return {"type": "http.disconnect"} + + async def send(message: dict[str, Any]) -> None: + if message["type"] != "http.response.body": + return + chunk = message.get("body", b"") + if isinstance(chunk, bytes) and f'"type":"{event_type}"'.encode() in chunk: + disconnect.set() + + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "root_path": "", + "headers": [(b"content-type", b"application/json"), (b"host", b"testserver")], + "client": ("testclient", 50000), + "server": ("testserver", 80), + } + await app(scope, receive, send) + + def _run_finished_interrupts(event: dict[str, Any]) -> list[dict[str, Any]]: """Return canonical interrupts from an SSE RUN_FINISHED event.""" assert "interrupt" not in event @@ -88,6 +132,57 @@ def _latest_messages_snapshot(response: Any) -> list[dict[str, Any]]: return snapshots[-1] +def _build_server_guard_endpoint( + streaming_chat_client_stub: Any, + *, + first_provider_response: str, + server_tool_enabled: bool, +) -> tuple[TestClient, Agent, FunctionTool, list[str]]: + server_executions: list[str] = [] + provider_calls = 0 + + def server_guard() -> str: + server_executions.append("executed") + return "server guard executed" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal provider_calls + del messages, options, kwargs + provider_calls += 1 + if provider_calls == 1 and first_provider_response == "text": + yield ChatResponseUpdate(contents=[Content.from_text(text="Client tools accepted.")], role="assistant") + return + function_call_number = 1 if first_provider_response == "function_call" else 2 + if provider_calls == function_call_number: + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + call_id="call-server-guard", + name="server_guard", + arguments={}, + ) + ], + role="assistant", + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + + server_tool = FunctionTool(name="server_guard", description="Server guard", func=server_guard) + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + tools=[server_tool] if server_tool_enabled else [], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, agent, path="/agent") + return TestClient(app), agent, server_tool, server_executions + + @pytest.fixture def build_chat_client(streaming_chat_client_stub, stream_from_updates_fixture): """Create a typed chat client stub for endpoint tests.""" @@ -128,6 +223,96 @@ async def test_add_endpoint_with_wrapped_agent(build_chat_client): assert response.headers["content-type"] == "text/event-stream; charset=utf-8" +async def test_endpoint_failed_client_tool_collision_does_not_affect_next_request( + streaming_chat_client_stub, +) -> None: + """A rejected client-tool declaration cannot suppress a server tool on the next request.""" + client, _, _, server_executions = _build_server_guard_endpoint( + streaming_chat_client_stub, + first_provider_response="function_call", + server_tool_enabled=True, + ) + + with client: + collision_response = client.post( + "/agent", + json={ + "runId": "run-collision", + "threadId": "attacker-thread", + "messages": [{"role": "user", "content": "Declare a colliding client tool"}], + "tools": [ + { + "name": "server_guard", + "description": "Client-controlled collision", + "parameters": {"type": "object", "properties": {}}, + } + ], + }, + ) + collision_events = _decode_sse_events(collision_response) + assert [event for event in collision_events if event.get("type") == "RUN_ERROR"] + + next_response = client.post( + "/agent", + json={ + "runId": "run-next", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Run the server guard"}], + }, + ) + + assert next_response.status_code == 200 + next_events = _decode_sse_events(next_response) + assert not [event for event in next_events if event.get("type") == "RUN_ERROR"] + assert server_executions == ["executed"] + assert "Done." in [event["delta"] for event in next_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + +async def test_endpoint_client_tools_do_not_persist_into_next_request( + streaming_chat_client_stub, +) -> None: + """Client tool declarations are request-scoped on a shared agent.""" + client, agent, server_tool, server_executions = _build_server_guard_endpoint( + streaming_chat_client_stub, + first_provider_response="text", + server_tool_enabled=False, + ) + + with client: + client_tool_response = client.post( + "/agent", + json={ + "runId": "run-client-tools", + "threadId": "first-thread", + "messages": [{"role": "user", "content": "Use a client tool"}], + "tools": [ + { + "name": "server_guard", + "description": "Client-side guard", + "parameters": {"type": "object", "properties": {}}, + } + ], + }, + ) + assert not [event for event in _decode_sse_events(client_tool_response) if event.get("type") == "RUN_ERROR"] + agent.default_options["tools"] = [server_tool] + + next_response = client.post( + "/agent", + json={ + "runId": "run-next", + "threadId": "second-thread", + "messages": [{"role": "user", "content": "Run the server guard"}], + }, + ) + + assert next_response.status_code == 200 + next_events = _decode_sse_events(next_response) + assert not [event for event in next_events if event.get("type") == "RUN_ERROR"] + assert server_executions == ["executed"] + assert "Done." in [event["delta"] for event in next_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + async def test_add_endpoint_with_workflow_protocol(): """Test adding endpoint with native Workflow support.""" @@ -3553,7 +3738,7 @@ def record_city(city: str) -> str: assert "outcome" not in hydrate_events[-1] -def _build_workflow_request_info_app() -> FastAPI: +def _build_flight_choice_workflow() -> Any: class FlightChoiceExecutor(Executor): def __init__(self) -> None: super().__init__(id="flight_choice") @@ -3572,9 +3757,20 @@ async def handle_choice(self, original_request: dict, response: dict, ctx: Workf del original_request await ctx.yield_output(f"Booked {response['airline']}") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + return WorkflowBuilder(start_executor=FlightChoiceExecutor()).build() + + +def _build_workflow_request_info_app( + *, + snapshot_scope_resolver: Callable[[AGUIRequest], str] | None = None, +) -> FastAPI: app = FastAPI() - workflow = WorkflowBuilder(start_executor=FlightChoiceExecutor()).build() - add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow") + add_agent_framework_fastapi_endpoint( + app, + _build_flight_choice_workflow(), + path="/workflow", + snapshot_scope_resolver=snapshot_scope_resolver, + ) return app @@ -3630,6 +3826,303 @@ async def test_endpoint_workflow_request_info_emits_canonical_interrupt_and_resu assert "outcome" not in [event for event in resume_events if event.get("type") == "RUN_FINISHED"][-1] +async def test_endpoint_workflow_request_info_rejects_resume_from_different_thread(): + """A workflow interrupt can only be resumed by the AG-UI thread that created it.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + assert attacker_response.status_code == 200 + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + victim_response = client.post( + "/workflow", + json={ + "runId": "run-victim-resume", + "threadId": "victim-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "United"}, + } + ], + }, + ) + + assert victim_response.status_code == 200 + victim_events = _decode_sse_events(victim_response) + assert not [event for event in victim_events if event.get("type") == "RUN_ERROR"] + text_deltas = [event["delta"] for event in victim_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert "Booked United" in text_deltas + + +async def test_endpoint_workflow_request_info_rejects_resume_from_different_scope(): + """A workflow interrupt can only be resumed within the Snapshot Scope that created it.""" + + def resolve_scope(request: AGUIRequest) -> str: + forwarded_props = request.forwarded_props + assert forwarded_props is not None + tenant = forwarded_props["tenant"] + assert isinstance(tenant, str) + return tenant + + app = _build_workflow_request_info_app(snapshot_scope_resolver=resolve_scope) + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "shared-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + assert pause_response.status_code == 200 + + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "shared-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + "forwardedProps": {"tenant": "tenant-b"}, + }, + ) + + assert attacker_response.status_code == 200 + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + +async def test_endpoint_workflow_request_info_rejects_cancellation_from_different_thread(): + """A different AG-UI thread cannot cancel another thread's workflow interrupt.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker-cancel", + "threadId": "attacker-thread", + "messages": [], + "resume": [{"interruptId": "flight-choice", "status": "cancelled"}], + }, + ) + + assert attacker_response.status_code == 200 + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + + victim_response = client.post( + "/workflow", + json={ + "runId": "run-victim-resume", + "threadId": "victim-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "United"}, + } + ], + }, + ) + victim_events = _decode_sse_events(victim_response) + assert not [event for event in victim_events if event.get("type") == "RUN_ERROR"] + text_deltas = [event["delta"] for event in victim_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert "Booked United" in text_deltas + + +async def test_endpoint_workflow_request_info_remains_owned_after_client_disconnect(): + """Disconnecting after an interrupt is visible does not release its thread ownership.""" + app = _build_workflow_request_info_app() + + await _post_until_sse_event_then_disconnect( + app, + "/workflow", + { + "runId": "run-pause", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + event_type="TOOL_CALL_END", + ) + + with TestClient(app) as client: + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + + +async def test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt(): + """An explicitly threaded endpoint cannot claim pending state created outside that endpoint.""" + workflow = _build_flight_choice_workflow() + _ = [ + event + async for event in workflow.run( + message=[Message(role="user", contents=[Content.from_text(text="Book me a flight")])], + stream=True, + ) + ] + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow") + + with TestClient(app) as client: + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + + +async def test_endpoint_workflow_checkpoint_resume_rejects_different_thread_after_restart(): + """Checkpoint restore preserves the AG-UI thread that owns a pending workflow interrupt.""" + storage = InMemoryCheckpointStorage() + first_app = FastAPI() + first_workflow = _build_flight_choice_workflow() + add_agent_framework_fastapi_endpoint( + first_app, + first_workflow, + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(first_app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + second_app = FastAPI() + add_agent_framework_fastapi_endpoint( + second_app, + _build_flight_choice_workflow(), + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(second_app) as client: + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + "forwardedProps": {"checkpointId": checkpoint_id}, + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + async def test_endpoint_workflow_request_info_cancelled_resume_completes_normally(): """Cancelled workflow resumes complete without output and do not wedge the next turn.""" app = _build_workflow_request_info_app() diff --git a/python/packages/ag-ui/tests/ag_ui/test_tooling.py b/python/packages/ag-ui/tests/ag_ui/test_tooling.py index bcdf285b0f..f62c55c7df 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_tooling.py +++ b/python/packages/ag-ui/tests/ag_ui/test_tooling.py @@ -9,7 +9,6 @@ from agent_framework_ag_ui._tooling import ( collect_server_tools, merge_tools, - register_additional_client_tools, ) @@ -52,21 +51,6 @@ def test_merge_tools_filters_duplicates() -> None: merge_tools(server, client) -def test_register_additional_client_tools_assigns_when_configured() -> None: - """register_additional_client_tools should set additional_tools on the chat client.""" - from agent_framework import BaseChatClient, normalize_function_invocation_configuration - - mock_chat_client = MagicMock(spec=BaseChatClient) - mock_chat_client.function_invocation_configuration = normalize_function_invocation_configuration(None) - - agent = Agent(client=mock_chat_client) - - tools = [DummyTool("x")] - register_additional_client_tools(cast(Any, agent), tools) - - assert mock_chat_client.function_invocation_configuration["additional_tools"] == tools - - def test_collect_server_tools_includes_mcp_tools_when_connected() -> None: """MCP tool functions should be included when the MCP tool is connected.""" mcp_function1 = DummyTool("mcp_function_1") @@ -156,29 +140,6 @@ class MockAgent: assert tools == [] -def test_register_additional_client_tools_no_tools() -> None: - """register_additional_client_tools does nothing with None tools.""" - mock_chat_client = MagicMock() - agent = Agent(client=mock_chat_client) - - # Should not raise - register_additional_client_tools(agent, None) - - -def test_register_additional_client_tools_no_chat_client() -> None: - """register_additional_client_tools does nothing when agent has no client.""" - from agent_framework_ag_ui._tooling import register_additional_client_tools - - class MockAgent: - pass - - agent = MockAgent() - tools = [DummyTool("x")] - - # Should not raise - register_additional_client_tools(cast(Any, agent), tools) - - def test_merge_tools_no_client_tools() -> None: """merge_tools returns None when no client tools.""" server = [DummyTool("a")] From 6e4fbbc24bdf53420c7ebb522ff1aa2b9d9167ff Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Fri, 14 Aug 2026 14:39:04 +0900 Subject: [PATCH 2/3] Python: Refine AG-UI continuation ownership Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aadad05d-2646-405f-8c62-a7a223abfc92 --- .../specs/004-python-function-calling-loop.md | 2 +- .../ag-ui/agent_framework_ag_ui/_workflow.py | 110 ++-------- .../ag-ui/tests/ag_ui/test_endpoint.py | 193 +++++++++++++++++- 3 files changed, 204 insertions(+), 101 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index eeb4180f7e..8f90781a67 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -468,7 +468,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | | AG-UI standard approval payload | Agent and workflow tool approvals emit canonical `tool_call` interrupts. `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. Hosted approvals remain decision-only. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments`, `test_workflow_endpoint_emits_canonical_tool_approval_interrupt`, `test_workflow_endpoint_accepts_canonical_tool_approval_resume`, `test_workflow_endpoint_applies_canonical_approval_edited_args`, `test_workflow_endpoint_accepts_legacy_partial_approval_edits`, `test_workflow_endpoint_hosted_approval_rejects_argument_edits` | | AG-UI cancellation | A cancelled interrupt executes zero times and completes normally, including an identical retry during retained cancellation state; resolved siblings in the same complete resume still execute once. Workflow cancellation clears both runner correlation and the owning agent executor's pending request so later approvals remain resumable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_replayed_cancellation_completes_idempotently`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally`, `test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval` | -| AG-UI shared workflow interrupt ownership | A direct shared `Workflow` request-info interrupt can only be resolved or cancelled by the Snapshot Scope and AG-UI thread that created it. Ownership is recorded before the interrupt is exposed and survives checkpoint restoration; explicitly threaded callers fail closed when ownership is unavailable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_workflow_request_info_rejects_resume_from_different_thread`, `test_endpoint_workflow_request_info_rejects_resume_from_different_scope`, `test_endpoint_workflow_request_info_rejects_cancellation_from_different_thread`, `test_endpoint_workflow_request_info_remains_owned_after_client_disconnect`, `test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt`, `test_endpoint_workflow_checkpoint_resume_rejects_different_thread_after_restart` | +| AG-UI shared workflow interrupt ownership | A direct shared `Workflow` request-info interrupt can only be resolved or cancelled by the Snapshot Scope and AG-UI thread that created it. Ownership follows the authoritative pending request occurrence, and explicitly threaded cold checkpoint resumes fail closed when ownership is unavailable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_workflow_request_info_rejects_resume_from_different_thread`, `test_endpoint_workflow_request_info_rejects_resume_from_different_scope`, `test_endpoint_workflow_request_info_rejects_cancellation_from_different_thread`, `test_endpoint_workflow_request_info_remains_owned_after_client_disconnect`, `test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt`, `test_endpoint_workflow_checkpoint_resume_rejects_threaded_resume_after_restart` | | AG-UI approval retention and capacity | Pending authority expires automatically, indeterminate outcomes remain non-retryable until their safety window permits reclamation, and one trusted scope cannot consume another scope's occurrence quota. | `packages/ag-ui/tests/ag_ui/test_approval_lifecycle.py::test_abandoned_pending_occurrence_expires_and_releases_capacity`, `test_indeterminate_occurrence_is_reclaimed_after_its_safety_window`, `test_capacity_is_enforced_per_trusted_scope` | | AG-UI local executor unavailable on resume | A claimed local occurrence whose executor disappeared releases its unstarted claim, reports temporary unavailability, and remains safely retryable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable` | | AG-UI forwarded execution interruption | A provider failure, cancellation, or stream close after forwarding an approval recovers the open occurrence as indeterminate when no idempotency key proves retry safety. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_hosted_approval_becomes_indeterminate_when_provider_stream_fails` | 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 6f1eb6a1ed..b1893e8af4 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -42,14 +42,14 @@ AGUIThreadSnapshotStore, ) from ._utils import generate_event_id, make_json_safe -from ._workflow_run import run_workflow_stream +from ._workflow_run import _pending_request_events, run_workflow_stream # pyright: ignore[reportPrivateUsage] logger = logging.getLogger(__name__) WorkflowFactory = Callable[[str], Workflow] WorkflowRequestOwner = tuple[str | None, str | None] -_CHECKPOINT_REQUEST_OWNER_KEY = "ag_ui_workflow_request_owner" +_REQUEST_OWNER_ATTRIBUTE = "_ag_ui_request_owner" def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: @@ -63,42 +63,6 @@ def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: return str(checkpoint_id) -def _checkpoint_request_owner(metadata: dict[str, Any]) -> WorkflowRequestOwner | None: - """Read AG-UI workflow request ownership from checkpoint metadata.""" - raw_owner = metadata.get(_CHECKPOINT_REQUEST_OWNER_KEY) - if not isinstance(raw_owner, dict): - return None - snapshot_scope = raw_owner.get("snapshot_scope") - thread_id = raw_owner.get("thread_id") - if snapshot_scope is not None and not isinstance(snapshot_scope, str): - return None - if thread_id is not None and not isinstance(thread_id, str): - return None - return snapshot_scope, thread_id - - -async def _persist_checkpoint_request_owner( - *, - checkpoint_storage: CheckpointStorage, - workflow: Workflow, - pending_interrupt_ids: set[str], - request_owner: WorkflowRequestOwner, -) -> None: - """Persist the owner of pending workflow requests on the latest checkpoint.""" - checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow.name) - if checkpoint is None: - return - checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} - if not pending_interrupt_ids.issubset(checkpoint_pending_ids): - return - checkpoint.metadata = dict(checkpoint.metadata) - checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = { - "snapshot_scope": request_owner[0], - "thread_id": request_owner[1], - } - await checkpoint_storage.save(checkpoint) - - class _WorkflowSnapshotBuilder: """Capture replayable workflow protocol output without retaining raw events.""" @@ -266,7 +230,6 @@ def __init__( self.workflow = workflow self._workflow_factory = workflow_factory - self._shared_workflow_request_owners: dict[str, WorkflowRequestOwner] = {} # Cache keyed by (snapshot_scope, thread_id): the Snapshot Scope is the # authorization boundary for both snapshots and in-memory workflow_factory # instances, so the same thread id under different scopes must never share @@ -360,35 +323,24 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: for interrupt in _normalize_resume_interrupts(resume_payload) if interrupt.get("id") is not None } - if self.workflow is not None and stored_snapshot is not None: - for interrupt in stored_snapshot.interrupt or []: - interrupt_id = interrupt.get("id") - if interrupt_id is not None: - self._shared_workflow_request_owners[str(interrupt_id)] = request_owner + 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 not None and checkpoint_storage is not None: checkpoint = await checkpoint_storage.load(checkpoint_id) checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} - checkpoint_owner = _checkpoint_request_owner(checkpoint.metadata) - if checkpoint_pending_ids and checkpoint_owner is None and request_owner != (None, None): - yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - yield RunErrorEvent( - message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", - code="WORKFLOW_RESUME_NOT_FOUND", - ) - return - if checkpoint_owner is not None and checkpoint_owner != request_owner: - yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - yield RunErrorEvent( - message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", - code="WORKFLOW_RESUME_NOT_FOUND", - ) - return - if checkpoint_owner is not None: - for interrupt_id in checkpoint_pending_ids: - self._shared_workflow_request_owners[interrupt_id] = checkpoint_owner + for interrupt_id in checkpoint_pending_ids: + request_event = live_pending_events.get(interrupt_id) + owner = getattr(request_event, _REQUEST_OWNER_ATTRIBUTE, None) + if owner != request_owner and (owner is not None or request_owner != (None, None)): + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", + code="WORKFLOW_RESUME_NOT_FOUND", + ) + return if self.workflow is not None: for interrupt_id in resume_interrupt_ids: - owner = self._shared_workflow_request_owners.get(interrupt_id) + request_event = live_pending_events.get(interrupt_id) + owner = getattr(request_event, _REQUEST_OWNER_ATTRIBUTE, None) if owner != request_owner and (owner is not None or request_owner != (None, None)): yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield RunErrorEvent( @@ -436,7 +388,6 @@ 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 - pending_interrupt_ids: set[str] = set() async for event in run_workflow_stream( input_data, workflow, checkpoint_storage=checkpoint_storage, checkpoint_id=checkpoint_id ): @@ -448,34 +399,9 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: and event.tool_call_name == "request_info" ): interrupt_id = str(event.tool_call_id) - self._shared_workflow_request_owners[interrupt_id] = request_owner - if checkpoint_storage is not None: - await _persist_checkpoint_request_owner( - checkpoint_storage=checkpoint_storage, - workflow=workflow, - pending_interrupt_ids={interrupt_id}, - request_owner=request_owner, - ) - if isinstance(event, RunFinishedEvent): - outcome = getattr(event, "outcome", None) - if getattr(outcome, "type", None) == "interrupt": - pending_interrupt_ids = { - str(interrupt.id) - for interrupt in getattr(outcome, "interrupts", []) - if getattr(interrupt, "id", None) - } - if self.workflow is not None: - for interrupt_id in resume_interrupt_ids: - self._shared_workflow_request_owners.pop(interrupt_id, None) - for interrupt_id in pending_interrupt_ids: - self._shared_workflow_request_owners[interrupt_id] = request_owner - if checkpoint_storage is not None and pending_interrupt_ids: - await _persist_checkpoint_request_owner( - checkpoint_storage=checkpoint_storage, - workflow=workflow, - pending_interrupt_ids=pending_interrupt_ids, - request_owner=request_owner, - ) + pending_event = (await _pending_request_events(workflow)).get(interrupt_id) + if pending_event is not None: + setattr(pending_event, _REQUEST_OWNER_ATTRIBUTE, request_owner) if isinstance(event, RunErrorEvent): run_error_emitted = True if getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED": diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 09c88b0172..433fe1ef55 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -40,6 +40,8 @@ from fastapi import FastAPI, Header, HTTPException from fastapi.params import Depends from fastapi.testclient import TestClient +from starlette.types import Message as ASGIMessage +from starlette.types import Receive, Scope, Send from agent_framework_ag_ui import ( AGUIRequest, @@ -70,7 +72,7 @@ async def _post_until_sse_event_then_disconnect( disconnect = asyncio.Event() body = json.dumps(payload).encode() - async def receive() -> dict[str, Any]: + async def receive() -> ASGIMessage: nonlocal request_sent if not request_sent: request_sent = True @@ -78,14 +80,14 @@ async def receive() -> dict[str, Any]: await disconnect.wait() return {"type": "http.disconnect"} - async def send(message: dict[str, Any]) -> None: + async def send(message: ASGIMessage) -> None: if message["type"] != "http.response.body": return chunk = message.get("body", b"") if isinstance(chunk, bytes) and f'"type":"{event_type}"'.encode() in chunk: disconnect.set() - scope = { + scope: Scope = { "type": "http", "asgi": {"version": "3.0"}, "http_version": "1.1", @@ -99,7 +101,7 @@ async def send(message: dict[str, Any]) -> None: "client": ("testclient", 50000), "server": ("testserver", 80), } - await app(scope, receive, send) + await asyncio.wait_for(app(scope, cast(Receive, receive), cast(Send, send)), timeout=5) def _run_finished_interrupts(event: dict[str, Any]) -> list[dict[str, Any]]: @@ -3936,6 +3938,108 @@ def resolve_scope(request: AGUIRequest) -> str: assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] +async def test_endpoint_workflow_request_info_stale_snapshot_does_not_replace_live_owner(): + """A stale scoped snapshot cannot replace a newer live interrupt owner.""" + + def resolve_scope(request: AGUIRequest) -> str: + forwarded_props = request.forwarded_props + assert forwarded_props is not None + tenant = forwarded_props["tenant"] + assert isinstance(tenant, str) + return tenant + + store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + _build_flight_choice_workflow(), + path="/workflow", + snapshot_store=store, + snapshot_scope_resolver=resolve_scope, + ) + + with TestClient(app) as client: + first_pause = client.post( + "/workflow", + json={ + "runId": "run-first-pause", + "threadId": "shared-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + assert first_pause.status_code == 200 + stale_snapshot = await store.get(scope="tenant-a", thread_id="shared-thread") + assert stale_snapshot is not None + + first_cancel = client.post( + "/workflow", + json={ + "runId": "run-first-cancel", + "threadId": "shared-thread", + "messages": [], + "resume": [{"interruptId": "flight-choice", "status": "cancelled"}], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + assert first_cancel.status_code == 200 + + second_pause = client.post( + "/workflow", + json={ + "runId": "run-second-pause", + "threadId": "shared-thread", + "messages": [{"role": "user", "content": "Book another flight"}], + "forwardedProps": {"tenant": "tenant-b"}, + }, + ) + assert second_pause.status_code == 200 + + await store.save(scope="tenant-a", thread_id="shared-thread", snapshot=stale_snapshot) + stale_resume = client.post( + "/workflow", + json={ + "runId": "run-stale-resume", + "threadId": "shared-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + + stale_events = _decode_sse_events(stale_resume) + stale_errors = [event for event in stale_events if event.get("type") == "RUN_ERROR"] + assert len(stale_errors) == 1 + assert stale_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + + live_resume = client.post( + "/workflow", + json={ + "runId": "run-live-resume", + "threadId": "shared-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "United"}, + } + ], + "forwardedProps": {"tenant": "tenant-b"}, + }, + ) + live_events = _decode_sse_events(live_resume) + assert not [event for event in live_events if event.get("type") == "RUN_ERROR"] + text_deltas = [event["delta"] for event in live_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert "Booked United" in text_deltas + + async def test_endpoint_workflow_request_info_rejects_cancellation_from_different_thread(): """A different AG-UI thread cannot cancel another thread's workflow interrupt.""" app = _build_workflow_request_info_app() @@ -4062,8 +4166,8 @@ async def test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt( assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" -async def test_endpoint_workflow_checkpoint_resume_rejects_different_thread_after_restart(): - """Checkpoint restore preserves the AG-UI thread that owns a pending workflow interrupt.""" +async def test_endpoint_workflow_checkpoint_resume_rejects_threaded_resume_after_restart(): + """An explicitly threaded cold checkpoint resume fails closed when ownership is unavailable.""" storage = InMemoryCheckpointStorage() first_app = FastAPI() first_workflow = _build_flight_choice_workflow() @@ -4123,6 +4227,64 @@ async def test_endpoint_workflow_checkpoint_resume_rejects_different_thread_afte assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] +async def test_endpoint_workflow_checkpoint_resume_without_thread_remains_supported(): + """Legacy unthreaded checkpoint resumes remain compatible after wrapper restart.""" + storage = InMemoryCheckpointStorage() + first_app = FastAPI() + first_workflow = _build_flight_choice_workflow() + add_agent_framework_fastapi_endpoint( + first_app, + first_workflow, + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(first_app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + second_app = FastAPI() + add_agent_framework_fastapi_endpoint( + second_app, + _build_flight_choice_workflow(), + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(second_app) as client: + resume_response = client.post( + "/workflow", + json={ + "runId": "run-resume", + "messages": [], + "forwardedProps": {"checkpointId": checkpoint_id}, + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + resume_events = _decode_sse_events(resume_response) + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + text_deltas = [event["delta"] for event in resume_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert "Booked KLM" in text_deltas + + async def test_endpoint_workflow_request_info_cancelled_resume_completes_normally(): """Cancelled workflow resumes complete without output and do not wedge the next turn.""" app = _build_workflow_request_info_app() @@ -6236,10 +6398,25 @@ async def fake_run_workflow_stream(input_data: Any, workflow: Any, **kwargs: Any @executor(id="noop") async def noop(message: Any, ctx: WorkflowContext[Any, Any]) -> None: - del message, ctx + del message + await ctx.request_info({"agent": "flights"}, str, request_id="interrupt-1") + workflow = WorkflowBuilder(start_executor=noop).build() + _ = [ + event + async for event in workflow.run( + message=[Message(role="user", contents=[Content.from_text(text="First question")])], + stream=True, + ) + ] + pending_events = await workflow_module._pending_request_events(workflow) + setattr( + pending_events["interrupt-1"], + workflow_module._REQUEST_OWNER_ATTRIBUTE, + ("tenant-a", "workflow-thread"), + ) runner = AgentFrameworkWorkflow( - workflow=WorkflowBuilder(start_executor=noop).build(), + workflow=workflow, snapshot_store=store, ) From 3011eb78b5ea43f0be33d039b9d93ac78a2d0904 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Fri, 14 Aug 2026 15:41:07 +0900 Subject: [PATCH 3/3] Python: Persist AG-UI checkpoint ownership Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: aadad05d-2646-405f-8c62-a7a223abfc92 --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 94 ++++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 244 ++++++++++++++++++ 2 files changed, 327 insertions(+), 11 deletions(-) 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 b1893e8af4..a5ccd1a9ca 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -23,7 +23,7 @@ ToolCallResultEvent, ToolCallStartEvent, ) -from agent_framework import CheckpointStorage, Workflow +from agent_framework import CheckpointID, CheckpointStorage, Workflow, WorkflowCheckpoint from agent_framework._telemetry import mark_feature_used from ._feature_usage import FeatureIndex @@ -50,6 +50,7 @@ WorkflowRequestOwner = tuple[str | None, str | None] _REQUEST_OWNER_ATTRIBUTE = "_ag_ui_request_owner" +_CHECKPOINT_REQUEST_OWNER_KEY = "ag_ui_workflow_request_owner" def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: @@ -63,6 +64,58 @@ def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: return str(checkpoint_id) +def _checkpoint_request_owner(metadata: dict[str, Any]) -> WorkflowRequestOwner | None: + """Read AG-UI workflow request ownership from checkpoint metadata.""" + raw_owner = metadata.get(_CHECKPOINT_REQUEST_OWNER_KEY) + if not isinstance(raw_owner, dict): + return None + snapshot_scope = raw_owner.get("snapshot_scope") + thread_id = raw_owner.get("thread_id") + if snapshot_scope is not None and not isinstance(snapshot_scope, str): + return None + if thread_id is not None and not isinstance(thread_id, str): + return None + return snapshot_scope, thread_id + + +class _OwnedWorkflowCheckpointStorage: + """Attach one AG-UI request owner to checkpoints in their original save.""" + + def __init__(self, storage: CheckpointStorage, owner: WorkflowRequestOwner) -> None: + self._storage = storage + self._owner = owner + + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + """Save a checkpoint with ownership for any pending request occurrences.""" + if checkpoint.pending_request_info_events: + checkpoint.metadata = dict(checkpoint.metadata) + checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = { + "snapshot_scope": self._owner[0], + "thread_id": self._owner[1], + } + return await self._storage.save(checkpoint) + + async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: + """Load a checkpoint from the underlying store.""" + return await self._storage.load(checkpoint_id) + + async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: + """List checkpoints from the underlying store.""" + return await self._storage.list_checkpoints(workflow_name=workflow_name) + + async def delete(self, checkpoint_id: CheckpointID) -> bool: + """Delete a checkpoint from the underlying store.""" + return await self._storage.delete(checkpoint_id) + + async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: + """Get the latest checkpoint from the underlying store.""" + return await self._storage.get_latest(workflow_name=workflow_name) + + async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: + """List checkpoint IDs from the underlying store.""" + return await self._storage.list_checkpoint_ids(workflow_name=workflow_name) + + class _WorkflowSnapshotBuilder: """Capture replayable workflow protocol output without retaining raw events.""" @@ -323,21 +376,38 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: for interrupt in _normalize_resume_interrupts(resume_payload) if interrupt.get("id") is not None } + workflow = self._resolve_workflow(thread_id, snapshot_scope) 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 not None and checkpoint_storage is not None: - checkpoint = await checkpoint_storage.load(checkpoint_id) - checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} - for interrupt_id in checkpoint_pending_ids: - request_event = live_pending_events.get(interrupt_id) + if self.workflow is not None and checkpoint_id is None: + for request_event in live_pending_events.values(): owner = getattr(request_event, _REQUEST_OWNER_ATTRIBUTE, None) if owner != request_owner and (owner is not None or request_owner != (None, None)): yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield RunErrorEvent( - message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", + message="No pending interrupt found for this AG-UI thread.", code="WORKFLOW_RESUME_NOT_FOUND", ) return - if self.workflow is not None: + if checkpoint_id is not None and checkpoint_storage is not None: + try: + checkpoint = await checkpoint_storage.load(checkpoint_id) + except Exception as exc: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"Could not load workflow checkpoint '{checkpoint_id}': {exc}", + code="WORKFLOW_CHECKPOINT_LOAD_FAILED", + ) + return + checkpoint_pending_ids = {str(request_id) for request_id in checkpoint.pending_request_info_events} + checkpoint_owner = _checkpoint_request_owner(checkpoint.metadata) + if checkpoint_pending_ids and checkpoint_owner != request_owner: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent( + message=f"No pending interrupt found for checkpointId '{checkpoint_id}'.", + code="WORKFLOW_RESUME_NOT_FOUND", + ) + return + if self.workflow is not None and checkpoint_id is None: for interrupt_id in resume_interrupt_ids: request_event = live_pending_events.get(interrupt_id) owner = getattr(request_event, _REQUEST_OWNER_ATTRIBUTE, None) @@ -373,7 +443,9 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: if effective_state: input_data["state"] = effective_state - workflow = self._resolve_workflow(thread_id, snapshot_scope) + run_checkpoint_storage = checkpoint_storage + if checkpoint_storage is not None: + run_checkpoint_storage = _OwnedWorkflowCheckpointStorage(checkpoint_storage, request_owner) builder_seed_messages = raw_messages if resume_payload is not None or (checkpoint_id is not None and not raw_messages): # Resume requests carry only the synthesized interrupt response, and a @@ -389,7 +461,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: snapshot_builder.state = cast(dict[str, Any], state_snapshot) run_error_emitted = False async for event in run_workflow_stream( - input_data, workflow, checkpoint_storage=checkpoint_storage, checkpoint_id=checkpoint_id + input_data, workflow, checkpoint_storage=run_checkpoint_storage, checkpoint_id=checkpoint_id ): if snapshot_builder is not None: snapshot_builder.observe(event) @@ -400,7 +472,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: ): interrupt_id = str(event.tool_call_id) pending_event = (await _pending_request_events(workflow)).get(interrupt_id) - if pending_event is not None: + if pending_event is not None and getattr(pending_event, _REQUEST_OWNER_ATTRIBUTE, None) is None: setattr(pending_event, _REQUEST_OWNER_ATTRIBUTE, request_owner) if isinstance(event, RunErrorEvent): run_error_emitted = True diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 433fe1ef55..6fefe2a39b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -3889,6 +3889,37 @@ async def test_endpoint_workflow_request_info_rejects_resume_from_different_thre assert "Booked United" in text_deltas +async def test_endpoint_workflow_request_info_rejects_replay_from_different_thread(): + """A different AG-UI thread cannot observe another thread's pending interrupt.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + }, + ) + + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + assert not [event for event in attacker_events if event.get("type") == "TOOL_CALL_START"] + + async def test_endpoint_workflow_request_info_rejects_resume_from_different_scope(): """A workflow interrupt can only be resumed within the Snapshot Scope that created it.""" @@ -4227,6 +4258,191 @@ async def test_endpoint_workflow_checkpoint_resume_rejects_threaded_resume_after assert not [event for event in attacker_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] +async def test_endpoint_workflow_checkpoint_resume_same_owner_after_restart(): + """Checkpoint ownership permits the originating thread to resume after restart.""" + storage = InMemoryCheckpointStorage() + first_app = FastAPI() + first_workflow = _build_flight_choice_workflow() + add_agent_framework_fastapi_endpoint( + first_app, + first_workflow, + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(first_app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + second_app = FastAPI() + add_agent_framework_fastapi_endpoint( + second_app, + _build_flight_choice_workflow(), + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(second_app) as client: + resume_response = client.post( + "/workflow", + json={ + "runId": "run-resume", + "threadId": "victim-thread", + "messages": [], + "forwardedProps": {"checkpointId": checkpoint_id}, + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + resume_events = _decode_sse_events(resume_response) + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + text_deltas = [event["delta"] for event in resume_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert "Booked KLM" in text_deltas + + +async def test_endpoint_workflow_checkpoint_resume_uses_checkpoint_owner_not_live_reused_id(): + """A live reused interrupt ID cannot authorize a different checkpoint occurrence.""" + storage = InMemoryCheckpointStorage() + first_app = FastAPI() + first_workflow = _build_flight_choice_workflow() + add_agent_framework_fastapi_endpoint( + first_app, + first_workflow, + path="/workflow", + checkpoint_storage=storage, + ) + with TestClient(first_app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-tenant-a", + "threadId": "tenant-a-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + tenant_a_checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + second_app = FastAPI() + second_workflow = _build_flight_choice_workflow() + add_agent_framework_fastapi_endpoint( + second_app, + second_workflow, + path="/workflow", + checkpoint_storage=storage, + ) + with TestClient(second_app) as client: + tenant_b_pause = client.post( + "/workflow", + json={ + "runId": "run-tenant-b", + "threadId": "tenant-b-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert tenant_b_pause.status_code == 200 + + cross_occurrence_resume = client.post( + "/workflow", + json={ + "runId": "run-cross-occurrence", + "threadId": "tenant-b-thread", + "messages": [], + "forwardedProps": {"checkpointId": tenant_a_checkpoint_id}, + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + cross_events = _decode_sse_events(cross_occurrence_resume) + cross_errors = [event for event in cross_events if event.get("type") == "RUN_ERROR"] + assert len(cross_errors) == 1 + assert cross_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + + +async def test_endpoint_workflow_factory_checkpoint_resume_rejects_different_thread_after_restart(): + """Workflow-factory checkpoint restore validates the exact checkpoint owner.""" + storage = InMemoryCheckpointStorage() + created_workflows: list[Any] = [] + + def workflow_factory(_thread_id: str) -> Any: + workflow = _build_flight_choice_workflow() + created_workflows.append(workflow) + return workflow + + first_app = FastAPI() + first_runner = AgentFrameworkWorkflow(workflow_factory=workflow_factory, checkpoint_storage=storage) + add_agent_framework_fastapi_endpoint(first_app, first_runner, path="/workflow") + with TestClient(first_app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-victim", + "threadId": "victim-thread", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + checkpoints = await storage.list_checkpoints(workflow_name=created_workflows[0].name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + assert pending_checkpoints + checkpoint_id = max(pending_checkpoints, key=lambda checkpoint: checkpoint.timestamp).checkpoint_id + + second_app = FastAPI() + second_runner = AgentFrameworkWorkflow(workflow_factory=workflow_factory, checkpoint_storage=storage) + add_agent_framework_fastapi_endpoint(second_app, second_runner, path="/workflow") + with TestClient(second_app) as client: + attacker_response = client.post( + "/workflow", + json={ + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], + "forwardedProps": {"checkpointId": checkpoint_id}, + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + }, + ) + + attacker_events = _decode_sse_events(attacker_response) + attacker_errors = [event for event in attacker_events if event.get("type") == "RUN_ERROR"] + assert len(attacker_errors) == 1 + assert attacker_errors[0]["code"] == "WORKFLOW_RESUME_NOT_FOUND" + + async def test_endpoint_workflow_checkpoint_resume_without_thread_remains_supported(): """Legacy unthreaded checkpoint resumes remain compatible after wrapper restart.""" storage = InMemoryCheckpointStorage() @@ -4285,6 +4501,34 @@ async def test_endpoint_workflow_checkpoint_resume_without_thread_remains_suppor assert "Booked KLM" in text_deltas +async def test_endpoint_workflow_checkpoint_load_failure_emits_protocol_error(): + """Checkpoint load failures emit RUN_STARTED before a useful RUN_ERROR.""" + storage = InMemoryCheckpointStorage() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + _build_flight_choice_workflow(), + path="/workflow", + checkpoint_storage=storage, + ) + + with TestClient(app) as client: + response = client.post( + "/workflow", + json={ + "runId": "run-missing-checkpoint", + "threadId": "thread-1", + "messages": [], + "forwardedProps": {"checkpointId": "missing-checkpoint"}, + }, + ) + + events = _decode_sse_events(response) + assert [event["type"] for event in events] == ["RUN_STARTED", "RUN_ERROR"] + assert events[-1]["code"] == "WORKFLOW_CHECKPOINT_LOAD_FAILED" + assert "missing-checkpoint" in events[-1]["message"] + + async def test_endpoint_workflow_request_info_cancelled_resume_completes_normally(): """Cancelled workflow resumes complete without output and do not wedge the next turn.""" app = _build_workflow_request_info_app()