diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 4470067c1ba..afdee048806 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -108,6 +108,13 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - AG-UI Thread and Run ids are client-owned protocol correlation ids. Service-session mode stores provider conversation or response ids privately in the thread snapshot. Set `service_session_id_from_thread_id=True` only for compatibility when the application intentionally uses a provider continuation id as its AG-UI Thread id. +- A workflow run resolves its Thread id once. A generated id for an omitted `threadId` is the id reported in + `RUN_STARTED`, keys the `workflow_factory` cache and Thread Snapshot, and owns the run's live `request_info` + interrupts, so a continuation must send it back. A `workflow_factory` request without `threadId` that carries a + `resume` entry or a tool or approval-response message (as `_extract_responses_from_messages` defines them) is + rejected with `WORKFLOW_RESUME_NOT_FOUND` before a workflow instance is built. Checkpoint ownership binds only + client-supplied Thread ids; a checkpoint of a run without one also records the generated id as an alias, so both + the reported id and legacy id-less restores resume it. - A configured Snapshot Scope and the client-owned Thread id are combined into a deterministic internal `AgentSession.session_id`. Keep the raw Thread id for protocol events and snapshot addressing so equal Thread ids in different trusted scopes cannot collide in context-provider state. The deprecated diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index ea677bd2220..07d95cd2f45 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -495,6 +495,13 @@ AG-UI events and snapshot operations. When using `AgentFrameworkWorkflow(workflo also scopes the in-memory workflow cache even without a snapshot store; provide it in multi-user deployments so two users who submit the same `threadId` do not share a live `Workflow` instance. +A workflow request that omits `threadId` is assigned a generated Thread id. It is reported in `RUN_STARTED`, keys the +`workflow_factory` instance and Thread Snapshot, and owns that run's pending `request_info` interrupts. Send it back as +`threadId` to continue the run; a continuation that omits it receives `WORKFLOW_RESUME_NOT_FOUND`. With +`workflow_factory`, a request without `threadId` that carries a `resume` entry, a tool result, or an approval response +is rejected the same way instead of starting a new run, so send `threadId` with every such request. A checkpoint +created by such a run resumes with `checkpointId` and either the reported `threadId` or, as before, no `threadId`. + Authenticate and authorize the current scope on every request, including hydration and approval resume, rather than trusting a scope supplied through Shared State or forwarded properties. A configured resolver must return a non-empty string. Returning `None`, an empty string, or another type fails the request with a generic HTTP 500 configuration 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 962ccb384ab..4b56f672098 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -35,7 +35,11 @@ from agent_framework._telemetry import mark_feature_used from ._feature_usage import FeatureIndex -from ._message_adapters import agent_framework_messages_to_agui, agui_messages_to_snapshot_format +from ._message_adapters import ( + agent_framework_messages_to_agui, + agui_messages_to_snapshot_format, + normalize_agui_input_messages, +) from ._run_common import ( _cancelled_resume_interrupt_ids, _extract_resume_payload, @@ -61,7 +65,11 @@ generate_event_id, make_json_safe, ) -from ._workflow_run import _pending_request_events, run_workflow_stream # pyright: ignore[reportPrivateUsage] +from ._workflow_run import ( # pyright: ignore[reportPrivateUsage] + _extract_responses_from_messages, + _pending_request_events, + run_workflow_stream, +) logger = logging.getLogger(__name__) @@ -232,6 +240,12 @@ def _checkpoint_id_from_input(input_data: dict[str, Any]) -> str | None: return str(checkpoint_id) +def _message_response_ids(raw_messages: list[dict[str, Any]]) -> set[str]: + """Return the request ids that tool or approval-response messages in a request answer.""" + messages, _ = normalize_agui_input_messages(raw_messages, sanitize_tool_history=False) + return set(_extract_responses_from_messages(messages)) + + 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) @@ -246,6 +260,15 @@ def _checkpoint_request_owner(metadata: dict[str, Any]) -> WorkflowRequestOwner return snapshot_scope, thread_id +def _checkpoint_generated_thread_id(metadata: dict[str, Any]) -> str | None: + """Read the generated thread id recorded for a checkpoint whose run omitted ``threadId``.""" + raw_owner = metadata.get(_CHECKPOINT_REQUEST_OWNER_KEY) + if not isinstance(raw_owner, dict): + return None + generated_thread_id = raw_owner.get("generated_thread_id") + return generated_thread_id if isinstance(generated_thread_id, str) else None + + def _workflow_checkpoint_storage(workflow: Workflow) -> CheckpointStorage | None: """Return the workflow's build-time checkpoint storage.""" runner = getattr(workflow, "_runner", None) @@ -258,21 +281,33 @@ def _workflow_checkpoint_storage(workflow: Workflow) -> CheckpointStorage | None class _OwnedWorkflowCheckpointStorage: """Attach one AG-UI request owner to checkpoints in their original save.""" - def __init__(self, storage: CheckpointStorage, owner: WorkflowRequestOwner) -> None: + def __init__( + self, storage: CheckpointStorage, owner: WorkflowRequestOwner, generated_thread_id: str | None = None + ) -> None: self._storage = storage self._owner = owner + self._generated_thread_id = generated_thread_id def _owns(self, checkpoint: WorkflowCheckpoint) -> bool: checkpoint_owner = _checkpoint_request_owner(checkpoint.metadata) - return checkpoint_owner == self._owner or (checkpoint_owner is None and self._owner == (None, None)) + if checkpoint_owner == self._owner or (checkpoint_owner is None and self._owner == (None, None)): + return True + # A run that omitted threadId also records the id it reported, so a continuation can echo that id. + generated_thread_id = _checkpoint_generated_thread_id(checkpoint.metadata) + return ( + checkpoint_owner is not None + and checkpoint_owner[1] is None + and generated_thread_id is not None + and self._owner == (checkpoint_owner[0], generated_thread_id) + ) async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: """Save a checkpoint with ownership for any pending request occurrences.""" checkpoint.metadata = dict(checkpoint.metadata) - checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = { - "snapshot_scope": self._owner[0], - "thread_id": self._owner[1], - } + owner_metadata: dict[str, str | None] = {"snapshot_scope": self._owner[0], "thread_id": self._owner[1]} + if self._generated_thread_id is not None: + owner_metadata["generated_thread_id"] = self._generated_thread_id + checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] = owner_metadata return await self._storage.save(checkpoint) async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: @@ -663,13 +698,16 @@ def __init__( self.snapshot_store = snapshot_store self.checkpoint_storage = checkpoint_storage + @staticmethod + def _supplied_thread_id_from_input(input_data: dict[str, Any]) -> str | None: + """Return the client-supplied thread id, treating an empty value as omitted.""" + thread_id = input_data.get("thread_id") or input_data.get("threadId") + return str(thread_id) if thread_id else None + @staticmethod def _thread_id_from_input(input_data: dict[str, Any]) -> str: """Resolve a stable thread id from AG-UI input payload.""" - thread_id = input_data.get("thread_id") or input_data.get("threadId") - if thread_id is not None: - return str(thread_id) - return str(uuid.uuid4()) + return AgentFrameworkWorkflow._supplied_thread_id_from_input(input_data) or str(uuid.uuid4()) def _resolve_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> Workflow: """Get the workflow instance for the current run.""" @@ -742,17 +780,39 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: checkpoint_id = _checkpoint_id_from_input(input_data) - 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) + supplied_thread_id = self._supplied_thread_id_from_input(input_data) + request_owner = (snapshot_scope, supplied_thread_id) + # Live interrupts bind to the resolved thread id, including a generated one that RUN_STARTED reports, so a + # continuation must echo it. Checkpoint ownership keeps binding only client-supplied ids and records a + # generated id as an alias, which keeps legacy id-less checkpoint restores working. + live_owner = (snapshot_scope, thread_id) 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 None + and checkpoint_id is None + and supplied_thread_id is None + and (resume_interrupt_ids or _message_response_ids(raw_messages)) + ): + # A factory run that omits threadId reports a new generated id, so nothing can be pending on it and a + # resume entry, tool result, or approval response cannot continue anything. + 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 workflow = self._resolve_workflow(thread_id, snapshot_scope) checkpoint_storage = self.checkpoint_storage or _workflow_checkpoint_storage(workflow) owned_checkpoint_storage = ( - _OwnedWorkflowCheckpointStorage(checkpoint_storage, request_owner) + _OwnedWorkflowCheckpointStorage( + checkpoint_storage, + request_owner, + generated_thread_id=thread_id if supplied_thread_id is None else None, + ) if checkpoint_storage is not None else None ) @@ -766,7 +826,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: 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)): + if owner != live_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.", @@ -801,7 +861,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: 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)): + if owner != live_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}'.", @@ -880,6 +940,7 @@ def append_conversational_responses(values: list[Any]) -> None: checkpoint_id=checkpoint_id, pending_request_events=effective_pending_events, on_conversational_responses=append_conversational_responses, + resolved_thread_id=thread_id, ): if snapshot_builder is not None: snapshot_builder.observe(event) @@ -891,7 +952,7 @@ def append_conversational_responses(values: list[Any]) -> None: 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) + setattr(pending_event, _REQUEST_OWNER_ATTRIBUTE, live_owner) if isinstance(event, RunErrorEvent): run_error_emitted = True if getattr(event, "code", None) == "WORKFLOW_RESUME_CANCELLED": diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index fa92cd400ce..fcb80bb49e8 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -1149,6 +1149,7 @@ async def run_workflow_stream( checkpoint_id: str | None = None, pending_request_events: dict[str, Any] | None = None, on_conversational_responses: Callable[[list[Any]], None] | None = None, + resolved_thread_id: str | None = None, ) -> AsyncGenerator[BaseEvent]: """Run a Workflow and emit AG-UI protocol events. @@ -1168,9 +1169,11 @@ async def run_workflow_stream( start-executor message, so resume stays consistent with the core API. pending_request_events: Optional effective pending requests already resolved by the wrapper. on_conversational_responses: Optional callback for validated non-control responses. + resolved_thread_id: Optional thread id already resolved by the caller, reported in protocol events instead of + generating a new one when the request omits it. """ supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") - thread_id = supplied_thread_id or str(uuid.uuid4()) + thread_id = resolved_thread_id or supplied_thread_id or str(uuid.uuid4()) run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4()) available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts") if available_interrupts: 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 2a2fc812832..0d7d9d61c83 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -516,13 +516,17 @@ async def approve( pause_response = client.post( "/workflow-approval", - json={"messages": [{"role": "user", "content": "Refund the order"}]}, + json={ + "threadId": "thread-approval", + "messages": [{"role": "user", "content": "Refund the order"}], + }, ) assert pause_response.status_code == 200 resume_response = client.post( "/workflow-approval", json={ + "threadId": "thread-approval", "messages": [], "resume": [{"interruptId": "approval-1", "status": "resolved", "payload": {"approved": True}}], }, @@ -1048,13 +1052,17 @@ async def approve( client = TestClient(app) pause_response = client.post( "/workflow-approval", - json={"messages": [{"role": "user", "content": "Refund the order"}]}, + json={ + "threadId": "thread-approval", + "messages": [{"role": "user", "content": "Refund the order"}], + }, ) assert pause_response.status_code == 200 resume_response = client.post( "/workflow-approval", json={ + "threadId": "thread-approval", "messages": [], "resume": [ { @@ -1115,13 +1123,17 @@ async def approve( client = TestClient(app) pause_response = client.post( "/workflow-approval", - json={"messages": [{"role": "user", "content": "Refund the order"}]}, + json={ + "threadId": "thread-approval", + "messages": [{"role": "user", "content": "Refund the order"}], + }, ) assert pause_response.status_code == 200 resume_response = client.post( "/workflow-approval", json={ + "threadId": "thread-approval", "messages": [], "resume": [ { @@ -1179,7 +1191,10 @@ async def approve( client = TestClient(app) pause_response = client.post( "/workflow-approval", - json={"messages": [{"role": "user", "content": "Refund the order"}]}, + json={ + "threadId": "thread-approval", + "messages": [{"role": "user", "content": "Refund the order"}], + }, ) assert pause_response.status_code == 200 @@ -1190,6 +1205,7 @@ async def approve( resume_response = client.post( "/workflow-approval", json={ + "threadId": "thread-approval", "messages": [], "resume": [ { @@ -8528,6 +8544,164 @@ async def test_endpoint_workflow_request_info_rejects_cancellation_from_differen assert "Booked United" in text_deltas +def _flight_choice_resume(airline: str, *, thread_id: str | None = None, **extra: Any) -> dict[str, Any]: + payload: dict[str, Any] = { + "runId": f"run-resume-{airline}", + "messages": [], + "resume": [{"interruptId": "flight-choice", "status": "resolved", "payload": {"airline": airline}}], + **extra, + } + if thread_id is not None: + payload["threadId"] = thread_id + return payload + + +def _reported_thread_ids(events: list[dict[str, Any]]) -> set[str]: + return {event["threadId"] for event in events if event.get("type") in {"RUN_STARTED", "RUN_FINISHED"}} + + +def _assert_resume_not_found(events: list[dict[str, Any]]) -> None: + assert [event["code"] for event in events if event.get("type") == "RUN_ERROR"] == ["WORKFLOW_RESUME_NOT_FOUND"] + assert not [event for event in events if event.get("type") in {"TEXT_MESSAGE_CONTENT", "TOOL_CALL_START"}] + + +async def test_endpoint_workflow_request_info_omitted_thread_binds_interrupt_to_reported_thread(): + """A run without threadId reports one generated id, and only that id continues its interrupt.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_events = _decode_sse_events( + client.post( + "/workflow", + json={"runId": "run-pause", "messages": [{"role": "user", "content": "Book me a flight"}]}, + ) + ) + reported = _reported_thread_ids(pause_events) + assert len(reported) == 1 + thread_id = reported.pop() + + for rejected in ( + _flight_choice_resume("KLM"), + _flight_choice_resume("KLM", thread_id="other-thread"), + {"runId": "run-observe", "messages": []}, + ): + _assert_resume_not_found(_decode_sse_events(client.post("/workflow", json=rejected))) + + resume_events = _decode_sse_events( + client.post("/workflow", json=_flight_choice_resume("United", thread_id=thread_id)) + ) + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + assert _reported_thread_ids(resume_events) == {thread_id} + text_deltas = [event["delta"] for event in resume_events if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert text_deltas.count("Booked United") == 1 + assert "Booked KLM" not in text_deltas + + +async def test_endpoint_workflow_request_info_omitted_thread_rejects_cancellation_from_other_omitted_thread(): + """A caller that also omits threadId cannot cancel another run's pending interrupt.""" + app = _build_workflow_request_info_app() + + with TestClient(app) as client: + pause_events = _decode_sse_events( + client.post( + "/workflow", + json={"runId": "run-pause", "messages": [{"role": "user", "content": "Book me a flight"}]}, + ) + ) + thread_id = next(iter(_reported_thread_ids(pause_events))) + + cancel_events = _decode_sse_events( + client.post( + "/workflow", + json={ + "runId": "run-cancel", + "messages": [], + "resume": [{"interruptId": "flight-choice", "status": "cancelled"}], + }, + ) + ) + _assert_resume_not_found(cancel_events) + + resume_events = _decode_sse_events( + client.post("/workflow", json=_flight_choice_resume("United", thread_id=thread_id)) + ) + 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 United" in text_deltas + + +async def test_endpoint_workflow_request_info_omitted_thread_rejects_resume_from_different_scope(): + """With a Snapshot Scope, neither an omitted nor the reported thread id lets another scope resume.""" + + 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_events = _decode_sse_events( + client.post( + "/workflow", + json={ + "runId": "run-pause", + "messages": [{"role": "user", "content": "Book me a flight"}], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + ) + thread_id = next(iter(_reported_thread_ids(pause_events))) + + for rejected in ( + _flight_choice_resume("KLM", forwardedProps={"tenant": "tenant-b"}), + _flight_choice_resume("KLM", thread_id=thread_id, forwardedProps={"tenant": "tenant-b"}), + ): + _assert_resume_not_found(_decode_sse_events(client.post("/workflow", json=rejected))) + + resume_events = _decode_sse_events( + client.post( + "/workflow", + json=_flight_choice_resume("United", thread_id=thread_id, forwardedProps={"tenant": "tenant-a"}), + ) + ) + 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 United" in text_deltas + + +async def test_endpoint_workflow_omitted_thread_snapshot_hydrates_by_reported_thread_id(): + """The Thread Snapshot of a run without threadId is stored under the id its events report.""" + store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + _build_flight_choice_workflow(), + path="/workflow", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + ) + + with TestClient(app) as client: + pause_events = _decode_sse_events( + client.post( + "/workflow", + json={"runId": "run-pause", "messages": [{"role": "user", "content": "Book me a flight"}]}, + ) + ) + thread_id = next(iter(_reported_thread_ids(pause_events))) + assert await store.get(scope="tenant-a", thread_id=thread_id) is not None + + hydrate_events = _decode_sse_events( + client.post("/workflow", json={"runId": "run-hydrate", "threadId": thread_id, "messages": []}) + ) + snapshots = [event for event in hydrate_events if event.get("type") == "MESSAGES_SNAPSHOT"] + assert snapshots + assert "Book me a flight" in json.dumps(snapshots[-1]["messages"]) + + 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() @@ -8617,6 +8791,25 @@ async def test_owned_checkpoint_storage_stamps_owner_without_pending_events() -> assert stored.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] == expected_owner +async def test_owned_checkpoint_storage_accepts_generated_thread_alias_only_within_scope() -> None: + """A checkpoint of a run without threadId is also owned by its generated id, in the same scope only.""" + storage = InMemoryCheckpointStorage() + writer = _OwnedWorkflowCheckpointStorage(storage, ("scope-1", None), generated_thread_id="generated-1") + checkpoint_id = await writer.save(WorkflowCheckpoint(workflow_name="owned-workflow", graph_signature_hash="sig")) + + for accepted_owner in (("scope-1", None), ("scope-1", "generated-1")): + loaded = await _OwnedWorkflowCheckpointStorage(storage, accepted_owner).load(checkpoint_id) + assert loaded.checkpoint_id == checkpoint_id + for rejected_owner in (("scope-1", "other-thread"), ("scope-2", "generated-1"), (None, "generated-1")): + with pytest.raises(PermissionError): + await _OwnedWorkflowCheckpointStorage(storage, rejected_owner).load(checkpoint_id) + + supplied = _OwnedWorkflowCheckpointStorage(storage, ("scope-1", "supplied-1"), generated_thread_id=None) + supplied_id = await supplied.save(WorkflowCheckpoint(workflow_name="owned-workflow", graph_signature_hash="sig")) + with pytest.raises(PermissionError): + await _OwnedWorkflowCheckpointStorage(storage, ("scope-1", None)).load(supplied_id) + + async def test_owned_checkpoint_storage_filters_and_rejects_foreign_checkpoints() -> None: """Owned storage never exposes or deletes another request owner's checkpoints.""" storage = InMemoryCheckpointStorage() @@ -9115,6 +9308,54 @@ async def test_endpoint_builder_checkpoint_storage_ignores_stale_runtime_owner() } +async def test_endpoint_workflow_checkpoint_resume_with_reported_thread_id(): + """After a restart, a run without threadId resumes from its checkpoint using the id its events reported.""" + 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_events = _decode_sse_events( + client.post( + "/workflow", + json={"runId": "run-pause", "messages": [{"role": "user", "content": "Book me a flight"}]}, + ) + ) + thread_id = next(event["threadId"] for event in pause_events if event.get("type") == "RUN_STARTED") + + checkpoints = await storage.list_checkpoints(workflow_name=first_workflow.name) + pending_checkpoints = [checkpoint for checkpoint in checkpoints if checkpoint.pending_request_info_events] + checkpoint = max(pending_checkpoints, key=lambda item: item.timestamp) + assert checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] == { + "snapshot_scope": None, + "thread_id": None, + "generated_thread_id": thread_id, + } + + def resume_payload(thread_id_value: str) -> dict[str, Any]: + return { + "runId": "run-resume", + "threadId": thread_id_value, + "messages": [], + "forwardedProps": {"checkpointId": checkpoint.checkpoint_id}, + "resume": [{"interruptId": "flight-choice", "status": "resolved", "payload": {"airline": "KLM"}}], + } + + 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: + rejected = _decode_sse_events(client.post("/workflow", json=resume_payload("other-thread"))) + _assert_resume_not_found(rejected) + + resumed = _decode_sse_events(client.post("/workflow", json=resume_payload(thread_id))) + assert not [event for event in resumed if event.get("type") == "RUN_ERROR"] + text_deltas = [event["delta"] for event in resumed if event.get("type") == "TEXT_MESSAGE_CONTENT"] + assert "Booked KLM" in text_deltas + + async def test_endpoint_workflow_checkpoint_resume_without_thread_remains_supported(): """Legacy unthreaded checkpoint resumes remain compatible after wrapper restart.""" storage = InMemoryCheckpointStorage() @@ -11806,7 +12047,10 @@ def approval_update(approval_id: str, call_id: str) -> AgentResponseUpdate: first_pause = client.post( "/workflow-agent-approval", - json={"messages": [{"role": "user", "content": "First refund"}]}, + json={ + "threadId": "thread-agent-approval", + "messages": [{"role": "user", "content": "First refund"}], + }, ) assert first_pause.status_code == 200 first_finished = [event for event in _decode_sse_events(first_pause) if event.get("type") == "RUN_FINISHED"] @@ -11815,6 +12059,7 @@ def approval_update(approval_id: str, call_id: str) -> AgentResponseUpdate: cancelled = client.post( "/workflow-agent-approval", json={ + "threadId": "thread-agent-approval", "messages": [], "resume": [{"interruptId": "approval-1", "status": "cancelled"}], }, @@ -11824,7 +12069,10 @@ def approval_update(approval_id: str, call_id: str) -> AgentResponseUpdate: agent.updates = [approval_update("approval-2", "refund-call-2")] second_pause = client.post( "/workflow-agent-approval", - json={"messages": [{"role": "user", "content": "Second refund"}]}, + json={ + "threadId": "thread-agent-approval", + "messages": [{"role": "user", "content": "Second refund"}], + }, ) assert second_pause.status_code == 200 second_finished = [event for event in _decode_sse_events(second_pause) if event.get("type") == "RUN_FINISHED"] @@ -11836,6 +12084,7 @@ def approval_update(approval_id: str, call_id: str) -> AgentResponseUpdate: resumed = client.post( "/workflow-agent-approval", json={ + "threadId": "thread-agent-approval", "messages": [], "resume": [{"interruptId": "approval-2", "status": "resolved", "payload": {"approved": True}}], }, diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 6f5df70d923..1c46e177565 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -24,6 +24,8 @@ ) from agent_framework_ag_ui import AgentFrameworkWorkflow +from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY +from agent_framework_ag_ui._workflow import _CHECKPOINT_REQUEST_OWNER_KEY class _HitlMessageRequestExecutor(Executor): @@ -61,6 +63,41 @@ def _interrupts_from_finished(event: Any) -> list[dict[str, Any]]: return cast(list[dict[str, Any]], interrupts) +class _ChoiceRecorder(Executor): + """Pauses on a stable request id and records every response it dispatches.""" + + def __init__(self, handled: list[str]) -> None: + super().__init__(id="choice_recorder") + self._handled = handled + + @handler + async def start(self, message: Any, ctx: WorkflowContext[Any, str]) -> None: + del message + await ctx.request_info({"message": "Choose an option"}, dict, request_id="choice") + + @response_handler + async def handle_choice( + self, original_request: dict[str, Any], response: dict[str, Any], ctx: WorkflowContext[Any, str] + ) -> None: + del original_request + self._handled.append(response["selection"]) + await ctx.yield_output(f"Chose {response['selection']}") + + +def _choice_resume(thread_id: str | None = None) -> dict[str, Any]: + payload: dict[str, Any] = { + "messages": [], + "resume": {"interrupts": [{"id": "choice", "value": {"selection": "a"}}]}, + } + if thread_id is not None: + payload["thread_id"] = thread_id + return payload + + +def _reported_thread_ids(events: list[Any]) -> list[str]: + return [event.thread_id for event in events if event.type == "RUN_STARTED"] + + async def test_workflow_wrapper_rejects_workflow_and_factory_at_once() -> None: """Workflow wrapper should reject ambiguous workflow source configuration.""" @@ -141,6 +178,205 @@ def workflow_factory(thread_id: str) -> Workflow: assert factory_calls["thread-a"] == 2 +async def test_workflow_wrapper_omitted_thread_id_binds_live_interrupt_to_reported_id() -> None: + """A generated thread id is reported consistently and is the only id that continues the interrupt.""" + handled: list[str] = [] + agent = AgentFrameworkWorkflow(workflow=WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build()) + + first_events = await _run(agent, {"messages": [{"role": "user", "content": "start"}]}) + first_finished = [event for event in first_events if event.type == "RUN_FINISHED"] + assert _interrupts_from_finished(first_finished[0])[0]["id"] == "choice" + reported = {event.thread_id for event in first_events if event.type in {"RUN_STARTED", "RUN_FINISHED"}} + assert len(reported) == 1 + thread_id = reported.pop() + + for rejected in (_choice_resume(), _choice_resume("other-thread")): + events = await _run(agent, rejected) + assert [event.code for event in events if event.type == "RUN_ERROR"] == ["WORKFLOW_RESUME_NOT_FOUND"] + assert handled == [] + + resumed = await _run(agent, _choice_resume(thread_id)) + assert "RUN_ERROR" not in [event.type for event in resumed] + assert _reported_thread_ids(resumed) == [thread_id] + assert handled == ["a"] + + +async def test_workflow_wrapper_factory_receives_reported_thread_id_when_omitted() -> None: + """Echoing the reported thread id reuses the factory instance that holds the pending interrupt.""" + handled: list[str] = [] + factory_thread_ids: list[str] = [] + + def workflow_factory(thread_id: str) -> Workflow: + factory_thread_ids.append(thread_id) + return WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build() + + agent = AgentFrameworkWorkflow(workflow_factory=workflow_factory) + + first_events = await _run(agent, {"messages": [{"role": "user", "content": "start"}]}) + reported = _reported_thread_ids(first_events) + assert len(reported) == 1 + assert factory_thread_ids == reported + + resumed = await _run(agent, _choice_resume(reported[0])) + assert "RUN_ERROR" not in [event.type for event in resumed] + assert handled == ["a"] + assert factory_thread_ids == reported + + await _run(agent, {"messages": [{"role": "user", "content": "start again"}]}) + assert len(factory_thread_ids) == 2 + assert factory_thread_ids[1] != reported[0] + + +async def test_workflow_wrapper_treats_empty_thread_ids_as_omitted() -> None: + """Empty snake_case and camelCase thread ids resolve to one generated id like an omitted id.""" + handled: list[str] = [] + agent = AgentFrameworkWorkflow(workflow=WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build()) + + first_events = await _run( + agent, {"thread_id": "", "threadId": "", "messages": [{"role": "user", "content": "start"}]} + ) + reported = _reported_thread_ids(first_events) + assert len(reported) == 1 + assert reported[0] + + resumed = await _run(agent, _choice_resume(reported[0])) + assert "RUN_ERROR" not in [event.type for event in resumed] + assert handled == ["a"] + + +async def _newest_pending_checkpoint(storage: InMemoryCheckpointStorage, workflow: Workflow) -> Any: + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + return max((item for item in checkpoints if item.pending_request_info_events), key=lambda item: item.timestamp) + + +@pytest.mark.parametrize( + "empty_ids", + [{"thread_id": ""}, {"threadId": ""}, {"thread_id": "", "threadId": ""}], + ids=["snake-case", "camel-case", "both"], +) +async def test_workflow_wrapper_empty_thread_ids_own_checkpoints_like_omitted(empty_ids: dict[str, str]) -> None: + """An empty thread id is not stamped as checkpoint owner, so id-less restores keep working.""" + storage = InMemoryCheckpointStorage() + handled: list[str] = [] + first_workflow = WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build() + first = AgentFrameworkWorkflow(workflow=first_workflow, checkpoint_storage=storage) + await _run(first, {"messages": [{"role": "user", "content": "start"}], **empty_ids}) + + checkpoint = await _newest_pending_checkpoint(storage, first_workflow) + assert checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY]["thread_id"] is None + + restarted = AgentFrameworkWorkflow( + workflow=WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build(), checkpoint_storage=storage + ) + resumed = await _run( + restarted, {**_choice_resume(), "forwarded_props": {"checkpoint_id": checkpoint.checkpoint_id}} + ) + assert "RUN_ERROR" not in [event.type for event in resumed] + assert handled == ["a"] + + +_FACTORY_CONTINUATIONS: dict[str, dict[str, Any]] = { + "resume-entry": _choice_resume(), + "tool-message": {"messages": [{"role": "tool", "toolCallId": "choice", "content": '{"selection": "a"}'}]}, + "approval-message": { + "messages": [ + { + "role": "assistant", + "content": "", + "toolCalls": [{"id": "choice", "type": "function", "function": {"name": "choose", "arguments": "{}"}}], + }, + {"role": "tool", "toolCallId": "choice", "content": '{"accepted": true}'}, + ] + }, +} + + +@pytest.mark.parametrize("continuation", list(_FACTORY_CONTINUATIONS.values()), ids=list(_FACTORY_CONTINUATIONS)) +async def test_workflow_wrapper_factory_rejects_continuations_without_thread_id( + continuation: dict[str, Any], +) -> None: + """A resume entry, tool result, or approval response cannot continue a factory run without a thread id.""" + handled: list[str] = [] + factory_thread_ids: list[str] = [] + + def workflow_factory(thread_id: str) -> Workflow: + factory_thread_ids.append(thread_id) + return WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build() + + agent = AgentFrameworkWorkflow(workflow_factory=workflow_factory) + + first_events = await _run(agent, {"messages": [{"role": "user", "content": "start"}]}) + reported = _reported_thread_ids(first_events) + assert factory_thread_ids == reported + + rejected = await _run(agent, deepcopy(continuation)) + assert [event.code for event in rejected if event.type == "RUN_ERROR"] == ["WORKFLOW_RESUME_NOT_FOUND"] + assert factory_thread_ids == reported + assert handled == [] + + resumed = await _run(agent, _choice_resume(reported[0])) + assert "RUN_ERROR" not in [event.type for event in resumed] + assert handled == ["a"] + + +async def test_workflow_wrapper_factory_runs_requests_without_responses_when_thread_id_is_omitted() -> None: + """Only continuations are rejected without a thread id; a new request still starts a run under a new id.""" + handled: list[str] = [] + factory_thread_ids: list[str] = [] + + def workflow_factory(thread_id: str) -> Workflow: + factory_thread_ids.append(thread_id) + return WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build() + + agent = AgentFrameworkWorkflow(workflow_factory=workflow_factory) + + first = _reported_thread_ids(await _run(agent, {"messages": [{"role": "user", "content": "start"}]})) + second = _reported_thread_ids(await _run(agent, {"messages": [{"role": "user", "content": "start again"}]})) + + assert factory_thread_ids == first + second + assert first != second + assert handled == [] + + +async def test_workflow_wrapper_checkpoint_resumes_with_reported_or_omitted_thread_id() -> None: + """A checkpoint of a run without threadId resumes with the reported id or, as before, with no id.""" + storage = InMemoryCheckpointStorage() + handled: list[str] = [] + + def restart() -> AgentFrameworkWorkflow: + workflow = WorkflowBuilder(start_executor=_ChoiceRecorder(handled)).build() + return AgentFrameworkWorkflow(workflow=workflow, checkpoint_storage=storage) + + first = restart() + first_events = await _run( + first, {"messages": [{"role": "user", "content": "start"}], _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a"} + ) + reported = _reported_thread_ids(first_events)[0] + assert first.workflow is not None + checkpoint = await _newest_pending_checkpoint(storage, first.workflow) + assert checkpoint.metadata[_CHECKPOINT_REQUEST_OWNER_KEY] == { + "snapshot_scope": "tenant-a", + "thread_id": None, + "generated_thread_id": reported, + } + forwarded = {"forwarded_props": {"checkpoint_id": checkpoint.checkpoint_id}} + + for wrong_thread_id, wrong_scope in (("other-thread", "tenant-a"), (reported, "tenant-b")): + rejected = await _run( + restart(), {**_choice_resume(wrong_thread_id), **forwarded, _SNAPSHOT_SCOPE_INPUT_KEY: wrong_scope} + ) + assert [event.code for event in rejected if event.type == "RUN_ERROR"] == ["WORKFLOW_RESUME_NOT_FOUND"] + assert handled == [] + + echoed = await _run(restart(), {**_choice_resume(reported), **forwarded, _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a"}) + assert "RUN_ERROR" not in [event.type for event in echoed] + assert handled == ["a"] + + legacy = await _run(restart(), {**_choice_resume(), **forwarded, _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a"}) + assert "RUN_ERROR" not in [event.type for event in legacy] + assert handled == ["a", "a"] + + async def test_workflow_wrapper_without_workflow_raises_not_implemented() -> None: """Without workflow/workflow_factory, run should raise NotImplementedError.""" agent = AgentFrameworkWorkflow()