diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 4027d9a8b1..8f90781a67 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 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` | @@ -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..a5ccd1a9ca 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, @@ -22,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 @@ -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 @@ -40,11 +42,15 @@ 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] + +_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: @@ -58,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.""" @@ -301,6 +359,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 +369,56 @@ 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 + } + 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 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="No pending interrupt found for this AG-UI thread.", + code="WORKFLOW_RESUME_NOT_FOUND", + ) + return + 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) + 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 +428,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, @@ -335,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 @@ -351,10 +461,19 @@ 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) + if ( + self.workflow is not None + and isinstance(event, ToolCallStartEvent) + and event.tool_call_name == "request_info" + ): + interrupt_id = str(event.tool_call_id) + pending_event = (await _pending_request_events(workflow)).get(interrupt_id) + 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 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..6fefe2a39b 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 @@ -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, @@ -58,6 +60,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() -> ASGIMessage: + 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: 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 = { + "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 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]]: """Return canonical interrupts from an SSE RUN_FINISHED event.""" assert "interrupt" not in event @@ -88,6 +134,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 +225,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 +3740,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 +3759,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,8 +3828,8 @@ 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_cancelled_resume_completes_normally(): - """Cancelled workflow resumes complete without output and do not wedge the next turn.""" +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: @@ -3639,45 +3837,60 @@ async def test_endpoint_workflow_request_info_cancelled_resume_completes_normall "/workflow", json={ "runId": "run-pause", - "threadId": "thread-flights", + "threadId": "victim-thread", "messages": [{"role": "user", "content": "Book me a flight"}], }, ) assert pause_response.status_code == 200 - resume_response = client.post( + attacker_response = client.post( "/workflow", json={ - "runId": "run-cancel", - "threadId": "thread-flights", + "runId": "run-attacker", + "threadId": "attacker-thread", "messages": [], - "resume": [{"interruptId": "flight-choice", "status": "cancelled"}], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], }, ) - assert resume_response.status_code == 200 - events = _decode_sse_events(resume_response) - assert [event.get("type") for event in events] == ["RUN_STARTED", "RUN_FINISHED"] - assert not [event for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + 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"] - next_response = client.post( + victim_response = client.post( "/workflow", json={ - "runId": "run-after-cancel", - "threadId": "thread-flights", - "messages": [{"role": "user", "content": "Book a different flight"}], + "runId": "run-victim-resume", + "threadId": "victim-thread", + "messages": [], + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "United"}, + } + ], }, ) - 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"] - next_finished = [event for event in next_events if event.get("type") == "RUN_FINISHED"] - assert _run_finished_interrupts(next_finished[-1])[0]["id"] == "flight-choice" + 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_new_input_with_pending_interrupt_emits_run_error(): - """New non-resume input on a workflow-interrupted thread must fail with RUN_ERROR.""" +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: @@ -3685,116 +3898,802 @@ async def test_endpoint_workflow_request_info_new_input_with_pending_interrupt_e "/workflow", json={ "runId": "run-pause", - "threadId": "thread-flights", + "threadId": "victim-thread", "messages": [{"role": "user", "content": "Book me a flight"}], }, ) assert pause_response.status_code == 200 - response = client.post( + attacker_response = client.post( "/workflow", json={ - "runId": "run-new-input", - "threadId": "thread-flights", - "messages": [{"role": "user", "content": "I prefer KLM"}], + "runId": "run-attacker", + "threadId": "attacker-thread", + "messages": [], }, ) - assert response.status_code == 200 - events = _decode_sse_events(response) - run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "WORKFLOW_RESUME_REQUIRED" - assert not [event for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + 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_malformed_resume_entry_emits_run_error(): - """Malformed workflow resume entries must fail as observable stream RUN_ERROR events.""" - app = _build_workflow_request_info_app() +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": "thread-flights", + "threadId": "shared-thread", "messages": [{"role": "user", "content": "Book me a flight"}], + "forwardedProps": {"tenant": "tenant-a"}, }, ) assert pause_response.status_code == 200 - response = client.post( + attacker_response = client.post( "/workflow", json={ - "runId": "run-malformed", - "threadId": "thread-flights", + "runId": "run-attacker", + "threadId": "shared-thread", "messages": [], - "forwardedProps": {"command": {"resume": [{"status": "resolved", "payload": {"airline": "KLM"}}]}}, + "resume": [ + { + "interruptId": "flight-choice", + "status": "resolved", + "payload": {"airline": "KLM"}, + } + ], + "forwardedProps": {"tenant": "tenant-b"}, }, ) - assert response.status_code == 200 - events = _decode_sse_events(response) - run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "WORKFLOW_RESUME_INVALID" + 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_invalid_response_payload_emits_run_error(): - """Workflow resume payloads that fail declared response-schema coercion must RUN_ERROR.""" - app = _build_workflow_request_info_app() +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: - pause_response = client.post( + first_pause = client.post( "/workflow", json={ - "runId": "run-pause", - "threadId": "thread-flights", + "runId": "run-first-pause", + "threadId": "shared-thread", "messages": [{"role": "user", "content": "Book me a flight"}], + "forwardedProps": {"tenant": "tenant-a"}, }, ) - assert pause_response.status_code == 200 + assert first_pause.status_code == 200 + stale_snapshot = await store.get(scope="tenant-a", thread_id="shared-thread") + assert stale_snapshot is not None - response = client.post( + first_cancel = client.post( "/workflow", json={ - "runId": "run-invalid-payload", - "threadId": "thread-flights", + "runId": "run-first-cancel", + "threadId": "shared-thread", "messages": [], - "resume": [{"interruptId": "flight-choice", "status": "resolved", "payload": "KLM"}], + "resume": [{"interruptId": "flight-choice", "status": "cancelled"}], + "forwardedProps": {"tenant": "tenant-a"}, }, ) + assert first_cancel.status_code == 200 - assert response.status_code == 200 - events = _decode_sse_events(response) - run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "WORKFLOW_RESUME_INVALID_RESPONSE" - + 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 -async def test_endpoint_with_workflow_as_agent_stream_output(build_chat_client): - """Test endpoint handles workflow-as-agent stream outputs.""" - app = FastAPI() - brainstorm_agent = Agent(name="brainstorm", instructions="Brainstorm ideas", client=build_chat_client("Idea")) - reviewer_agent = Agent(name="reviewer", instructions="Review ideas", client=build_chat_client("Review")) - agent = SequentialBuilder(participants=[brainstorm_agent, reviewer_agent]).build().as_agent() + 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"}, + }, + ) - add_agent_framework_fastapi_endpoint(app, agent, path="/workflow-like") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + 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" - client = TestClient(app) - response = client.post("/workflow-like", json={"messages": [{"role": "user", "content": "Hello"}]}) + 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 - assert response.status_code == 200 - content = response.content.decode("utf-8") - lines = [line for line in content.split("\n") if line.startswith("data: ")] - event_types = [json.loads(line[6:]).get("type") for line in lines] - assert "RUN_STARTED" in event_types - assert "TEXT_MESSAGE_CONTENT" in event_types - assert "RUN_FINISHED" in event_types +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 -async def test_endpoint_error_handling(build_chat_client): + 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_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() + 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_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() + 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_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() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "thread-flights", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + resume_response = client.post( + "/workflow", + json={ + "runId": "run-cancel", + "threadId": "thread-flights", + "messages": [], + "resume": [{"interruptId": "flight-choice", "status": "cancelled"}], + }, + ) + + assert resume_response.status_code == 200 + events = _decode_sse_events(resume_response) + assert [event.get("type") for event in events] == ["RUN_STARTED", "RUN_FINISHED"] + assert not [event for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + next_response = client.post( + "/workflow", + json={ + "runId": "run-after-cancel", + "threadId": "thread-flights", + "messages": [{"role": "user", "content": "Book a different flight"}], + }, + ) + + 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"] + next_finished = [event for event in next_events if event.get("type") == "RUN_FINISHED"] + assert _run_finished_interrupts(next_finished[-1])[0]["id"] == "flight-choice" + + +async def test_endpoint_workflow_request_info_new_input_with_pending_interrupt_emits_run_error(): + """New non-resume input on a workflow-interrupted thread must fail with RUN_ERROR.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "thread-flights", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + response = client.post( + "/workflow", + json={ + "runId": "run-new-input", + "threadId": "thread-flights", + "messages": [{"role": "user", "content": "I prefer KLM"}], + }, + ) + + assert response.status_code == 200 + events = _decode_sse_events(response) + run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] + assert len(run_errors) == 1 + assert run_errors[0]["code"] == "WORKFLOW_RESUME_REQUIRED" + assert not [event for event in events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + + +async def test_endpoint_workflow_request_info_malformed_resume_entry_emits_run_error(): + """Malformed workflow resume entries must fail as observable stream RUN_ERROR events.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "thread-flights", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + response = client.post( + "/workflow", + json={ + "runId": "run-malformed", + "threadId": "thread-flights", + "messages": [], + "forwardedProps": {"command": {"resume": [{"status": "resolved", "payload": {"airline": "KLM"}}]}}, + }, + ) + + assert response.status_code == 200 + events = _decode_sse_events(response) + run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] + assert len(run_errors) == 1 + assert run_errors[0]["code"] == "WORKFLOW_RESUME_INVALID" + + +async def test_endpoint_workflow_request_info_invalid_response_payload_emits_run_error(): + """Workflow resume payloads that fail declared response-schema coercion must RUN_ERROR.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_response = client.post( + "/workflow", + json={ + "runId": "run-pause", + "threadId": "thread-flights", + "messages": [{"role": "user", "content": "Book me a flight"}], + }, + ) + assert pause_response.status_code == 200 + + response = client.post( + "/workflow", + json={ + "runId": "run-invalid-payload", + "threadId": "thread-flights", + "messages": [], + "resume": [{"interruptId": "flight-choice", "status": "resolved", "payload": "KLM"}], + }, + ) + + assert response.status_code == 200 + events = _decode_sse_events(response) + run_errors = [event for event in events if event.get("type") == "RUN_ERROR"] + assert len(run_errors) == 1 + assert run_errors[0]["code"] == "WORKFLOW_RESUME_INVALID_RESPONSE" + + +async def test_endpoint_with_workflow_as_agent_stream_output(build_chat_client): + """Test endpoint handles workflow-as-agent stream outputs.""" + app = FastAPI() + brainstorm_agent = Agent(name="brainstorm", instructions="Brainstorm ideas", client=build_chat_client("Idea")) + reviewer_agent = Agent(name="reviewer", instructions="Review ideas", client=build_chat_client("Review")) + agent = SequentialBuilder(participants=[brainstorm_agent, reviewer_agent]).build().as_agent() + + add_agent_framework_fastapi_endpoint(app, agent, path="/workflow-like") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] + + client = TestClient(app) + response = client.post("/workflow-like", json={"messages": [{"role": "user", "content": "Hello"}]}) + + assert response.status_code == 200 + content = response.content.decode("utf-8") + lines = [line for line in content.split("\n") if line.startswith("data: ")] + event_types = [json.loads(line[6:]).get("type") for line in lines] + + assert "RUN_STARTED" in event_types + assert "TEXT_MESSAGE_CONTENT" in event_types + assert "RUN_FINISHED" in event_types + + +async def test_endpoint_error_handling(build_chat_client): """Test endpoint error handling during request parsing.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) @@ -5743,10 +6642,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, ) 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")]