From 916291eaf45a23a578edaadbc9888a022f69ad7d Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 13 Jul 2026 13:38:05 +0900 Subject: [PATCH 1/5] Python: bridge AG-UI request state into sessions Decisions: - Project resolved AG-UI Shared State into the per-run AgentSession without typed restoration. - Preserve existing local/service session identifiers and keep AG-UI state out of provider metadata. Files changed: - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/tests/ag_ui/test_endpoint.py Verification: - uv run poe check -P ag-ui - uv run poe test -P ag-ui (912 passed) Notes: - Scoped cross-run Session Continuation State remains for the next dependent issue. --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 1 + .../ag-ui/tests/ag_ui/test_endpoint.py | 79 +++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index a2ccbc1290b..27b2c8e1581 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -1821,6 +1821,7 @@ async def run_agent_stream( session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id) else: session = AgentSession(session_id=thread_id) + session.state.update(flow.current_state) _restore_tool_approval_state(session, approval_state_store, approval_thread_id) # Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation) 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 3edfde77d48..b49b9bad4d8 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -16,11 +16,15 @@ from agent_framework import ( Agent, AgentResponseUpdate, + AgentSession, ChatResponseUpdate, Content, + ContextProvider, Executor, FunctionTool, Message, + SessionContext, + SupportsAgentRun, ToolApprovalMiddleware, WorkflowBuilder, WorkflowContext, @@ -334,6 +338,81 @@ async def test_endpoint_with_state_schema(build_chat_client): assert response.status_code == 200 +async def test_endpoint_bridges_request_state_to_context_provider_without_snapshots(streaming_chat_client_stub): + """Request Shared State is ordinary per-run context available to context providers.""" + + class RequestContextProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, state + identity = session.state["identity"] + assert isinstance(identity, dict) + context.extend_messages( + self, + [ + Message( + role="system", + contents=[ + Content.from_text( + text=( + f"agent={session.state['agent_id']} " + f"user={session.state['user_id']} " + f"tenant={session.state['tenant_id']} " + f"type={identity['type']}" + ) + ) + ], + ) + ], + ) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del kwargs + assert options.get("metadata") is None + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + expected = "agent=agent-123 user=user-456 tenant=tenant-789 type=message" + provider = RequestContextProvider("request-context") + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[provider], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, agent, path="/request-context", keepalive_seconds=None) + + response = TestClient(app).post( + "/request-context", + json={ + "messages": [{"role": "user", "content": "Who am I?"}], + "state": { + "agent_id": "agent-123", + "user_id": "user-456", + "tenant_id": "tenant-789", + "identity": {"type": "message"}, + }, + }, + ) + + assert response.status_code == 200 + text_deltas = [ + event["delta"] for event in _decode_sse_events(response) if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert text_deltas == [expected] + + async def test_endpoint_with_default_state_seed(build_chat_client): """Test endpoint seeds default state when client omits it.""" app = FastAPI() From e3068626bcb1423402ac346bc84883c83a951e42 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 13 Jul 2026 13:47:06 +0900 Subject: [PATCH 2/5] Python: persist scoped AG-UI session continuity Decisions: - Store private Session Continuation State atomically in scoped thread snapshots and restore it through the core AgentSession contract. - Exclude Shared State keys, all HistoryProvider buckets, and tool approval state; request overlays evict colliding private values. - Finalize interrupted response streams before snapshotting so provider after_run mutations are included. Files changed: - packages/ag-ui/agent_framework_ag_ui/_agent_run.py - packages/ag-ui/agent_framework_ag_ui/_snapshots.py - packages/ag-ui/tests/ag_ui/test_endpoint.py - packages/ag-ui/tests/ag_ui/test_snapshots.py - packages/ag-ui/AGENTS.md Verification: - uv run poe test -P ag-ui (921 passed) - uv run poe check -P ag-ui - uv run poe typing -P ag-ui Notes: - Lifecycle, isolation, and broader storage guidance remain for the next dependent issue. --- python/packages/ag-ui/AGENTS.md | 3 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 62 +- .../ag-ui/agent_framework_ag_ui/_snapshots.py | 11 +- .../ag-ui/tests/ag_ui/test_endpoint.py | 577 ++++++++++++++++++ .../ag-ui/tests/ag_ui/test_snapshots.py | 26 +- 5 files changed, 669 insertions(+), 10 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 3e31b8d46eb..4c58c746d2d 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -15,7 +15,8 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. ## Types - **`AGUIRequest`** / **`AGUIChatOptions`** - Request types -- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Replayable thread snapshot model and scoped async store protocol +- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Thread snapshot model with client-replayable data, + private Session Continuation State, and a scoped async store protocol - **`availableInterrupts` / `resume`** - Optional canonical AG-UI `Interrupt` and `ResumeEntry` protocol data - **`AgentState`** / **`RunMetadata`** - State management types - **`PredictStateConfig`** - Configuration for state prediction diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 27b2c8e1581..931285e8620 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -30,6 +30,7 @@ from agent_framework import ( AgentSession, Content, + HistoryProvider, Message, SupportsAgentRun, ) @@ -581,6 +582,7 @@ def _restore_tool_approval_state( thread_id: str, ) -> None: """Restore only core tool-approval state into the per-run AgentSession.""" + session.state.pop(_TOOL_APPROVAL_STATE_KEY, None) if approval_state_store is None: return stored_state = approval_state_store.tool_approval_states.get(thread_id) @@ -1624,8 +1626,9 @@ async def _save_thread_snapshot( messages: list[dict[str, Any]], state: dict[str, Any] | None, interrupt: list[dict[str, Any]] | None, + session_state: dict[str, Any] | None, ) -> None: - """Save the latest replayable AG-UI Thread Snapshot when persistence is configured.""" + """Save the latest AG-UI Thread Snapshot when persistence is configured.""" if config.snapshot_store is None or scope is None: return @@ -1633,7 +1636,12 @@ async def _save_thread_snapshot( await config.snapshot_store.save( scope=scope, thread_id=thread_id, - snapshot=AGUIThreadSnapshot(messages=messages, state=state, interrupt=interrupt), + snapshot=AGUIThreadSnapshot( + messages=messages, + state=state, + interrupt=interrupt, + session_state=session_state, + ), ) except Exception: # The run itself already streamed successfully; a transient store failure @@ -1646,6 +1654,42 @@ async def _save_thread_snapshot( ) +def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThreadSnapshot | None) -> None: + """Restore typed private state from trusted snapshot storage.""" + if snapshot is None or snapshot.session_state is None: + return + restored = AgentSession.from_dict( + { + "type": "session", + "session_id": session.session_id, + "state": snapshot.session_state, + } + ) + session.state.update(restored.state) + + +def _serialize_session_continuation_state( + session: AgentSession, + agent: SupportsAgentRun, + *, + shared_state_keys: set[str], +) -> dict[str, Any] | None: + """Serialize server-owned state while preserving each AG-UI State Authority.""" + context_providers = cast(list[Any], getattr(agent, "context_providers", [])) + excluded_keys = { + *shared_state_keys, + _TOOL_APPROVAL_STATE_KEY, + *(provider.source_id for provider in context_providers if isinstance(provider, HistoryProvider)), + } + continuation_state = {key: value for key, value in session.state.items() if key not in excluded_keys} + if not continuation_state: + return None + + serialized_session = AgentSession(session_id=session.session_id) + serialized_session.state.update(continuation_state) + return cast(dict[str, Any], serialized_session.to_dict()["state"]) + + async def run_agent_stream( input_data: dict[str, Any], agent: SupportsAgentRun, @@ -1821,6 +1865,7 @@ async def run_agent_stream( session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id) else: session = AgentSession(session_id=thread_id) + _restore_session_continuation_state(session, stored_snapshot) session.state.update(flow.current_state) _restore_tool_approval_state(session, approval_state_store, approval_thread_id) @@ -1895,6 +1940,11 @@ async def run_agent_stream( messages=persisted_messages, state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None, interrupt=None, + session_state=_serialize_session_continuation_state( + session, + agent, + shared_state_keys=set(flow.current_state), + ), ) _save_tool_approval_state(session, approval_state_store, approval_thread_id) yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) @@ -2001,6 +2051,9 @@ async def run_agent_stream( if flow.waiting_for_approval: break + if flow.waiting_for_approval and isinstance(stream, ResponseStream): + await stream.get_final_response() + # If no updates at all, still emit RunStarted if not run_started_emitted: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) @@ -2189,6 +2242,11 @@ async def run_agent_stream( messages=persisted_messages, state=latest_state_snapshot, interrupt=flow.interrupts or None, + session_state=_serialize_session_continuation_state( + session, + agent, + shared_state_keys=set(flow.current_state), + ), ) _save_tool_approval_state(session, approval_state_store, approval_thread_id) yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=flow.interrupts) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py index 4993b1aa2ae..e862d3b52c7 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py @@ -32,22 +32,25 @@ @dataclass(slots=True) class AGUIThreadSnapshot: - """Replayable AG-UI Thread state. + """AG-UI Thread state. - AG-UI Thread Snapshots intentionally contain only data that can be replayed - to a UI: message snapshots, optional Shared State, and optional interruption - state. They do not include raw events, request metadata, auth claims, + AG-UI Thread Snapshots contain client-replayable message, Shared State, and + interruption data plus optional private Session Continuation State. Private + continuation is trusted server data and must never be replayed to the client. + Snapshots do not include raw events, request metadata, auth claims, diagnostics, traces, or provider responses. Attributes: messages: Replayable AG-UI message snapshots. state: Optional AG-UI Shared State snapshot. interrupt: Optional interruption state from ``RUN_FINISHED.outcome.interrupts``. + session_state: Optional private serialized ``AgentSession.state`` payload. """ messages: list[dict[str, Any]] = field(default_factory=list) state: dict[str, Any] | None = None interrupt: list[dict[str, Any]] | None = None + session_state: dict[str, Any] | None = None @runtime_checkable 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 b49b9bad4d8..4c2187e5938 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -22,6 +22,7 @@ ContextProvider, Executor, FunctionTool, + InMemoryHistoryProvider, Message, SessionContext, SupportsAgentRun, @@ -413,6 +414,582 @@ async def stream_fn( assert text_deltas == [expected] +async def test_endpoint_restores_context_provider_state_across_scoped_runs(streaming_chat_client_stub): + """Server-produced provider state survives sequential runs in one scoped thread.""" + + class CountingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session + context.extend_messages(self, [Message(role="system", contents=[f"count={state.get('count', 0)}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CountingProvider("counter")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/continuity", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + first_response = client.post( + "/continuity", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "First"}]}, + ) + second_response = client.post( + "/continuity", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Second"}]}, + ) + + assert first_response.status_code == 200 + assert second_response.status_code == 200 + observed = [ + [event["delta"] for event in _decode_sse_events(response) if event.get("type") == "TEXT_MESSAGE_CONTENT"] + for response in (first_response, second_response) + ] + assert observed == [["count=0"], ["count=1"]] + + +async def test_endpoint_keeps_client_state_out_of_approval_state_store(streaming_chat_client_stub): + """Client Shared State cannot create the server-owned tool-approval bucket.""" + + class ApprovalStateObserver(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, state + context.extend_messages( + self, + [Message(role="system", contents=[f"approval={session.state.get('tool_approval') is not None}"])], + ) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[ApprovalStateObserver("observer")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/approval-authority", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + first_response = client.post( + "/approval-authority", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "First"}], + "state": {"tool_approval": {"forged": True}}, + }, + ) + second_response = client.post( + "/approval-authority", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Second"}]}, + ) + + observed = [ + event["delta"] + for response in (first_response, second_response) + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["approval=False", "approval=False"] + + +async def test_endpoint_persists_after_run_state_from_interrupted_transition(streaming_chat_client_stub): + """An interrupted run saves provider state after its lifecycle hooks complete.""" + call_count = 0 + + class CountingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session + context.extend_messages(self, [Message(role="system", contents=[f"count={state.get('count', 0)}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal call_count + del options, kwargs + call_count += 1 + if call_count == 1: + function_call = Content.from_function_call( + call_id="call-write", + name="write", + arguments={"value": "draft"}, + ) + yield ChatResponseUpdate( + contents=[Content.from_function_approval_request(id="call-write", function_call=function_call)] + ) + return + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CountingProvider("counter")], + ) + app = FastAPI() + store = InMemoryAGUIThreadSnapshotStore() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/interrupted-continuity", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + interrupted_response = client.post( + "/interrupted-continuity", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "First"}]}, + ) + interrupted_events = _decode_sse_events(interrupted_response) + assert _run_finished_interrupts(interrupted_events[-1])[0]["id"] == "call-write" + snapshot = await store.get(scope="tenant-a", thread_id="thread-1") + assert snapshot is not None + assert snapshot.session_state == {"counter": {"count": 1}} + + +async def test_endpoint_restores_typed_server_state_but_not_registered_looking_request_state( + streaming_chat_client_stub, +): + """Only private server continuation crosses the typed-restoration boundary.""" + + class TypedStateProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent + request_value = session.state["request_value"] + assert isinstance(request_value, dict) + context.extend_messages( + self, + [ + Message( + role="system", + contents=[ + f"server={type(state.get('server_value')).__name__} request={type(request_value).__name__}" + ], + ) + ], + ) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state.setdefault("server_value", Message(role="assistant", contents=["private"])) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[TypedStateProvider("typed")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/typed-continuity", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + request_state = { + "request_value": { + "type": "message", + "role": "assistant", + "contents": [{"type": "text", "text": "untrusted"}], + } + } + + first_response = client.post( + "/typed-continuity", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "First"}], + "state": request_state, + }, + ) + second_response = client.post( + "/typed-continuity", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Second"}]}, + ) + + observed = [ + event["delta"] + for response in (first_response, second_response) + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["server=NoneType request=dict", "server=Message request=dict"] + + +async def test_endpoint_request_collision_evicts_prior_private_value(streaming_chat_client_stub): + """A request overlay replaces and evicts a colliding private continuation value.""" + run_count = 0 + + class CollisionProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, state + context.extend_messages(self, [Message(role="system", contents=[f"mode={session.state.get('mode')}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + nonlocal run_count + del agent, context, state + run_count += 1 + if run_count == 1: + session.state["mode"] = "server-private" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CollisionProvider("observer")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/collision", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + responses = [ + client.post( + "/collision", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "First"}]}, + ), + client.post( + "/collision", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "Second"}], + "state": {"mode": "request"}, + }, + ), + client.post( + "/collision", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Third"}]}, + ), + ] + + observed = [ + event["delta"] + for response in responses + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["mode=None", "mode=request", "mode=request"] + + +async def test_endpoint_excludes_history_provider_state_from_continuation(streaming_chat_client_stub): + """Snapshot messages remain the sole conversation-history authority.""" + captured_messages: list[list[tuple[str, str]]] = [] + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + 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=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[InMemoryHistoryProvider(source_id="history")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/history-authority", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + first_response = client.post( + "/history-authority", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "First"}]}, + ) + second_response = client.post( + "/history-authority", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Second"}]}, + ) + + assert first_response.status_code == 200 + assert second_response.status_code == 200 + assert captured_messages == [ + [("user", "First")], + [("user", "First"), ("assistant", "Reply 1"), ("user", "Second")], + ] + + +async def test_endpoint_session_continuity_requires_scoped_snapshot_configuration(streaming_chat_client_stub): + """Server-produced state remains per-run when scoped snapshots are not configured.""" + + class CountingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session + context.extend_messages(self, [Message(role="system", contents=[f"count={state.get('count', 0)}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CountingProvider("counter")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, agent, path="/stateless", keepalive_seconds=None) + client = TestClient(app) + + responses = [ + client.post( + "/stateless", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": prompt}]}, + ) + for prompt in ("First", "Second") + ] + + observed = [ + event["delta"] + for response in responses + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["count=0", "count=0"] + + +async def test_endpoint_failed_run_keeps_previous_completed_continuation(streaming_chat_client_stub): + """A failed run cannot replace the last completed private continuation.""" + call_count = 0 + + class CountingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session + context.extend_messages(self, [Message(role="system", contents=[f"count={state.get('count', 0)}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal call_count + del options, kwargs + call_count += 1 + if call_count == 2: + raise RuntimeError("model failed") + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CountingProvider("counter")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/failed-continuity", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + responses = [ + client.post( + "/failed-continuity", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": prompt}]}, + ) + for prompt in ("First", "Second", "Third") + ] + + assert any(event.get("type") == "RUN_ERROR" for event in _decode_sse_events(responses[1])) + observed = [ + event["delta"] + for response in (responses[0], responses[2]) + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["count=0", "count=1"] + + async def test_endpoint_with_default_state_seed(build_chat_client): """Test endpoint seeds default state when client omits it.""" app = FastAPI() diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py index 00dcd211da1..e998d2f29e7 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -7,9 +7,9 @@ from agent_framework_ag_ui import AGUIThreadSnapshot, AGUIThreadSnapshotStore, InMemoryAGUIThreadSnapshotStore -def test_thread_snapshot_model_contains_only_replayable_snapshot_fields() -> None: - """The public snapshot model is limited to messages, Shared State, and interruption state.""" - assert [field.name for field in fields(AGUIThreadSnapshot)] == ["messages", "state", "interrupt"] +def test_thread_snapshot_model_contains_replayable_and_private_snapshot_fields() -> None: + """The public snapshot model carries replayable data and optional private continuation.""" + assert [field.name for field in fields(AGUIThreadSnapshot)] == ["messages", "state", "interrupt", "session_state"] def test_in_memory_snapshot_store_satisfies_snapshot_store_protocol() -> None: @@ -39,6 +39,26 @@ async def test_in_memory_snapshot_store_replaces_latest_snapshot() -> None: assert snapshot.state == {"count": 2} +async def test_in_memory_snapshot_store_defensively_copies_private_continuation() -> None: + """Private continuation cannot be mutated through saved or returned references.""" + store = InMemoryAGUIThreadSnapshotStore() + session_state = {"provider": {"count": 1}} + snapshot = AGUIThreadSnapshot(session_state=session_state) + + await store.save(scope="tenant-a", thread_id="thread-1", snapshot=snapshot) + session_state["provider"]["count"] = 2 + stored = await store.get(scope="tenant-a", thread_id="thread-1") + + assert stored is not None + assert stored.session_state is not None + assert stored.session_state == {"provider": {"count": 1}} + stored.session_state["provider"]["count"] = 3 + + reread = await store.get(scope="tenant-a", thread_id="thread-1") + assert reread is not None + assert reread.session_state == {"provider": {"count": 1}} + + async def test_in_memory_snapshot_store_keeps_scopes_separate() -> None: """The same AG-UI Thread id in different Snapshot Scopes addresses different snapshots.""" store = InMemoryAGUIThreadSnapshotStore() From ee5ebba4ae4cba835d1b5004d691c8463988938d Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 13 Jul 2026 13:53:11 +0900 Subject: [PATCH 3/5] Python: document AG-UI session continuity lifecycle Decisions: - Keep scoped thread snapshots as the single reset and continuity boundary, with missing request Shared State preserving private continuation. - Document trusted typed-restoration storage, State Authorities, custom-store round trips, and one-active-run last-writer-wins consistency. - Verify failure, hydration privacy, scope/thread isolation, and reset mechanics through public endpoint and store seams. Files changed: - packages/ag-ui/README.md - packages/ag-ui/tests/ag_ui/test_endpoint.py - packages/ag-ui/tests/ag_ui/test_snapshots.py Verification: - uv run pytest -q (6 passed) - uv run poe test -P ag-ui - uv run poe syntax -P ag-ui -C - uv run poe typing -P ag-ui - uv run poe check -P ag-ui - uv run poe markdown-code-lint Notes: - No runtime capability probe, secondary state store, locking, or configuration flag was added. - No blockers remain for this lifecycle and guidance slice. --- python/packages/ag-ui/README.md | 35 +++- .../ag-ui/tests/ag_ui/test_endpoint.py | 192 +++++++++++++++++- .../ag-ui/tests/ag_ui/test_snapshots.py | 40 +++- 3 files changed, 254 insertions(+), 13 deletions(-) diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index b34f3545acf..0d7247f20b2 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -337,6 +337,24 @@ Endpoint configuration requires `snapshot_scope_resolver` whenever a snapshot st the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFrameworkWorkflow`. The resolver returns the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key. +For hosted agents, request Shared State is also available through `AgentSession.state` during that run, whether or +not snapshot persistence is configured. Request values are untrusted per-run context: they overlay restored values, +are not passed through typed session restoration, and are excluded from private Session Continuation State. + +When scoped snapshots are configured, each category has one State Authority: + +| State category | State Authority | +| --- | --- | +| Conversation history | AG-UI Thread Snapshot messages | +| AG-UI Shared State and request context | The current AG-UI request and replayable snapshot state | +| Approval State | The Approval State Store | +| Other server-produced provider working state | Private Session Continuation State | + +Session Continuation State is stored atomically in the optional `AGUIThreadSnapshot.session_state` field and restored +through the core `AgentSession` typed serialization contract. It is never accepted from an AG-UI request or emitted +during hydration. Deleting a scoped thread snapshot resets its replayable and private state together, and clearing a +Snapshot Scope removes all such records in that scope. Missing or empty request Shared State is not a reset command. + AG-UI Thread ids identify AG-UI Threads; they do not authorize snapshot access. Do not treat a thread id as a bearer credential or tenant boundary. Production applications must authenticate and authorize every AG-UI endpoint request and choose a Snapshot Scope that represents the app's real access boundary, such as an authenticated user, tenant, @@ -348,14 +366,21 @@ It is not an authentication, tenant authorization, or distributed durability mec responsible for endpoint authentication, tenant authorization, and deployment/storage architecture that matches their availability and worker topology requirements. -Stored snapshots are untrusted application data with confidentiality impact. They may contain sensitive user text, -model output, tool results, function arguments, UI payloads, Shared State, and interrupt data. The built-in -`InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable production -storage. It is cleared on process restart and is not shared across workers. +Snapshot storage is trusted, integrity-protected server storage because private continuation is eligible for typed +core restoration. Snapshots also have confidentiality impact: they may contain sensitive user text, model output, +tool results, function arguments, UI payloads, Shared State, interrupt data, and private provider working state. The +built-in `InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable +production storage. It is cleared on process restart and is not shared across workers. No file-backed AG-UI snapshot store is provided by the package. Applications that need durable persistence should provide an app-owned implementation of the `AGUIThreadSnapshotStore` protocol and own storage hardening, including -encryption, access control, retention, audit, data residency, and deletion behavior. +encryption, integrity protection, access control, retention, audit, data residency, and deletion behavior. Existing +custom stores remain source-compatible because `session_state` is optional, but they provide Session State Continuity +only when they round-trip that field unchanged with the rest of the snapshot. + +The supported consistency model is one active run per `(Snapshot Scope, threadId)`. Concurrent writes to the same +scoped thread remain last-writer-wins. Applications that require stronger consistency must serialize those runs using +coordination appropriate to their deployment; a process-local lock does not provide distributed consistency. ## Architecture 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 4c2187e5938..03a6831b2c6 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -484,6 +484,85 @@ async def stream_fn( assert observed == [["count=0"], ["count=1"]] +async def test_endpoint_continuation_is_scoped_without_request_state_reset(streaming_chat_client_stub): + """Missing or empty request state preserves continuation without crossing scope or thread boundaries.""" + + class CountingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session + context.extend_messages(self, [Message(role="system", contents=[f"count={state.get('count', 0)}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CountingProvider("counter")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/isolated-continuity", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda request: cast("dict[str, Any]", request.forwarded_props)["scope"], + keepalive_seconds=None, + ) + client = TestClient(app) + + request_keys: list[tuple[str, str, dict[str, Any] | None]] = [ + ("tenant-a", "thread-1", None), + ("tenant-a", "thread-1", {}), + ("tenant-b", "thread-1", None), + ("tenant-a", "thread-2", None), + ("tenant-a", "thread-1", None), + ] + responses = [] + for index, (scope, thread_id, request_state) in enumerate(request_keys): + request: dict[str, Any] = { + "thread_id": thread_id, + "messages": [{"role": "user", "content": f"Turn {index}"}], + "forwardedProps": {"scope": scope}, + } + if request_state is not None: + request["state"] = request_state + responses.append(client.post("/isolated-continuity", json=request)) + + observed = [ + event["delta"] + for response in responses + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["count=0", "count=1", "count=0", "count=0", "count=2"] + + async def test_endpoint_keeps_client_state_out_of_approval_state_store(streaming_chat_client_stub): """Client Shared State cannot create the server-owned tool-approval bucket.""" @@ -3243,13 +3322,30 @@ async def test_agent_endpoint_hydrates_stored_thread_snapshot_without_invoking_a app = FastAPI() call_count = 0 + class PrivateStateProvider(ContextProvider): + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["private_secret"] = "must-not-be-replayed" + 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)) + agent = Agent( + name="test", + instructions="Test agent", + client=streaming_chat_client_stub(stream_fn), + context_providers=[PrivateStateProvider("private")], + ) store = InMemoryAGUIThreadSnapshotStore() add_agent_framework_fastapi_endpoint( app, @@ -3285,6 +3381,7 @@ async def stream_fn(messages: Any, options: Any, **kwargs: Any): message.get("role") == "assistant" and message.get("content") == "Stored reply" for message in events[2]["messages"] ) + assert b"must-not-be-replayed" not in hydrate_response.content async def test_agent_endpoint_hydrates_snapshots_by_scope_and_thread(streaming_chat_client_stub): @@ -4578,6 +4675,18 @@ async def save(self, *, scope: str, thread_id: str, snapshot: Any) -> None: raise RuntimeError("store down") +class _FailNextSaveStore(InMemoryAGUIThreadSnapshotStore): + """Store that can fail one save without replacing its previous snapshot.""" + + fail_next_save = False + + async def save(self, *, scope: str, thread_id: str, snapshot: Any) -> None: + if self.fail_next_save: + self.fail_next_save = False + raise RuntimeError("store down") + await super().save(scope=scope, thread_id=thread_id, snapshot=snapshot) + + 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() @@ -4607,6 +4716,87 @@ async def stream_fn(messages: Any, options: Any, **kwargs: Any): assert "RUN_ERROR" not in event_types +async def test_agent_endpoint_snapshot_save_failure_keeps_previous_continuation( + streaming_chat_client_stub, + caplog, +): + """A failed snapshot save is logged and leaves the previous completed continuation authoritative.""" + + class CountingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session + context.extend_messages(self, [Message(role="system", contents=[f"count={state.get('count', 0)}"])]) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + store = _FailNextSaveStore() + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[CountingProvider("counter")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/snapshot-failure", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + first_response = client.post( + "/snapshot-failure", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "First"}]}, + ) + store.fail_next_save = True + failed_save_response = client.post( + "/snapshot-failure", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Second"}]}, + ) + third_response = client.post( + "/snapshot-failure", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Third"}]}, + ) + + observed = [ + event["delta"] + for response in (first_response, failed_save_response, third_response) + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == ["count=0", "count=1", "count=1"] + assert "RUN_ERROR" not in [event.get("type") for event in _decode_sse_events(failed_save_response)] + assert "Failed to save AG-UI Thread Snapshot" in caplog.text + + 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.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py index e998d2f29e7..e2ea85e37df 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshots.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshots.py @@ -10,6 +10,7 @@ def test_thread_snapshot_model_contains_replayable_and_private_snapshot_fields() -> None: """The public snapshot model carries replayable data and optional private continuation.""" assert [field.name for field in fields(AGUIThreadSnapshot)] == ["messages", "state", "interrupt", "session_state"] + assert AGUIThreadSnapshot().session_state is None def test_in_memory_snapshot_store_satisfies_snapshot_store_protocol() -> None: @@ -24,12 +25,20 @@ async def test_in_memory_snapshot_store_replaces_latest_snapshot() -> None: await store.save( scope="tenant-a", thread_id="thread-1", - snapshot=AGUIThreadSnapshot(messages=[{"id": "first"}], state={"count": 1}), + snapshot=AGUIThreadSnapshot( + messages=[{"id": "first"}], + state={"count": 1}, + session_state={"provider": {"count": 1}}, + ), ) await store.save( scope="tenant-a", thread_id="thread-1", - snapshot=AGUIThreadSnapshot(messages=[{"id": "second"}], state={"count": 2}), + snapshot=AGUIThreadSnapshot( + messages=[{"id": "second"}], + state={"count": 2}, + session_state={"provider": {"count": 2}}, + ), ) snapshot = await store.get(scope="tenant-a", thread_id="thread-1") @@ -37,6 +46,7 @@ async def test_in_memory_snapshot_store_replaces_latest_snapshot() -> None: assert snapshot is not None assert snapshot.messages == [{"id": "second"}] assert snapshot.state == {"count": 2} + assert snapshot.session_state == {"provider": {"count": 2}} async def test_in_memory_snapshot_store_defensively_copies_private_continuation() -> None: @@ -87,19 +97,35 @@ async def test_in_memory_snapshot_store_deletes_and_clears_snapshots() -> None: """Delete removes one scoped thread key, while clear can remove a scope or the whole store.""" store = InMemoryAGUIThreadSnapshotStore() - await store.save(scope="tenant-a", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "a1"}])) - await store.save(scope="tenant-a", thread_id="thread-2", snapshot=AGUIThreadSnapshot(messages=[{"id": "a2"}])) - await store.save(scope="tenant-b", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "b1"}])) + await store.save( + scope="tenant-a", + thread_id="thread-1", + snapshot=AGUIThreadSnapshot(messages=[{"id": "a1"}], session_state={"private": "a1"}), + ) + await store.save( + scope="tenant-a", + thread_id="thread-2", + snapshot=AGUIThreadSnapshot(messages=[{"id": "a2"}], session_state={"private": "a2"}), + ) + await store.save( + scope="tenant-b", + thread_id="thread-1", + snapshot=AGUIThreadSnapshot(messages=[{"id": "b1"}], session_state={"private": "b1"}), + ) assert await store.delete(scope="tenant-a", thread_id="thread-1") is True assert await store.delete(scope="tenant-a", thread_id="thread-1") is False assert await store.get(scope="tenant-a", thread_id="thread-1") is None - assert await store.get(scope="tenant-a", thread_id="thread-2") is not None + tenant_a_thread_2 = await store.get(scope="tenant-a", thread_id="thread-2") + assert tenant_a_thread_2 is not None + assert tenant_a_thread_2.session_state == {"private": "a2"} await store.clear(scope="tenant-a") assert await store.get(scope="tenant-a", thread_id="thread-2") is None - assert await store.get(scope="tenant-b", thread_id="thread-1") is not None + tenant_b_thread_1 = await store.get(scope="tenant-b", thread_id="thread-1") + assert tenant_b_thread_1 is not None + assert tenant_b_thread_1.session_state == {"private": "b1"} await store.clear() From 36b6afd6e956a33ce2ea1e2a8fcbd0f17882fe99 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 13 Jul 2026 14:21:31 +0900 Subject: [PATCH 4/5] Python: harden AG-UI session continuity --- python/packages/ag-ui/README.md | 8 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 77 +++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 222 +++++++++++++++++- 3 files changed, 292 insertions(+), 15 deletions(-) diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 0d7247f20b2..62e55de7904 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -338,8 +338,10 @@ the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFramewo the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key. For hosted agents, request Shared State is also available through `AgentSession.state` during that run, whether or -not snapshot persistence is configured. Request values are untrusted per-run context: they overlay restored values, -are not passed through typed session restoration, and are excluded from private Session Continuation State. +not snapshot persistence is configured. Request values are untrusted per-run context: they overlay ordinary restored +values, are not passed through typed session restoration, and are excluded from private Session Continuation State. +Keys owned by configured context providers or reserved for approval and message-injection middleware are not copied +into `AgentSession.state`; their server-owned values take precedence over client Shared State. When scoped snapshots are configured, each category has one State Authority: @@ -354,6 +356,8 @@ Session Continuation State is stored atomically in the optional `AGUIThreadSnaps through the core `AgentSession` typed serialization contract. It is never accepted from an AG-UI request or emitted during hydration. Deleting a scoped thread snapshot resets its replayable and private state together, and clearing a Snapshot Scope removes all such records in that scope. Missing or empty request Shared State is not a reset command. +If private continuation cannot be restored or serialized, the endpoint logs the failure and continues without that +continuation so stale or unsupported provider state cannot permanently block the thread or suppress `RUN_FINISHED`. AG-UI Thread ids identify AG-UI Threads; they do not authorize snapshot access. Do not treat a thread id as a bearer credential or tenant boundary. Production applications must authenticate and authorize every AG-UI endpoint request diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 931285e8620..341e1edcbe5 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -31,6 +31,8 @@ AgentSession, Content, HistoryProvider, + InMemoryHistoryProvider, + MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, Message, SupportsAgentRun, ) @@ -1658,16 +1660,34 @@ def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThr """Restore typed private state from trusted snapshot storage.""" if snapshot is None or snapshot.session_state is None: return - restored = AgentSession.from_dict( - { - "type": "session", - "session_id": session.session_id, - "state": snapshot.session_state, - } - ) + try: + restored = AgentSession.from_dict( + { + "type": "session", + "session_id": session.session_id, + "state": snapshot.session_state, + } + ) + except Exception: + logger.exception( + "Failed to restore AG-UI Session Continuation State for session_id=%s; continuing without it.", + session.session_id, + ) + return session.state.update(restored.state) +def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]: + """Return session-state namespaces that client Shared State cannot own.""" + context_providers = cast(list[Any], getattr(agent, "context_providers", [])) + return { + _TOOL_APPROVAL_STATE_KEY, + InMemoryHistoryProvider.DEFAULT_SOURCE_ID, + MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, + *(provider.source_id for provider in context_providers), + } + + def _serialize_session_continuation_state( session: AgentSession, agent: SupportsAgentRun, @@ -1690,6 +1710,36 @@ def _serialize_session_continuation_state( return cast(dict[str, Any], serialized_session.to_dict()["state"]) +def _safe_serialize_session_continuation_state( + session: AgentSession, + agent: SupportsAgentRun, + *, + shared_state_keys: set[str], +) -> dict[str, Any] | None: + """Return JSON-safe continuation state without failing a completed run.""" + try: + serialized_state = _serialize_session_continuation_state( + session, + agent, + shared_state_keys=shared_state_keys, + ) + if serialized_state is None: + return None + safe_state = make_json_safe(serialized_state) + if isinstance(safe_state, dict): + return cast(dict[str, Any], safe_state) + logger.warning( + "Ignoring AG-UI Session Continuation State with unsupported serialized type: %s", + type(safe_state).__name__, + ) + except Exception: + logger.exception( + "Failed to serialize AG-UI Session Continuation State for session_id=%s; saving snapshot without it.", + session.session_id, + ) + return None + + async def run_agent_stream( input_data: dict[str, Any], agent: SupportsAgentRun, @@ -1866,7 +1916,10 @@ async def run_agent_stream( else: session = AgentSession(session_id=thread_id) _restore_session_continuation_state(session, stored_snapshot) - session.state.update(flow.current_state) + protected_session_state_keys = _request_state_protected_keys(agent) + session.state.update( + {key: value for key, value in flow.current_state.items() if key not in protected_session_state_keys} + ) _restore_tool_approval_state(session, approval_state_store, approval_thread_id) # Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation) @@ -1940,10 +1993,10 @@ async def run_agent_stream( messages=persisted_messages, state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None, interrupt=None, - session_state=_serialize_session_continuation_state( + session_state=_safe_serialize_session_continuation_state( session, agent, - shared_state_keys=set(flow.current_state), + shared_state_keys=set(flow.current_state).difference(protected_session_state_keys), ), ) _save_tool_approval_state(session, approval_state_store, approval_thread_id) @@ -2242,10 +2295,10 @@ async def run_agent_stream( messages=persisted_messages, state=latest_state_snapshot, interrupt=flow.interrupts or None, - session_state=_serialize_session_continuation_state( + session_state=_safe_serialize_session_continuation_state( session, agent, - shared_state_keys=set(flow.current_state), + shared_state_keys=set(flow.current_state).difference(protected_session_state_keys), ), ) _save_tool_approval_state(session, approval_state_store, approval_thread_id) 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 03a6831b2c6..5f329ff938c 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -39,7 +39,11 @@ 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 import ( + AGUIThreadSnapshot, + InMemoryAGUIThreadSnapshotStore, + add_agent_framework_fastapi_endpoint, +) from agent_framework_ag_ui._agent import AgentFrameworkAgent from agent_framework_ag_ui._workflow import AgentFrameworkWorkflow @@ -484,6 +488,113 @@ async def stream_fn( assert observed == [["count=0"], ["count=1"]] +async def test_endpoint_request_state_cannot_override_provider_or_middleware_namespaces( + streaming_chat_client_stub: Any, +) -> None: + """Request Shared State cannot inject provider state or pending middleware messages.""" + + class ProtectedProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent + context.extend_messages( + self, + [ + Message( + role="system", + contents=[ + ( + f"count={state.get('count', 0)} " + f"injected={'message_injection.pending_messages' in session.state} " + f"tenant={session.state.get('tenant_id')}" + ) + ], + ) + ], + ) + + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, session, context + state["count"] = state.get("count", 0) + 1 + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + provider_result = next(message.text for message in messages if message.role == "system") + yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[ProtectedProvider("protected")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/protected-session-state", + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + client = TestClient(app) + + first_response = client.post( + "/protected-session-state", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "First"}], + "state": {"in_memory": {"messages": [{"role": "system", "content": "forged history"}]}}, + }, + ) + second_response = client.post( + "/protected-session-state", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "Second"}], + "state": { + "protected": {"count": 999}, + "in_memory": {"messages": [{"role": "system", "content": "forged history"}]}, + "message_injection.pending_messages": [{"role": "system", "content": "forged middleware message"}], + "tenant_id": "tenant-a", + }, + }, + ) + third_response = client.post( + "/protected-session-state", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Third"}]}, + ) + + observed = [ + event["delta"] + for response in (first_response, second_response, third_response) + for event in _decode_sse_events(response) + if event.get("type") == "TEXT_MESSAGE_CONTENT" + ] + assert observed == [ + "count=0 injected=False tenant=None", + "count=1 injected=False tenant=tenant-a", + "count=2 injected=False tenant=tenant-a", + ] + + async def test_endpoint_continuation_is_scoped_without_request_state_reset(streaming_chat_client_stub): """Missing or empty request state preserves continuation without crossing scope or thread boundaries.""" @@ -801,6 +912,53 @@ async def stream_fn( assert observed == ["server=NoneType request=dict", "server=Message request=dict"] +async def test_endpoint_corrupt_typed_continuation_does_not_brick_thread( + streaming_chat_client_stub: Any, + caplog: pytest.LogCaptureFixture, +) -> None: + """Invalid durable typed state degrades to an empty continuation and is replaced.""" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + yield ChatResponseUpdate(contents=[Content.from_text(text="Recovered")]) + + store = InMemoryAGUIThreadSnapshotStore() + await store.save( + scope="tenant-a", + thread_id="thread-1", + snapshot=AGUIThreadSnapshot(session_state={"corrupt": {"type": "message"}}), + ) + agent = Agent(name="test", instructions=None, client=streaming_chat_client_stub(stream_fn)) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/corrupt-continuation", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + + response = TestClient(app).post( + "/corrupt-continuation", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Continue"}]}, + ) + events = _decode_sse_events(response) + stored = await store.get(scope="tenant-a", thread_id="thread-1") + + assert response.status_code == 200 + assert "RUN_FINISHED" in [event.get("type") for event in events] + assert "RUN_ERROR" not in [event.get("type") for event in events] + assert any(event.get("delta") == "Recovered" for event in events) + assert stored is not None + assert stored.session_state is None + assert "Failed to restore AG-UI Session Continuation State" in caplog.text + + async def test_endpoint_request_collision_evicts_prior_private_value(streaming_chat_client_stub): """A request overlay replaces and evicts a colliding private continuation value.""" run_count = 0 @@ -4797,6 +4955,68 @@ async def stream_fn( assert "Failed to save AG-UI Thread Snapshot" in caplog.text +async def test_endpoint_unsafe_continuation_serialization_does_not_fail_completed_run( + streaming_chat_client_stub: Any, + caplog: pytest.LogCaptureFixture, +) -> None: + """An unsupported provider value cannot suppress RUN_FINISHED or replayable snapshot data.""" + + class ExplodingState: + def to_dict(self) -> dict[str, Any]: + raise TypeError("cannot serialize") + + class UnsafeStateProvider(ContextProvider): + async def after_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, context, state + session.state["exploding"] = ExplodingState() + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + yield ChatResponseUpdate(contents=[Content.from_text(text="Completed")]) + + store = InMemoryAGUIThreadSnapshotStore() + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[UnsafeStateProvider("unsafe")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/unsafe-continuation", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "tenant-a", + keepalive_seconds=None, + ) + + response = TestClient(app).post( + "/unsafe-continuation", + json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Run"}]}, + ) + events = _decode_sse_events(response) + stored = await store.get(scope="tenant-a", thread_id="thread-1") + + assert "RUN_FINISHED" in [event.get("type") for event in events] + assert "RUN_ERROR" not in [event.get("type") for event in events] + assert stored is not None + assert stored.messages + assert stored.session_state is None + assert "Failed to serialize AG-UI Session Continuation State" in caplog.text + + 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.""" From 00d44f593c671957c3683625fcae3fa5eb2fd622 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Mon, 13 Jul 2026 15:22:50 +0900 Subject: [PATCH 5/5] Python: isolate AG-UI request state --- python/packages/ag-ui/README.md | 11 ++-- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 6 +- .../ag-ui/tests/ag_ui/test_endpoint.py | 56 +++++++++++++++++++ 3 files changed, 67 insertions(+), 6 deletions(-) diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 62e55de7904..692ebcfd063 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -370,11 +370,12 @@ It is not an authentication, tenant authorization, or distributed durability mec responsible for endpoint authentication, tenant authorization, and deployment/storage architecture that matches their availability and worker topology requirements. -Snapshot storage is trusted, integrity-protected server storage because private continuation is eligible for typed -core restoration. Snapshots also have confidentiality impact: they may contain sensitive user text, model output, -tool results, function arguments, UI payloads, Shared State, interrupt data, and private provider working state. The -built-in `InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable -production storage. It is cleared on process restart and is not shared across workers. +Snapshot storage is treated as trusted server-side storage because private continuation is eligible for typed core +restoration; applications are responsible for providing its integrity protection. Snapshots also have confidentiality +impact: they may contain sensitive user text, model output, tool results, function arguments, UI payloads, Shared State, +interrupt data, and private provider working state. The built-in `InMemoryAGUIThreadSnapshotStore` is in-memory only, +process-local, bounded, latest-only, and not durable production storage. It is cleared on process restart and is not +shared across workers. No file-backed AG-UI snapshot store is provided by the package. Applications that need durable persistence should provide an app-owned implementation of the `AGUIThreadSnapshotStore` protocol and own storage hardening, including diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 341e1edcbe5..84820bba408 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -1918,7 +1918,11 @@ async def run_agent_stream( _restore_session_continuation_state(session, stored_snapshot) protected_session_state_keys = _request_state_protected_keys(agent) session.state.update( - {key: value for key, value in flow.current_state.items() if key not in protected_session_state_keys} + { + key: copy.deepcopy(value) + for key, value in flow.current_state.items() + if key not in protected_session_state_keys + } ) _restore_tool_approval_state(session, approval_state_store, approval_thread_id) 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 5f329ff938c..ac5ba412565 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -418,6 +418,62 @@ async def stream_fn( assert text_deltas == [expected] +async def test_endpoint_provider_mutation_does_not_change_shared_state_snapshot( + streaming_chat_client_stub: Any, +) -> None: + """Provider mutations through the session view do not alter replayable Shared State.""" + + class MutatingProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, context, state + identity = session.state["identity"] + assert isinstance(identity, dict) + identity["type"] = "provider-mutated" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + yield ChatResponseUpdate(contents=[Content.from_text(text="Completed")]) + + agent = Agent( + name="test", + instructions=None, + client=streaming_chat_client_stub(stream_fn), + context_providers=[MutatingProvider("mutating")], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/isolated-request-state", + state_schema={"identity": {"type": "object"}}, + keepalive_seconds=None, + ) + + response = TestClient(app).post( + "/isolated-request-state", + json={ + "messages": [{"role": "user", "content": "Run"}], + "state": {"identity": {"type": "request"}}, + }, + ) + + state_snapshots = [ + event["snapshot"] for event in _decode_sse_events(response) if event.get("type") == "STATE_SNAPSHOT" + ] + assert state_snapshots == [{"identity": {"type": "request"}}] + + async def test_endpoint_restores_context_provider_state_across_scoped_runs(streaming_chat_client_stub): """Server-produced provider state survives sequential runs in one scoped thread."""