# Copyright (c) Microsoft. All rights reserved. """Tests for FastAPI endpoint creation (_endpoint.py).""" import json from typing import Any import pytest from ag_ui.core import MessagesSnapshotEvent, RunStartedEvent, StateSnapshotEvent from agent_framework import ( Agent, ChatResponseUpdate, Content, WorkflowBuilder, WorkflowContext, executor, ) from agent_framework.orchestrations import SequentialBuilder from fastapi import FastAPI, Header, HTTPException from fastapi.params import Depends from fastapi.testclient import TestClient from agent_framework_ag_ui import InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint from agent_framework_ag_ui._agent import AgentFrameworkAgent from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow def _decode_sse_events(response: Any) -> list[dict[str, Any]]: content = response.content.decode("utf-8") return [json.loads(line[6:]) for line in content.splitlines() if line.startswith("data: ")] def _latest_messages_snapshot(response: Any) -> list[dict[str, Any]]: snapshots = [ event["messages"] for event in _decode_sse_events(response) if event.get("type") == "MESSAGES_SNAPSHOT" ] assert snapshots return snapshots[-1] @pytest.fixture def build_chat_client(streaming_chat_client_stub, stream_from_updates_fixture): """Create a typed chat client stub for endpoint tests.""" def _build(response_text: str = "Test response"): updates = [ChatResponseUpdate(contents=[Content.from_text(text=response_text)])] return streaming_chat_client_stub(stream_from_updates_fixture(updates)) return _build async def test_add_endpoint_with_agent_protocol(build_chat_client): """Test adding endpoint with raw SupportsAgentRun.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent") client = TestClient(app) response = client.post("/test-agent", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" async def test_add_endpoint_with_wrapped_agent(build_chat_client): """Test adding endpoint with pre-wrapped AgentFrameworkAgent.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped") add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent") client = TestClient(app) response = client.post("/wrapped-agent", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" async def test_add_endpoint_with_workflow_protocol(): """Test adding endpoint with native Workflow support.""" @executor(id="start") async def start(message: Any, ctx: WorkflowContext) -> None: await ctx.yield_output("Workflow response") app = FastAPI() workflow = WorkflowBuilder(start_executor=start).build() add_agent_framework_fastapi_endpoint(app, workflow, path="/workflow") client = TestClient(app) response = client.post("/workflow", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" 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_with_state_schema(build_chat_client): """Test endpoint with state_schema parameter.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) state_schema = {"document": {"type": "string"}} add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema) client = TestClient(app) response = client.post( "/stateful", json={"messages": [{"role": "user", "content": "Hello"}], "state": {"document": ""}} ) assert response.status_code == 200 async def test_endpoint_with_default_state_seed(build_chat_client): """Test endpoint seeds default state when client omits it.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) state_schema = {"proverbs": {"type": "array"}} default_state = {"proverbs": ["Keep the original."]} add_agent_framework_fastapi_endpoint( app, agent, path="/default-state", state_schema=state_schema, default_state=default_state, ) client = TestClient(app) response = client.post("/default-state", 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: ")] snapshots = [json.loads(line[6:]) for line in lines if json.loads(line[6:]).get("type") == "STATE_SNAPSHOT"] assert snapshots, "Expected a STATE_SNAPSHOT event" assert snapshots[0]["snapshot"]["proverbs"] == default_state["proverbs"] async def test_endpoint_with_predict_state_config(build_chat_client): """Test endpoint with predict_state_config parameter.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}} add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config) client = TestClient(app) response = client.post("/predictive", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 async def test_endpoint_request_logging(build_chat_client): """Test that endpoint logs request details.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/logged") client = TestClient(app) response = client.post( "/logged", json={ "messages": [{"role": "user", "content": "Test"}], "run_id": "run-123", "thread_id": "thread-456", }, ) assert response.status_code == 200 async def test_endpoint_event_streaming(build_chat_client): """Test that endpoint streams events correctly.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client("Streamed response")) add_agent_framework_fastapi_endpoint(app, agent, path="/stream") client = TestClient(app) response = client.post("/stream", 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.strip()] found_run_started = False found_text_content = False found_run_finished = False for line in lines: if line.startswith("data: "): event_data = json.loads(line[6:]) if event_data.get("type") == "RUN_STARTED": found_run_started = True elif event_data.get("type") == "TEXT_MESSAGE_CONTENT": found_text_content = True elif event_data.get("type") == "RUN_FINISHED": found_run_finished = True assert found_run_started assert found_text_content assert found_run_finished 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") 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()) add_agent_framework_fastapi_endpoint(app, agent, path="/failing") client = TestClient(app) # Send invalid JSON to trigger parsing error before streaming response = client.post("/failing", data=b"invalid json", headers={"content-type": "application/json"}) # type: ignore # Pydantic validation now returns 422 for invalid request body assert response.status_code == 422 async def test_endpoint_multiple_paths(build_chat_client): """Test adding multiple endpoints with different paths.""" app = FastAPI() agent1 = Agent(name="agent1", instructions="First agent", client=build_chat_client("Response 1")) agent2 = Agent(name="agent2", instructions="Second agent", client=build_chat_client("Response 2")) add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1") add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2") client = TestClient(app) response1 = client.post("/agent1", json={"messages": [{"role": "user", "content": "Hi"}]}) response2 = client.post("/agent2", json={"messages": [{"role": "user", "content": "Hi"}]}) assert response1.status_code == 200 assert response2.status_code == 200 async def test_endpoint_default_path(build_chat_client): """Test endpoint with default path.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent) client = TestClient(app) response = client.post("/", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 async def test_endpoint_response_headers(build_chat_client): """Test that endpoint sets correct response headers.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/headers") client = TestClient(app) response = client.post("/headers", json={"messages": [{"role": "user", "content": "Test"}]}) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" assert "cache-control" in response.headers assert response.headers["cache-control"] == "no-cache" async def test_endpoint_empty_messages(streaming_chat_client_stub): """Empty messages keep the existing no-op run behavior when snapshot persistence is not configured.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 yield ChatResponseUpdate(contents=[Content.from_text(text="Should not run")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) add_agent_framework_fastapi_endpoint(app, agent, path="/empty") client = TestClient(app) response = client.post("/empty", json={"messages": []}) assert response.status_code == 200 assert call_count == 0 assert [event.get("type") for event in _decode_sse_events(response)] == ["RUN_STARTED", "RUN_FINISHED"] async def test_endpoint_complex_input(build_chat_client): """Test endpoint with complex input data.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/complex") client = TestClient(app) response = client.post( "/complex", json={ "messages": [ {"role": "user", "content": "First message", "id": "msg-1"}, {"role": "assistant", "content": "Response", "id": "msg-2"}, {"role": "user", "content": "Follow-up", "id": "msg-3"}, ], "run_id": "complex-run-123", "thread_id": "complex-thread-456", "state": {"custom_field": "value"}, }, ) assert response.status_code == 200 async def test_endpoint_openapi_schema(build_chat_client): """Test that endpoint generates proper OpenAPI schema with request model.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/schema-test") client = TestClient(app) response = client.get("/openapi.json") assert response.status_code == 200 openapi_spec = response.json() # Verify the endpoint exists in the schema assert "/schema-test" in openapi_spec["paths"] endpoint_spec = openapi_spec["paths"]["/schema-test"]["post"] # Verify request body schema is defined assert "requestBody" in endpoint_spec request_body = endpoint_spec["requestBody"] assert "content" in request_body assert "application/json" in request_body["content"] # Verify schema references AGUIRequest model schema_ref = request_body["content"]["application/json"]["schema"] assert "$ref" in schema_ref assert "AGUIRequest" in schema_ref["$ref"] # Verify AGUIRequest model is in components assert "components" in openapi_spec assert "schemas" in openapi_spec["components"] assert "AGUIRequest" in openapi_spec["components"]["schemas"] # Verify AGUIRequest has required fields agui_request_schema = openapi_spec["components"]["schemas"]["AGUIRequest"] assert "properties" in agui_request_schema assert "messages" in agui_request_schema["properties"] assert "run_id" in agui_request_schema["properties"] assert "thread_id" in agui_request_schema["properties"] assert "state" in agui_request_schema["properties"] assert "required" in agui_request_schema assert "messages" in agui_request_schema["required"] async def test_endpoint_default_tags(build_chat_client): """Test that endpoint uses default 'AG-UI' tag.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/default-tags") client = TestClient(app) response = client.get("/openapi.json") assert response.status_code == 200 openapi_spec = response.json() endpoint_spec = openapi_spec["paths"]["/default-tags"]["post"] assert "tags" in endpoint_spec assert endpoint_spec["tags"] == ["AG-UI"] async def test_endpoint_custom_tags(build_chat_client): """Test that endpoint accepts custom tags.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/custom-tags", tags=["Custom", "Agent"]) client = TestClient(app) response = client.get("/openapi.json") assert response.status_code == 200 openapi_spec = response.json() endpoint_spec = openapi_spec["paths"]["/custom-tags"]["post"] assert "tags" in endpoint_spec assert endpoint_spec["tags"] == ["Custom", "Agent"] async def test_endpoint_missing_required_field(build_chat_client): """Test that endpoint validates required fields with Pydantic.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) add_agent_framework_fastapi_endpoint(app, agent, path="/validation") client = TestClient(app) # Missing required 'messages' field should trigger validation error response = client.post("/validation", json={"run_id": "test-123"}) assert response.status_code == 422 error_detail = response.json() assert "detail" in error_detail async def test_endpoint_internal_error_handling(build_chat_client): """Test endpoint error handling when an exception occurs before streaming starts.""" from unittest.mock import patch app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) # Use default_state to trigger the code path that can raise an exception add_agent_framework_fastapi_endpoint(app, agent, path="/error-test", default_state={"key": "value"}) client = TestClient(app) # Mock copy.deepcopy to raise an exception during default_state processing with patch("agent_framework_ag_ui._endpoint.copy.deepcopy") as mock_deepcopy: mock_deepcopy.side_effect = Exception("Simulated internal error") response = client.post("/error-test", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 500 assert response.json() == {"detail": "An internal error has occurred."} async def test_endpoint_streaming_error_emits_run_error_event(): """Streaming exceptions should emit RUN_ERROR instead of terminating silently.""" class FailingStreamWorkflow(AgentFrameworkWorkflow): async def run(self, input_data: dict[str, Any]): del input_data yield RunStartedEvent(run_id="run-1", thread_id="thread-1") raise RuntimeError("stream exploded") app = FastAPI() add_agent_framework_fastapi_endpoint(app, FailingStreamWorkflow(), path="/stream-error") client = TestClient(app) response = client.post("/stream-error", 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 "RUN_ERROR" in event_types async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client): """Test that endpoint blocks requests when authentication dependency fails.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) async def require_api_key(x_api_key: str | None = Header(None)): if x_api_key != "secret-key": raise HTTPException(status_code=401, detail="Unauthorized") add_agent_framework_fastapi_endpoint(app, agent, path="/protected", dependencies=[Depends(require_api_key)]) client = TestClient(app) # Request without API key should be rejected response = client.post("/protected", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 401 assert response.json()["detail"] == "Unauthorized" async def test_endpoint_with_dependencies_allows_authorized(build_chat_client): """Test that endpoint allows requests when authentication dependency passes.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) async def require_api_key(x_api_key: str | None = Header(None)): if x_api_key != "secret-key": raise HTTPException(status_code=401, detail="Unauthorized") add_agent_framework_fastapi_endpoint(app, agent, path="/protected", dependencies=[Depends(require_api_key)]) client = TestClient(app) # Request with valid API key should succeed response = client.post( "/protected", json={"messages": [{"role": "user", "content": "Hello"}]}, headers={"x-api-key": "secret-key"}, ) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" async def test_endpoint_with_multiple_dependencies(build_chat_client): """Test that endpoint supports multiple dependencies.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) execution_order: list[str] = [] async def first_dependency(): execution_order.append("first") async def second_dependency(): execution_order.append("second") add_agent_framework_fastapi_endpoint( app, agent, path="/multi-deps", dependencies=[Depends(first_dependency), Depends(second_dependency)], ) client = TestClient(app) response = client.post("/multi-deps", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 assert "first" in execution_order assert "second" in execution_order async def test_endpoint_without_dependencies_is_accessible(build_chat_client): """Test that endpoint without dependencies remains accessible (backward compatibility).""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) # No dependencies parameter - should be accessible without auth add_agent_framework_fastapi_endpoint(app, agent, path="/open") client = TestClient(app) response = client.post("/open", json={"messages": [{"role": "user", "content": "Hello"}]}) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" async def test_endpoint_invalid_agent_type_raises_typeerror(): """Passing an invalid agent type raises TypeError.""" app = FastAPI() with pytest.raises(TypeError, match="must be SupportsAgentRun"): add_agent_framework_fastapi_endpoint(app, agent="not_an_agent") # type: ignore[arg-type] async def test_endpoint_requires_snapshot_scope_resolver_when_store_configured(build_chat_client): """Snapshot persistence setup must require an explicit Snapshot Scope resolver.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) store = InMemoryAGUIThreadSnapshotStore() with pytest.raises(ValueError, match="snapshot_scope_resolver is required"): add_agent_framework_fastapi_endpoint(app, agent, path="/snapshots", snapshot_store=store) async def test_endpoint_requires_snapshot_scope_resolver_when_wrapped_runner_has_store(build_chat_client): """Pre-wrapped runners with snapshot stores must also provide a Snapshot Scope resolver.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) wrapped_agent = AgentFrameworkAgent(agent=agent, snapshot_store=InMemoryAGUIThreadSnapshotStore()) with pytest.raises(ValueError, match="snapshot_scope_resolver is required"): add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/snapshots") async def test_endpoint_accepts_snapshot_store_with_scope_resolver(build_chat_client): """Endpoint behavior remains the normal event stream when snapshot persistence is explicitly configured.""" app = FastAPI() agent = Agent(name="test", instructions="Test agent", client=build_chat_client()) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) response = client.post( "/snapshots", json={"messages": [{"role": "user", "content": "Hello"}], "thread_id": "thread-1"}, ) assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" async def test_agent_endpoint_hydrates_stored_thread_snapshot_without_invoking_agent(streaming_chat_client_stub): """A Hydrate Request replays stored agent messages and state without invoking the wrapped agent.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 yield ChatResponseUpdate(contents=[Content.from_text(text="Stored reply")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"recipe": {"type": "string"}}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"role": "user", "content": "Hello"}], "state": {"recipe": "pasta"}, }, ) assert first_response.status_code == 200 assert call_count == 1 hydrate_response = client.post("/snapshots", json={"thread_id": "thread-1", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 1 events = _decode_sse_events(hydrate_response) event_types = [event.get("type") for event in events] assert event_types == ["RUN_STARTED", "STATE_SNAPSHOT", "MESSAGES_SNAPSHOT", "RUN_FINISHED"] assert events[1]["snapshot"] == {"recipe": "pasta"} assert any(message.get("role") == "user" and message.get("content") == "Hello" for message in events[2]["messages"]) assert any( message.get("role") == "assistant" and message.get("content") == "Stored reply" for message in events[2]["messages"] ) async def test_agent_endpoint_hydrates_snapshots_by_scope_and_thread(streaming_chat_client_stub): """Hydration uses Snapshot Scope and AG-UI Thread id together when reading stored snapshots.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 yield ChatResponseUpdate(contents=[Content.from_text(text="Tenant A reply")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"tenant": {"type": "string"}}, snapshot_store=store, snapshot_scope_resolver=lambda request: request.forwarded_props["tenant"], ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"role": "user", "content": "Hello tenant A"}], "state": {"tenant": "tenant-a"}, "forwardedProps": {"tenant": "tenant-a"}, }, ) assert first_response.status_code == 200 assert call_count == 1 tenant_b_response = client.post( "/snapshots", json={"thread_id": "thread-1", "messages": [], "forwardedProps": {"tenant": "tenant-b"}}, ) assert tenant_b_response.status_code == 200 assert call_count == 1 assert [event.get("type") for event in _decode_sse_events(tenant_b_response)] == [ "RUN_STARTED", "RUN_FINISHED", ] tenant_a_response = client.post( "/snapshots", json={"thread_id": "thread-1", "messages": [], "forwardedProps": {"tenant": "tenant-a"}}, ) assert tenant_a_response.status_code == 200 assert call_count == 1 tenant_a_events = _decode_sse_events(tenant_a_response) assert [event.get("type") for event in tenant_a_events] == [ "RUN_STARTED", "STATE_SNAPSHOT", "MESSAGES_SNAPSHOT", "RUN_FINISHED", ] assert tenant_a_events[1]["snapshot"] == {"tenant": "tenant-a"} assert any(message.get("content") == "Tenant A reply" for message in tenant_a_events[2]["messages"]) async def test_agent_endpoint_prepends_stored_snapshot_for_new_user_turn(streaming_chat_client_stub): """A normal agent turn with a known thread id prepends stored history and keeps the new user input.""" app = FastAPI() captured_messages: list[list[tuple[str, str]]] = [] async def stream_fn(messages: Any, options: Any, **kwargs: Any): del options, kwargs captured_messages.append([(message.role, message.text) for message in messages]) yield ChatResponseUpdate(contents=[Content.from_text(text=f"Reply {len(captured_messages)}")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"recipe": {"type": "string"}}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"id": "user-1", "role": "user", "content": "Plan dinner"}], "state": {"recipe": "pasta"}, }, ) assert first_response.status_code == 200 second_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"id": "user-2", "role": "user", "content": "Add dessert"}], }, ) assert second_response.status_code == 200 assert len(captured_messages) == 2 assert captured_messages[1] == [ ("user", "Plan dinner"), ("assistant", "Reply 1"), ( "system", ( "Current state of the application:\n" '{\n "recipe": "pasta"\n}\n\n' "When modifying state, you MUST include ALL existing data plus your changes.\n" "For example, if adding one new item to a list, include ALL existing items PLUS the new item.\n" "Never replace existing data - always preserve and append or merge." ), ), ("user", "Add dessert"), ] events = _decode_sse_events(second_response) state_snapshots = [event for event in events if event.get("type") == "STATE_SNAPSHOT"] assert state_snapshots[0]["snapshot"] == {"recipe": "pasta"} async def test_agent_endpoint_deduplicates_full_history_and_merges_fresh_state(streaming_chat_client_stub): """Stored prior history is authoritative while incoming full history and fresh state remain supported.""" app = FastAPI() captured_messages: list[list[tuple[str, str]]] = [] async def stream_fn(messages: Any, options: Any, **kwargs: Any): del options, kwargs captured_messages.append([(message.role, message.text) for message in messages]) yield ChatResponseUpdate(contents=[Content.from_text(text=f"Reply {len(captured_messages)}")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"recipe": {"type": "string"}, "theme": {"type": "string"}}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"id": "user-1", "role": "user", "content": "Plan dinner"}], "state": {"recipe": "pasta", "theme": "dark"}, }, ) assert first_response.status_code == 200 first_snapshot = _latest_messages_snapshot(first_response) second_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [*first_snapshot, {"id": "user-2", "role": "user", "content": "Add dessert"}], "state": {"recipe": "salad"}, }, ) assert second_response.status_code == 200 second_non_system_messages = [message for message in captured_messages[1] if message[0] != "system"] assert second_non_system_messages == [ ("user", "Plan dinner"), ("assistant", "Reply 1"), ("user", "Add dessert"), ] second_events = _decode_sse_events(second_response) second_state_snapshots = [event for event in second_events if event.get("type") == "STATE_SNAPSHOT"] assert second_state_snapshots[0]["snapshot"] == {"recipe": "salad", "theme": "dark"} second_snapshot = _latest_messages_snapshot(second_response) conflicting_history = [message.copy() for message in second_snapshot] conflicting_history[0]["content"] = "Tampered dinner plan" conflicting_history[1]["content"] = "Tampered reply" third_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [*conflicting_history, {"id": "user-3", "role": "user", "content": "Pick wine"}], }, ) assert third_response.status_code == 200 third_texts = [text for role, text in captured_messages[2] if role != "system"] assert third_texts == ["Plan dinner", "Reply 1", "Add dessert", "Reply 2", "Pick wine"] assert "Tampered dinner plan" not in third_texts assert "Tampered reply" not in third_texts third_state_snapshots = [ event for event in _decode_sse_events(third_response) if event.get("type") == "STATE_SNAPSHOT" ] assert third_state_snapshots[0]["snapshot"] == {"recipe": "salad", "theme": "dark"} async def test_agent_endpoint_hydrates_interrupted_thread_without_invoking_agent(streaming_chat_client_stub): """Hydrating an interrupted agent replays state, messages, and interrupt metadata without resuming it.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 yield ChatResponseUpdate( contents=[ Content.from_function_call( name="draft_steps", call_id="draft-call", arguments=json.dumps({"steps": [{"description": "Draft outline"}]}), ) ], role="assistant", ) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "agent-thread", "messages": [{"role": "user", "content": "Draft the plan"}], "state": {"steps": []}, }, ) assert first_response.status_code == 200 assert call_count == 1 first_events = _decode_sse_events(first_response) first_finished = [event for event in first_events if event.get("type") == "RUN_FINISHED"] assert first_finished[-1]["interrupt"][0]["value"]["function_call"]["call_id"] == "draft-call" hydrate_response = client.post("/snapshots", json={"thread_id": "agent-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 1 events = _decode_sse_events(hydrate_response) assert [event.get("type") for event in events] == [ "RUN_STARTED", "STATE_SNAPSHOT", "MESSAGES_SNAPSHOT", "RUN_FINISHED", ] assert events[1]["snapshot"] == {"steps": [{"description": "Draft outline"}]} assert events[-1]["interrupt"][0]["value"]["function_call"]["name"] == "draft_steps" async def test_agent_endpoint_run_error_does_not_overwrite_previous_snapshot(streaming_chat_client_stub): """A failing agent turn leaves the last good AG-UI Thread Snapshot available for hydration.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 if call_count == 1: yield ChatResponseUpdate(contents=[Content.from_text(text="Stable reply")]) return raise RuntimeError("agent exploded") agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={"thread_id": "agent-thread", "messages": [{"role": "user", "content": "Start"}]}, ) assert first_response.status_code == 200 assert call_count == 1 error_response = client.post( "/snapshots", json={"thread_id": "agent-thread", "messages": [{"role": "user", "content": "Break the run"}]}, ) assert error_response.status_code == 200 assert call_count == 2 assert "RUN_ERROR" in [event.get("type") for event in _decode_sse_events(error_response)] hydrate_response = client.post("/snapshots", json={"thread_id": "agent-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 2 messages = _latest_messages_snapshot(hydrate_response) assert any(message.get("role") == "assistant" and message.get("content") == "Stable reply" for message in messages) assert not any(message.get("content") == "Break the run" for message in messages) async def test_workflow_endpoint_hydrates_emitted_snapshots_without_invoking_workflow(): """A workflow Hydrate Request replays emitted snapshots without invoking the wrapped workflow.""" app = FastAPI() call_count = 0 @executor(id="snapshotter") async def snapshotter(message: Any, ctx: WorkflowContext) -> None: nonlocal call_count del message call_count += 1 await ctx.yield_output(StateSnapshotEvent(snapshot={"active_agent": "flights"})) await ctx.yield_output( MessagesSnapshotEvent( messages=[{"id": "assistant-snapshot", "role": "assistant", "content": "Stored workflow reply"}] ) ) workflow = WorkflowBuilder(start_executor=snapshotter).build() store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, workflow, path="/workflow-snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": [{"role": "user", "content": "Start workflow"}]}, ) assert first_response.status_code == 200 assert call_count == 1 hydrate_response = client.post("/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 1 events = _decode_sse_events(hydrate_response) assert [event.get("type") for event in events] == [ "RUN_STARTED", "STATE_SNAPSHOT", "MESSAGES_SNAPSHOT", "RUN_FINISHED", ] assert events[1]["snapshot"] == {"active_agent": "flights"} assert events[2]["messages"] == [ {"id": "assistant-snapshot", "role": "assistant", "content": "Stored workflow reply"} ] async def test_workflow_endpoint_hydrates_synthesized_text_and_tool_snapshot(): """Workflow text and tool output are synthesized into replayable snapshot messages.""" app = FastAPI() call_count = 0 @executor(id="responder") async def responder(message: Any, ctx: WorkflowContext) -> None: nonlocal call_count del message call_count += 1 await ctx.yield_output("Workflow answer") await ctx.yield_output( [ Content.from_function_call( name="lookup_weather", call_id="call-1", arguments='{"city":"SF"}', ), Content.from_function_result(call_id="call-1", result="72F"), ] ) await ctx.yield_output({"diagnostic": "not persisted"}) workflow = WorkflowBuilder(start_executor=responder).build() store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, workflow, path="/workflow-snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/workflow-snapshots", json={ "thread_id": "workflow-thread", "messages": [{"id": "user-1", "role": "user", "content": "Start workflow"}], }, ) assert first_response.status_code == 200 assert call_count == 1 hydrate_response = client.post("/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 1 events = _decode_sse_events(hydrate_response) assert [event.get("type") for event in events] == ["RUN_STARTED", "MESSAGES_SNAPSHOT", "RUN_FINISHED"] messages = events[1]["messages"] assert any(message.get("role") == "user" and message.get("content") == "Start workflow" for message in messages) assert any( message.get("role") == "assistant" and message.get("content") == "Workflow answer" for message in messages ) tool_call_messages = [ message for message in messages if message.get("role") == "assistant" and message.get("toolCalls") ] assert len(tool_call_messages) == 1 tool_call = tool_call_messages[0]["toolCalls"][0] assert tool_call["id"] == "call-1" assert tool_call["function"] == {"name": "lookup_weather", "arguments": '{"city":"SF"}'} assert any( message.get("role") == "tool" and message.get("toolCallId") == "call-1" and message.get("content") == "72F" for message in messages ) async def test_workflow_endpoint_hydrates_interrupted_thread_without_invoking_workflow(): """Hydrating an interrupted workflow replays state, messages, and interrupt metadata without resuming it.""" app = FastAPI() call_count = 0 @executor(id="requester") async def requester(message: Any, ctx: WorkflowContext) -> None: nonlocal call_count del message call_count += 1 await ctx.yield_output(StateSnapshotEvent(snapshot={"step": "approval"})) await ctx.request_info( {"message": "Approve workflow step", "options": ["Approve", "Reject"]}, dict, request_id="workflow-approval", ) workflow = WorkflowBuilder(start_executor=requester).build() store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, workflow, path="/workflow-snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": [{"role": "user", "content": "Start workflow"}]}, ) assert first_response.status_code == 200 assert call_count == 1 first_finished = [event for event in _decode_sse_events(first_response) if event.get("type") == "RUN_FINISHED"] assert first_finished[-1]["interrupt"][0]["id"] == "workflow-approval" hydrate_response = client.post("/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 1 events = _decode_sse_events(hydrate_response) assert [event.get("type") for event in events] == [ "RUN_STARTED", "STATE_SNAPSHOT", "MESSAGES_SNAPSHOT", "RUN_FINISHED", ] assert events[1]["snapshot"] == {"step": "approval"} assert events[-1]["interrupt"][0]["id"] == "workflow-approval" assert events[-1]["interrupt"][0]["value"]["message"] == "Approve workflow step" async def test_workflow_endpoint_run_error_does_not_overwrite_previous_snapshot(): """A failing workflow turn leaves the last good AG-UI Thread Snapshot available for hydration.""" app = FastAPI() call_count = 0 @executor(id="responder") async def responder(message: Any, ctx: WorkflowContext) -> None: nonlocal call_count del message call_count += 1 if call_count == 1: await ctx.yield_output("Stable workflow reply") return raise RuntimeError("workflow exploded") workflow = WorkflowBuilder(start_executor=responder).build() store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, workflow, path="/workflow-snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": [{"role": "user", "content": "Start workflow"}]}, ) assert first_response.status_code == 200 assert call_count == 1 error_response = client.post( "/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": [{"role": "user", "content": "Break workflow"}]}, ) assert error_response.status_code == 200 assert call_count == 2 assert "RUN_ERROR" in [event.get("type") for event in _decode_sse_events(error_response)] hydrate_response = client.post("/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 2 messages = _latest_messages_snapshot(hydrate_response) assert any( message.get("role") == "assistant" and message.get("content") == "Stable workflow reply" for message in messages ) assert not any(message.get("content") == "Break workflow" for message in messages) async def test_endpoint_encoding_failure_emits_run_error(): """Event encoding failure emits RUN_ERROR event in the SSE stream.""" from unittest.mock import patch class SimpleWorkflow(AgentFrameworkWorkflow): async def run(self, input_data: dict[str, Any]): del input_data yield RunStartedEvent(run_id="run-1", thread_id="thread-1") app = FastAPI() add_agent_framework_fastapi_endpoint(app, SimpleWorkflow(), path="/encode-fail") client = TestClient(app) with patch("ag_ui.encoder.EventEncoder.encode") as mock_encode: # First call fails (the RUN_STARTED event), second call succeeds (the error event) mock_encode.side_effect = [ValueError("encode boom"), 'data: {"type":"RUN_ERROR"}\n\n'] response = client.post("/encode-fail", json={"messages": [{"role": "user", "content": "go"}]}) assert response.status_code == 200 content = response.content.decode("utf-8") assert "RUN_ERROR" in content async def test_endpoint_double_encoding_failure_terminates(): """When both event and error encoding fail, stream terminates gracefully.""" from unittest.mock import patch class SimpleWorkflow(AgentFrameworkWorkflow): async def run(self, input_data: dict[str, Any]): del input_data yield RunStartedEvent(run_id="run-1", thread_id="thread-1") app = FastAPI() add_agent_framework_fastapi_endpoint(app, SimpleWorkflow(), path="/double-fail") client = TestClient(app) with patch("ag_ui.encoder.EventEncoder.encode") as mock_encode: # Both calls fail - event encode and error event encode mock_encode.side_effect = ValueError("always fails") response = client.post("/double-fail", json={"messages": [{"role": "user", "content": "go"}]}) # Should still get 200 (SSE stream), just with no events assert response.status_code == 200 async def test_agent_endpoint_confirm_changes_clears_persisted_interrupt(streaming_chat_client_stub): """A confirm_changes response persists the completed turn and clears the stored interrupt.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 yield ChatResponseUpdate( contents=[ Content.from_function_call( name="draft_steps", call_id="draft-call", arguments=json.dumps({"steps": [{"description": "Draft outline"}]}), ) ], role="assistant", ) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "agent-thread", "messages": [{"id": "user-1", "role": "user", "content": "Draft the plan"}], "state": {"steps": []}, }, ) assert first_response.status_code == 200 assert call_count == 1 first_events = _decode_sse_events(first_response) first_finished = [event for event in first_events if event.get("type") == "RUN_FINISHED"] assert first_finished[-1]["interrupt"] confirm_call_id = first_finished[-1]["interrupt"][0]["id"] confirm_response = client.post( "/snapshots", json={ "thread_id": "agent-thread", "messages": [], "resume": {"interrupts": [{"id": confirm_call_id, "value": json.dumps({"accepted": True, "steps": []})}]}, }, ) assert confirm_response.status_code == 200 assert call_count == 1 confirm_event_types = [event.get("type") for event in _decode_sse_events(confirm_response)] assert "TEXT_MESSAGE_CONTENT" in confirm_event_types hydrate_response = client.post("/snapshots", json={"thread_id": "agent-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 1 events = _decode_sse_events(hydrate_response) assert not events[-1].get("interrupt") messages = _latest_messages_snapshot(hydrate_response) assert any( message.get("role") == "assistant" and message.get("content") == "Changes confirmed and applied successfully!" for message in messages ) assert any(message.get("role") == "user" and message.get("content") == "Draft the plan" for message in messages) async def test_agent_endpoint_default_state_does_not_reset_persisted_state(streaming_chat_client_stub): """Endpoint defaults fill missing keys but never override persisted Shared State.""" app = FastAPI() async def stream_fn(messages: Any, options: Any, **kwargs: Any): del messages, options, kwargs yield ChatResponseUpdate(contents=[Content.from_text(text="Reply")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"recipe": {"type": "string"}}, default_state={"recipe": ""}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) fresh_response = client.post( "/snapshots", json={"thread_id": "thread-fresh", "messages": [{"id": "user-0", "role": "user", "content": "Hi"}]}, ) assert fresh_response.status_code == 200 fresh_state_snapshots = [ event for event in _decode_sse_events(fresh_response) if event.get("type") == "STATE_SNAPSHOT" ] assert fresh_state_snapshots[0]["snapshot"] == {"recipe": ""} first_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"id": "user-1", "role": "user", "content": "Plan dinner"}], "state": {"recipe": "pasta"}, }, ) assert first_response.status_code == 200 second_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"id": "user-2", "role": "user", "content": "Add dessert"}], }, ) assert second_response.status_code == 200 second_state_snapshots = [ event for event in _decode_sse_events(second_response) if event.get("type") == "STATE_SNAPSHOT" ] assert second_state_snapshots[0]["snapshot"] == {"recipe": "pasta"} hydrate_response = client.post("/snapshots", json={"thread_id": "thread-1", "messages": []}) assert hydrate_response.status_code == 200 hydrate_events = _decode_sse_events(hydrate_response) hydrate_state_snapshots = [event for event in hydrate_events if event.get("type") == "STATE_SNAPSHOT"] assert hydrate_state_snapshots[0]["snapshot"] == {"recipe": "pasta"} async def test_agent_endpoint_persists_turn_output_when_intermediate_snapshot_suppressed(streaming_chat_client_stub): """A no-confirmation predictive turn persists tool output even when the outbound snapshot is suppressed.""" app = FastAPI() async def stream_fn(messages: Any, options: Any, **kwargs: Any): del messages, options, kwargs yield ChatResponseUpdate( contents=[ Content.from_function_call( name="write_doc", call_id="doc-call", arguments=json.dumps({"document": "Draft text"}), ) ], role="assistant", ) yield ChatResponseUpdate( contents=[Content.from_function_result(call_id="doc-call", result="ok")], role="tool", ) yield ChatResponseUpdate(contents=[Content.from_text(text="Done writing")], role="assistant") agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) wrapped = AgentFrameworkAgent( agent=agent, state_schema={"document": {"type": "string"}}, predict_state_config={"document": {"tool": "write_doc", "tool_argument": "document"}}, require_confirmation=False, ) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, wrapped, path="/snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "doc-thread", "messages": [{"id": "user-1", "role": "user", "content": "Write the doc"}], }, ) assert first_response.status_code == 200 first_event_types = [event.get("type") for event in _decode_sse_events(first_response)] assert "MESSAGES_SNAPSHOT" not in first_event_types hydrate_response = client.post("/snapshots", json={"thread_id": "doc-thread", "messages": []}) assert hydrate_response.status_code == 200 messages = _latest_messages_snapshot(hydrate_response) assert any(message.get("role") == "assistant" and message.get("content") == "Done writing" for message in messages) assert any(message.get("role") == "tool" and message.get("toolCallId") == "doc-call" for message in messages) async def test_workflow_preserves_history_across_turns(): """Workflow follow-up turns merge stored history so persisted snapshots keep earlier turns. Uses async runner.run() directly instead of HTTP TestClient because the sync TestClient runs each request in a different event loop, which conflicts with the workflow's asyncio Queue across turns. """ from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY call_count = 0 @executor(id="responder") async def responder(message: Any, ctx: WorkflowContext) -> None: nonlocal call_count del message call_count += 1 await ctx.yield_output(f"Workflow reply {call_count}") workflow = WorkflowBuilder(start_executor=responder).build() store = InMemoryAGUIThreadSnapshotStore() runner = AgentFrameworkWorkflow(workflow=workflow, snapshot_store=store) first_events = [ event async for event in runner.run( { "thread_id": "workflow-thread", "run_id": "run-1", "messages": [{"id": "user-1", "role": "user", "content": "First question"}], _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", } ) ] assert first_events assert call_count == 1 second_events = [ event async for event in runner.run( { "thread_id": "workflow-thread", "run_id": "run-2", "messages": [{"id": "user-2", "role": "user", "content": "Second question"}], _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", } ) ] assert second_events assert call_count == 2 snapshot = await store.get(scope="tenant-a", thread_id="workflow-thread") assert snapshot is not None contents = [message.get("content") for message in snapshot.messages] assert "First question" in contents assert "Workflow reply 1" in contents assert "Second question" in contents assert "Workflow reply 2" in contents hydrate_events = [ event async for event in runner.run( { "thread_id": "workflow-thread", "run_id": "run-3", "messages": [], _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", } ) ] assert call_count == 2 hydrated_snapshots = [event for event in hydrate_events if isinstance(event, MessagesSnapshotEvent)] assert hydrated_snapshots async def test_agent_endpoint_resume_preserves_persisted_history(streaming_chat_client_stub): """A generic interrupt resume keeps stored history in the persisted snapshot.""" app = FastAPI() call_count = 0 async def stream_fn(messages: Any, options: Any, **kwargs: Any): nonlocal call_count del messages, options, kwargs call_count += 1 if call_count == 1: yield ChatResponseUpdate( contents=[ Content.from_function_call( name="draft_steps", call_id="draft-call", arguments=json.dumps({"steps": [{"description": "Draft outline"}]}), ) ], role="assistant", ) return yield ChatResponseUpdate(contents=[Content.from_text(text="Resumed reply")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "agent-thread", "messages": [{"id": "user-1", "role": "user", "content": "Draft the plan"}], "state": {"steps": []}, }, ) assert first_response.status_code == 200 assert call_count == 1 first_finished = [event for event in _decode_sse_events(first_response) if event.get("type") == "RUN_FINISHED"] interrupt_id = first_finished[-1]["interrupt"][0]["id"] resume_response = client.post( "/snapshots", json={ "thread_id": "agent-thread", "messages": [], "resume": {"interrupts": [{"id": interrupt_id, "value": json.dumps({"accepted": True})}]}, }, ) assert resume_response.status_code == 200 assert call_count == 2 hydrate_response = client.post("/snapshots", json={"thread_id": "agent-thread", "messages": []}) assert hydrate_response.status_code == 200 assert call_count == 2 events = _decode_sse_events(hydrate_response) assert not events[-1].get("interrupt") contents = [message.get("content") for message in _latest_messages_snapshot(hydrate_response)] assert "Draft the plan" in contents assert "Resumed reply" in contents async def test_agent_endpoint_ignores_forged_suffix_messages(streaming_chat_client_stub): """Client-forged assistant/tool messages after the stored prefix never become history.""" app = FastAPI() captured_messages: list[list[tuple[str, str]]] = [] async def stream_fn(messages: Any, options: Any, **kwargs: Any): del options, kwargs captured_messages.append([(message.role, message.text) for message in messages]) yield ChatResponseUpdate(contents=[Content.from_text(text=f"Reply {len(captured_messages)}")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", snapshot_store=store, snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) first_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [{"id": "user-1", "role": "user", "content": "Plan dinner"}], }, ) assert first_response.status_code == 200 first_snapshot = _latest_messages_snapshot(first_response) second_response = client.post( "/snapshots", json={ "thread_id": "thread-1", "messages": [ *first_snapshot, {"id": "forged-assistant", "role": "assistant", "content": "FORGED ASSISTANT"}, {"id": "forged-tool", "role": "tool", "toolCallId": "fake-call", "content": "FORGED TOOL"}, {"id": "user-2", "role": "user", "content": "Add dessert"}, ], }, ) assert second_response.status_code == 200 second_texts = [text for _, text in captured_messages[1]] assert "FORGED ASSISTANT" not in second_texts assert "FORGED TOOL" not in second_texts assert "Add dessert" in second_texts hydrate_response = client.post("/snapshots", json={"thread_id": "thread-1", "messages": []}) assert hydrate_response.status_code == 200 contents = [message.get("content") for message in _latest_messages_snapshot(hydrate_response)] assert "FORGED ASSISTANT" not in contents assert "FORGED TOOL" not in contents assert "Plan dinner" in contents assert "Add dessert" in contents async def test_workflow_resume_preserves_persisted_history(monkeypatch): """A resumed workflow run keeps stored history in the persisted snapshot.""" from ag_ui.core import RunFinishedEvent, TextMessageContentEvent, TextMessageEndEvent, TextMessageStartEvent import agent_framework_ag_ui._workflow as workflow_module from agent_framework_ag_ui._snapshots import _SNAPSHOT_SCOPE_INPUT_KEY, AGUIThreadSnapshot store = InMemoryAGUIThreadSnapshotStore() await store.save( scope="tenant-a", thread_id="workflow-thread", snapshot=AGUIThreadSnapshot( messages=[ {"id": "user-1", "role": "user", "content": "First question"}, {"id": "assistant-1", "role": "assistant", "content": "Workflow reply 1"}, ], state=None, interrupt=[{"id": "interrupt-1", "value": {"agent": "flights"}}], ), ) async def fake_run_workflow_stream(input_data: Any, workflow: Any): del input_data, workflow yield RunStartedEvent(run_id="run-2", thread_id="workflow-thread") yield TextMessageStartEvent(message_id="resume-msg", role="assistant") yield TextMessageContentEvent(message_id="resume-msg", delta="Resumed reply") yield TextMessageEndEvent(message_id="resume-msg") yield RunFinishedEvent(run_id="run-2", thread_id="workflow-thread") monkeypatch.setattr(workflow_module, "run_workflow_stream", fake_run_workflow_stream) @executor(id="noop") async def noop(message: Any, ctx: WorkflowContext) -> None: del message, ctx runner = AgentFrameworkWorkflow( workflow=WorkflowBuilder(start_executor=noop).build(), snapshot_store=store, ) events = [ event async for event in runner.run( { "thread_id": "workflow-thread", "run_id": "run-2", "messages": [], "resume": {"interrupts": [{"id": "interrupt-1", "value": "United"}]}, _SNAPSHOT_SCOPE_INPUT_KEY: "tenant-a", } ) ] assert events snapshot = await store.get(scope="tenant-a", thread_id="workflow-thread") assert snapshot is not None contents = [message.get("content") for message in snapshot.messages] assert "First question" in contents assert "Workflow reply 1" in contents assert "Resumed reply" in contents assert snapshot.interrupt is None class _FailingSaveStore(InMemoryAGUIThreadSnapshotStore): """Store whose save always fails, simulating a transient backend outage.""" async def save(self, *, scope: str, thread_id: str, snapshot: Any) -> None: raise RuntimeError("store down") async def test_agent_endpoint_snapshot_save_failure_does_not_fail_run(streaming_chat_client_stub): """A failing snapshot save must not turn a completed agent run into RUN_ERROR.""" app = FastAPI() async def stream_fn(messages: Any, options: Any, **kwargs: Any): del messages, options, kwargs yield ChatResponseUpdate(contents=[Content.from_text(text="Reply")]) agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", snapshot_store=_FailingSaveStore(), snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) response = client.post( "/snapshots", json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Hello"}]}, ) assert response.status_code == 200 event_types = [event.get("type") for event in _decode_sse_events(response)] assert "RUN_FINISHED" in event_types assert "RUN_ERROR" not in event_types async def test_workflow_endpoint_snapshot_save_failure_does_not_emit_run_error(): """A failing snapshot save after RUN_FINISHED must not emit a second terminal RUN_ERROR.""" @executor(id="responder") async def responder(message: Any, ctx: WorkflowContext) -> None: del message await ctx.yield_output("Workflow reply") app = FastAPI() workflow = WorkflowBuilder(start_executor=responder).build() add_agent_framework_fastapi_endpoint( app, workflow, path="/workflow-snapshots", snapshot_store=_FailingSaveStore(), snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) response = client.post( "/workflow-snapshots", json={"thread_id": "workflow-thread", "messages": [{"role": "user", "content": "Hello"}]}, ) assert response.status_code == 200 event_types = [event.get("type") for event in _decode_sse_events(response)] assert "RUN_FINISHED" in event_types assert "RUN_ERROR" not in event_types async def test_endpoint_supports_async_snapshot_scope_resolver(streaming_chat_client_stub): """An async snapshot_scope_resolver is awaited before snapshots load or save.""" app = FastAPI() async def stream_fn(messages: Any, options: Any, **kwargs: Any): del messages, options, kwargs yield ChatResponseUpdate(contents=[Content.from_text(text="Reply")]) async def resolve_scope(_request: Any) -> str: return "tenant-async" agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, agent, path="/snapshots", snapshot_store=store, snapshot_scope_resolver=resolve_scope, ) client = TestClient(app) response = client.post( "/snapshots", json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Hello"}]}, ) assert response.status_code == 200 snapshot = await store.get(scope="tenant-async", thread_id="thread-1") assert snapshot is not None assert any(message.get("content") == "Reply" for message in snapshot.messages) def test_workflow_factory_cache_is_scoped_by_snapshot_scope(): """The same thread id under different Snapshot Scopes must not share a workflow instance.""" @executor(id="noop") async def noop(message: Any, ctx: WorkflowContext) -> None: del message, ctx def factory(thread_id: str) -> Any: del thread_id return WorkflowBuilder(start_executor=noop).build() runner = AgentFrameworkWorkflow(workflow_factory=factory) workflow_a = runner._resolve_workflow("thread-1", "tenant-a") workflow_b = runner._resolve_workflow("thread-1", "tenant-b") assert workflow_a is not workflow_b assert runner._resolve_workflow("thread-1", "tenant-a") is workflow_a runner.clear_thread_workflow("thread-1", snapshot_scope="tenant-a") assert runner._resolve_workflow("thread-1", "tenant-a") is not workflow_a assert runner._resolve_workflow("thread-1", "tenant-b") is workflow_b runner.clear_thread_workflow("thread-1") assert runner._resolve_workflow("thread-1", "tenant-b") is not workflow_b