diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 59bb2cfc17..683db67bdd 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -59,6 +59,16 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. migrations and must warn because it disables that isolation; it is unsafe for shared multi-tenant deployments. A configured endpoint resolver must return a non-empty string. Reject invalid results before accessing state or invoking the runner; absence of a resolver, not an invalid result, selects intentionally unscoped operation. +- Request `state` is client-owned Shared State and is merged into `AgentSession.state`, minus a protected set: + tool-approval state, history/context-provider namespaces, message-injection state, and provider-owned keys. + Agents declare the latter through a `service_session_state_keys` attribute, but that resolves against the agent + object AG-UI is handed, so a wrapper agent that does not forward it would silently drop the protection. + `_RESERVED_SERVICE_SESSION_STATE_KEYS` therefore reserves such keys unconditionally; add a key there whenever + it names a remote resource that the server's own credentialed call addresses. +- Provider-owned keys identify remote resources that the server's own credentialed call addresses — for example + the Foundry hosted-agent session ID, which selects a VM-isolated sandbox with a persistent filesystem. Never let + request state choose one. `AgentSession.service_session_id` is deliberately a separate attribute rather than a + `state` entry, so conversation continuation is unreachable from client input by construction; keep it that way. - `confirm_changes` snapshot cleanup resolves the synthetic confirmation back to its original `function_call_id`; it must never concatenate unrelated tool results or record accepted changes without a matching real result. - SSE keepalive is endpoint-owned transport behavior configured through diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index eece20fdcc..e6ecb54a82 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 @@ -8,7 +8,7 @@ import json import logging import uuid -from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Collection, Mapping, Sequence from dataclasses import dataclass, field from functools import partial from typing import TYPE_CHECKING, Any, cast @@ -139,6 +139,18 @@ _COLLECTED_APPROVAL_RESPONSES_KEY = "collected_approval_responses" _PROVIDER_SERVICE_SESSION_ID_STATE_KEY = "__ag_ui_provider_service_session_id" +# Provider-owned session-state keys reserved unconditionally, independent of what the agent declares through +# ``service_session_state_keys``. Agents are expected to declare their own keys, but that resolves through the +# agent object AG-UI is handed, so a wrapper agent that does not forward the attribute would silently drop the +# protection. Keys naming a remote resource that the server's own credentialed call addresses are reserved here +# instead. The Foundry key is spelled literally rather than imported, because agent-framework-ag-ui does not +# depend on agent-framework-foundry and must protect the key even when it is absent or outdated. +# +# agent_framework_foundry.FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY. Selects the Foundry hosted-agent runtime session, +# a VM-isolated sandbox with a persistent filesystem, so a client-supplied value would run the server's +# credentialed call inside another session's sandbox. +_RESERVED_SERVICE_SESSION_STATE_KEYS: frozenset[str] = frozenset({"foundry_hosted_agent_session_id"}) + @dataclass class _LocalApprovalOccurrence: @@ -2324,11 +2336,16 @@ def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]: def _provider_service_session_state_keys(agent: SupportsAgentRun) -> set[str]: - """Return provider-owned session-state keys that must not cross stateless runs.""" - keys = getattr(agent, "service_session_state_keys", ()) - if not isinstance(keys, (list, tuple, set, frozenset)): - return set() - return {key for key in keys if isinstance(key, str)} + """Return provider-owned session-state keys that must not cross stateless runs. + + Always includes :data:`_RESERVED_SERVICE_SESSION_STATE_KEYS`, so the reserved keys stay server-owned even + when ``agent`` declares nothing, such as a wrapper agent that does not forward the attribute. + """ + keys = set(_RESERVED_SERVICE_SESSION_STATE_KEYS) + declared = getattr(agent, "service_session_state_keys", ()) + if isinstance(declared, (list, tuple, set, frozenset)): + keys.update(key for key in cast(Collection[Any], declared) if isinstance(key, str)) + return keys def _serialize_session_continuation_state( diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index e76fad0715..a6fab9957a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -2922,3 +2922,89 @@ def capture_state(*args: Any, **kwargs: Any) -> Any: {"client_value": "available"}, {"private": "preserved", "client_value": "available"}, ] + + +async def test_reserved_service_session_keys_protected_without_agent_declaration(): + """Reserved provider keys stay server-owned even if the agent declares nothing. + + ``service_session_state_keys`` is resolved from the agent object AG-UI is handed, so an agent that does not + declare it, such as a wrapper that forgets to forward the attribute or an older provider package, would + otherwise silently fall back to accepting a client-supplied value. Reserving the key in + ``_RESERVED_SERVICE_SESSION_STATE_KEYS`` is the backstop that keeps those combinations safe. + """ + from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + + from agent_framework_ag_ui import AgentFrameworkAgent + + stub = StubAgent() + assert not hasattr(stub, "service_session_state_keys") + + observed_state: list[dict[str, Any]] = [] + original_run = stub.run + + def capture_state(*args: Any, **kwargs: Any) -> Any: + observed_state.append(dict(kwargs["session"].state)) + return original_run(*args, **kwargs) + + stub.run = capture_state # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + + agent = AgentFrameworkAgent(agent=stub) + _ = [ + event + async for event in agent.run( + { + "thread_id": "frontend-thread", + "run_id": "attacker-run", + "messages": [{"role": "user", "content": "Hello"}], + "state": { + "foundry_hosted_agent_session_id": "victim-sandbox", + "client_value": "available", + }, + } + ) + ] + + assert observed_state == [{"client_value": "available"}] + + +async def test_request_state_cannot_assign_service_session_id(): + """Chat history held by the provider is addressed by ``service_session_id``, never by client state. + + ``AgentSession`` keeps ``service_session_id`` as its own attribute rather than a ``state`` entry, and only + trusted snapshot storage or a provider response may set it. This test locks in that boundary so a future + refactor cannot start sourcing conversation continuation from request Shared State. + """ + from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + + from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore + + for use_service_session in (False, True): + stub = StubAgent() + observed: list[Any] = [] + original_run = stub.run + + def capture(*args: Any, **kwargs: Any) -> Any: + observed.append(kwargs["session"].service_session_id) + return original_run(*args, **kwargs) + + stub.run = capture # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + + agent = AgentFrameworkAgent( + agent=stub, + use_service_session=use_service_session, + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + ) + _ = [ + event + async for event in agent.run( + { + "thread_id": "frontend-thread", + "run_id": "attacker-run", + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "Hello"}], + "state": {"__ag_ui_provider_service_session_id": "victim-conversation"}, + } + ) + ] + + assert observed == [None], f"client state selected a conversation (use_service_session={use_service_session})" diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 8175c071a7..393a710938 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -658,6 +658,14 @@ class RawFoundryAgent( """ service_session_state_keys: ClassVar[frozenset[str]] = frozenset({FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY}) + """Session-state keys this agent owns, which untrusted input must never supply. + + Holds the Foundry hosted-agent session ID. Despite the attribute name, that value is not a conversation or + an ``AgentSession.service_session_id``: it identifies the hosted agent's *runtime session*, a VM-isolated + sandbox with a persistent filesystem. Hosts such as AG-UI read this to reject client-supplied values, since + honouring one would run the server's own credentialed call inside another session's sandbox. Hosts are + expected to reserve this key independently as well, because a wrapper agent may not forward this attribute. + """ def __init__( self, diff --git a/python/packages/foundry/agent_framework_foundry/_constants.py b/python/packages/foundry/agent_framework_foundry/_constants.py index 66f2f70cf4..dbafaca994 100644 --- a/python/packages/foundry/agent_framework_foundry/_constants.py +++ b/python/packages/foundry/agent_framework_foundry/_constants.py @@ -1,3 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY = "foundry_hosted_agent_session_id" +"""``AgentSession.state`` key holding the Foundry hosted-agent session ID. + +This is the hosted agent's runtime session, a VM-isolated sandbox with a persistent filesystem, sent as +``extra_body["agent_session_id"]``. It is distinct from ``AgentSession.service_session_id``, which continues the +model-side response or conversation chain. The value is server-owned; see +``RawFoundryAgent.service_session_state_keys``. +""" diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 4c3b6876e6..3fe9654149 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -6,7 +6,7 @@ import json import os import sys -from collections.abc import Awaitable, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from importlib import import_module from types import SimpleNamespace from typing import Any, cast @@ -24,6 +24,7 @@ ChatMiddleware, ChatResponse, ChatResponseUpdate, + Content, FunctionInvocationContext, FunctionMiddleware, Message, @@ -31,6 +32,7 @@ WorkflowBuilder, tool, ) +from agent_framework._types import ResponseStream from agent_framework.foundry import FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore from agent_framework_openai._chat_client import RawOpenAIChatClient @@ -1095,6 +1097,134 @@ def my_func() -> str: assert agent.default_options.get("tools") is not None +async def test_agui_request_state_cannot_choose_foundry_hosted_agent_sandbox() -> None: + """A client must not be able to pick which Foundry sandbox the server's credentialed call addresses. + + Reproduces the full reported chain without network access: an ordinary AG-UI request carries + ``foundry_hosted_agent_session_id`` in its client-owned ``state``, and the resulting session is then used to + prepare the outbound Foundry Responses call. ``agent_session_id`` selects a VM-isolated sandbox with a + persistent filesystem, so honouring a caller-supplied value would run the server's own credentialed request + inside another session's sandbox. + """ + mock_project = MagicMock() + mock_project.get_openai_client.return_value = MagicMock() + foundry_agent = RawFoundryAgent(project_client=mock_project, agent_name="test-agent") + + observed_state: list[dict[str, Any]] = [] + + def fake_run(messages: Any = None, **kwargs: Any) -> Any: + session = kwargs.get("session") + observed_state.append(dict(session.state) if session is not None else {}) + + async def _stream() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(text="ok")], role="assistant") + + return ResponseStream(_stream(), finalizer=AgentResponse.from_updates) + + foundry_agent.run = fake_run # type: ignore[method-assign] # ty: ignore[invalid-assignment] + + runner = AgentFrameworkAgent(agent=foundry_agent) + events = [ + event + async for event in runner.run({ + "runId": "attacker-run", + "threadId": "attacker-thread", + "messages": [{"role": "user", "content": "hello"}], + "state": { + FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY: "VICTIM-SANDBOX-9f31c2", + "benign": "ok", + }, + }) + ] + + assert not [event for event in events if getattr(event, "type", None) == "RUN_ERROR"] + # The sandbox handle is stripped, while ordinary client Shared State still reaches the agent. + assert observed_state == [{"benign": "ok"}] + + # The outbound Foundry call therefore pins no sandbox at all. + session = AgentSession() + session.state.update(observed_state[0]) + with patch( + "agent_framework._agents.RawAgent._prepare_run_context", + new=AsyncMock(return_value={"ok": True}), + ) as mock_prepare_run_context: + await foundry_agent._prepare_run_context( + messages="hi", + session=session, + tools=None, + options={}, + compaction_strategy=None, + tokenizer=None, + function_invocation_kwargs=None, + client_kwargs=None, + ) + + assert mock_prepare_run_context.await_args + assert "agent_session_id" not in mock_prepare_run_context.await_args.kwargs["options"].get("extra_body", {}) + + +async def test_agui_request_state_cannot_overwrite_established_foundry_sandbox() -> None: + """A server-established sandbox handle stays authoritative when a later request tries to replace it.""" + store = InMemoryAGUIThreadSnapshotStore() + mock_project = MagicMock() + mock_openai = MagicMock() + mock_openai.conversations.create = AsyncMock(return_value=SimpleNamespace(id="conv_server_owned")) + mock_project.get_openai_client.return_value = mock_openai + foundry_agent = RawFoundryAgent(project_client=mock_project, agent_name="test-agent") + + observed_state: list[dict[str, Any]] = [] + + def fake_run(messages: Any = None, **kwargs: Any) -> Any: + session = kwargs.get("session") + observed_state.append(dict(session.state) if session is not None else {}) + + async def _stream() -> AsyncIterator[AgentResponseUpdate]: + # The provider establishes the real sandbox handle, exactly as a live response would. + if session is not None: + session.state[FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY] = "server-owned-sandbox" + yield AgentResponseUpdate(contents=[Content.from_text(text="ok")], role="assistant") + + return ResponseStream(_stream(), finalizer=AgentResponse.from_updates) + + foundry_agent.run = fake_run # type: ignore[method-assign] # ty: ignore[invalid-assignment] + + runner = AgentFrameworkAgent( + agent=foundry_agent, + use_service_session=True, + snapshot_store=store, + ) + payload: dict[str, Any] = { + "threadId": "victim-thread", + "__ag_ui_snapshot_scope": "scope", + "messages": [{"role": "user", "content": "hello"}], + } + _ = [event async for event in runner.run({**payload, "runId": "run-1"})] + + attacker_payload = { + **payload, + "runId": "run-2", + "state": {FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY: "attacker-chosen-sandbox"}, + } + _ = [event async for event in runner.run(attacker_payload)] + + assert observed_state[0] == {} + # The second turn resumes the server's own sandbox and ignores the attacker's value. + assert observed_state[1] == {FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY: "server-owned-sandbox"} + + +def test_foundry_agents_declare_hosted_agent_session_id_as_server_owned() -> None: + """The hosted-agent session ID must stay server-owned so hosts reject client-supplied values. + + ``agent_session_id`` selects the Foundry hosted-agent runtime session, which is a VM-isolated sandbox with a + persistent filesystem. A caller-supplied value would redirect the server's own credentialed call into another + session's sandbox, so hosts discover this key through ``service_session_state_keys`` and refuse to let + untrusted input populate it. Without this test nothing fails if the declaration is dropped. + """ + assert FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY in RawFoundryAgent.service_session_state_keys + # FoundryAgent is the recommended production class, so it must inherit the same protection. + assert FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY in FoundryAgent.service_session_state_keys + + async def test_raw_foundry_agent_prepare_run_context_injects_agent_session_id_from_state() -> None: """Test that hosted-agent session state is sent separately from response continuation."""