Python: Reserve HA session id as reserved for AGUI

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

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(

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})"

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,

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``.
+"""

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,13 +24,15 @@
     ChatMiddleware,
     ChatResponse,
     ChatResponseUpdate,
+    Content,
     FunctionInvocationContext,
     FunctionMiddleware,
     Message,
     MiddlewareTermination,
     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."""