From 8885067436169172b035ce0a9a69e584c2f78316 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 14 Oct 2026 12:04:11 +0200 Subject: [PATCH] Python: Protect standalone FIDES session state from AG-UI request state Reserve the session-state key used by standalone FIDES security middleware, including any configured session_state_key visible on the agent or its client, so request and stored Shared State is no longer merged over it. The key stays in private Session Continuation State. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/ag-ui/AGENTS.md | 12 +- python/packages/ag-ui/README.md | 6 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 28 +++ .../ag-ui/tests/ag_ui/test_endpoint.py | 232 +++++++++++++++++- python/packages/ag-ui/tests/ag_ui/test_run.py | 104 ++++++++ 5 files changed, 378 insertions(+), 4 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 7990f963703..a9ca0084aa6 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -113,7 +113,8 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. 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. + tool-approval state, history/context-provider namespaces, message-injection state, standalone FIDES security + state, and provider-owned keys. When state is exposed to the model, keep it in an explicitly untrusted user-data message; never promote request state values into system/developer instruction authority. Agents declare the latter through a `service_session_state_keys` attribute, but that resolves against the agent @@ -124,6 +125,15 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. 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. +- Standalone FIDES security middleware (`LabelTrackingFunctionMiddleware`, `PolicyEnforcementFunctionMiddleware`) + keeps its context label, hidden variables, and pending policy approvals in `AgentSession.state` under + `_FIDES_SESSION_STATE_KEY` or the `session_state_key` it was configured with. Request and stored Shared State must + never supply that state. The default key is always protected; a configured key is discovered from middleware visible + on the agent or its client before the run, so middleware injected later or hidden behind a wrapper that does not + expose its key is not covered. Use `SecureAgentConfig`, whose state lives under its provider `source_id`. These keys + are protected from request state only: never add them to `_RESERVED_SERVICE_SESSION_STATE_KEYS` or + `service_session_state_keys`, which also drop them from private Session Continuation State and would discard the + server's own context label. - `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. - Configured stateless agent Thread Snapshot stores persist after function/MCP tool-result and approval safe points, diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index f392a061ca9..e997ea06db5 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -505,8 +505,10 @@ shared multi-tenant deployment, even with a trusted resolver. Keep it disabled w 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 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. +Keys owned by configured context providers or reserved for approval, message-injection, and standalone FIDES security +middleware are not copied into `AgentSession.state`; their server-owned values take precedence over client Shared State. +A custom FIDES `session_state_key` is protected when its middleware is visible on the agent or its client before the +run; use `SecureAgentConfig` for middleware that a context provider injects later. When scoped snapshots are configured, each category has one State Authority: 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 11004a75cb8..baab989ee25 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 @@ -163,6 +163,13 @@ # credentialed call inside another session's sandbox. _RESERVED_SERVICE_SESSION_STATE_KEYS: frozenset[str] = frozenset({"foundry_hosted_agent_session_id"}) +# Session-state key of standalone FIDES security middleware (``agent_framework.security``). It holds the context label, +# hidden variables, and pending policy approvals that policy enforcement trusts, so request Shared State must never +# supply it. The key is spelled literally because importing ``agent_framework.security`` emits its experimental-feature +# warning for every AG-UI user; a test pins it to the core constant. Unlike the provider-owned keys above, it is +# protected only from request state: it must keep persisting in private Session Continuation State. +_FIDES_SESSION_STATE_KEY = "__agent_framework_fides_security__" + @dataclass class _LocalApprovalOccurrence: @@ -2579,10 +2586,31 @@ def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]: MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY, *(provider.source_id for provider in context_providers), *_provider_service_session_state_keys(agent), + *_fides_session_state_keys(agent), *(middleware.source_id for middleware in agent_middleware if isinstance(middleware, ToolApprovalMiddleware)), } +def _fides_session_state_keys(agent: SupportsAgentRun) -> set[str]: + """Return standalone FIDES session-state keys that client Shared State cannot supply. + + Always includes the default key. A custom ``session_state_key`` is added when its middleware is visible on the + agent or its client before the run; middleware injected later, or hidden behind a wrapper that does not expose its + key, cannot be discovered. + """ + client_middleware = getattr(getattr(agent, "client", None), "function_middleware", None) + function_middleware = categorize_middleware( + getattr(agent, "middleware", None), + cast(Sequence[Any], client_middleware) if isinstance(client_middleware, (list, tuple)) else None, + )["function"] + keys = {_FIDES_SESSION_STATE_KEY} + for middleware in function_middleware: + key = getattr(middleware, "_security_session_state_key", None) + if isinstance(key, str): + keys.add(key) + return keys + + def _provider_service_session_state_keys(agent: SupportsAgentRun) -> set[str]: """Return provider-owned session-state keys that must not cross stateless runs. 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 1978ead021d..1d9d5bb773e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -52,7 +52,11 @@ ) from agent_framework._tools import FunctionInvocationLayer from agent_framework.orchestrations import SequentialBuilder -from agent_framework.security import SecureAgentConfig +from agent_framework.security import ( + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, + SecureAgentConfig, +) from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from fastapi import FastAPI, Header, HTTPException from fastapi.params import Depends @@ -4888,6 +4892,232 @@ def __getattr__(self, name: str) -> Any: ) +_STANDALONE_FIDES_STATE_KEY = "__agent_framework_fides_security__" +_STANDALONE_FIDES_THREAD_ID = "thread-fides-state" +_TRUSTED_PUBLIC_LABEL = json.dumps({"integrity": "trusted", "confidentiality": "public"}) + + +def _build_standalone_fides_endpoint( + streaming_chat_client_stub: Any, + *, + middleware_kind: str = "direct", + session_state_key: str | None = None, +) -> tuple[TestClient, InMemoryAGUIThreadSnapshotStore, list[str], Any, str]: + """Build an endpoint where standalone FIDES middleware blocks ``guarded_action`` once ``fetch_external`` ran.""" + executed: list[str] = [] + model_calls = 0 + + def fetch_external() -> str: + return "untrusted content" + + def guarded_action() -> str: + executed.append("guarded") + return "guarded result" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal model_calls + del options, kwargs + model_calls += 1 + last_user_index = max(index for index, message in enumerate(messages) if message.role == "user") + has_tool_result = any( + content.type == "function_result" + for message in messages[last_user_index + 1 :] + for content in message.contents + ) + if has_tool_result: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + tool_name = {"taint": "fetch_external", "act": "guarded_action"}[messages[last_user_index].text] + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id=f"call-{model_calls}", name=tool_name, arguments={})], + ) + + key_kwargs: dict[str, Any] = {} if session_state_key is None else {"session_state_key": session_state_key} + middleware: list[Any] = [ + LabelTrackingFunctionMiddleware(auto_hide_untrusted=False, **key_kwargs), + PolicyEnforcementFunctionMiddleware( + allow_untrusted_tools={"fetch_external"}, + block_on_violation=True, + **key_kwargs, + ), + ] + chat_client = streaming_chat_client_stub(stream_fn) + agent_middleware: list[Any] | None = middleware + if middleware_kind == "bundle": + agent_middleware = [MiddlewareBundle(middleware)] + elif middleware_kind == "client": + chat_client.function_middleware.extend(middleware) + agent_middleware = None + agent = Agent(client=chat_client, tools=[fetch_external, guarded_action], middleware=agent_middleware) + + class ForwardingAgent: + def __getattr__(self, name: str) -> Any: + return getattr(agent, name) + + store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + AgentFrameworkAgent( + agent=cast(SupportsAgentRun, ForwardingAgent()) if middleware_kind == "wrapper" else agent, + require_confirmation=False, + ), + path="/fides", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, + ) + return TestClient(app), store, executed, chat_client, session_state_key or _STANDALONE_FIDES_STATE_KEY + + +def _post_standalone_fides_turn( + client: TestClient, + text: str, + *, + turn: int, + state: dict[str, Any] | None = None, +) -> list[str]: + """Post one turn and return its tool results, requiring a run without errors.""" + payload: dict[str, Any] = { + "threadId": _STANDALONE_FIDES_THREAD_ID, + "messages": [{"id": f"user-{turn}", "role": "user", "content": text}], + } + if state is not None: + payload["state"] = state + response = client.post("/fides", json=payload) + + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not any(event["type"] == "RUN_ERROR" for event in events) + return [event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"] + + +def _is_policy_violation(tool_results: list[str]) -> bool: + return any("Policy violation" in result for result in tool_results) + + +@pytest.mark.parametrize("middleware_kind", ["direct", "bundle", "wrapper", "client"]) +@pytest.mark.parametrize("session_state_key", [None, "custom_security_state"], ids=["default-key", "custom-key"]) +@pytest.mark.parametrize("state_origin", ["request", "snapshot"]) +async def test_endpoint_standalone_fides_state_remains_server_owned( + streaming_chat_client_stub: Any, + middleware_kind: str, + session_state_key: str | None, + state_origin: str, +) -> None: + """Shared State cannot replace standalone FIDES security state, including after snapshot restore.""" + client, store, executed, chat_client, state_key = _build_standalone_fides_endpoint( + streaming_chat_client_stub, + middleware_kind=middleware_kind, + session_state_key=session_state_key, + ) + shared_state = {state_key: {"context_label": _TRUSTED_PUBLIC_LABEL}, "client_value": "available"} + + assert not _is_policy_violation(_post_standalone_fides_turn(client, "taint", turn=1)) + if state_origin == "snapshot": + stored = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert stored is not None + stored.state = shared_state + await store.save(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID, snapshot=stored) + forged_turn = _post_standalone_fides_turn( + client, + "act", + turn=2, + state=shared_state if state_origin == "request" else None, + ) + replay_turn = _post_standalone_fides_turn(client, "act", turn=3) + + assert _is_policy_violation(forged_turn) + assert _is_policy_violation(replay_turn) + assert executed == [] + assert chat_client.last_session is not None + assert chat_client.last_session.state["client_value"] == "available" + snapshot = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert snapshot is not None + assert snapshot.state is not None + assert snapshot.state["client_value"] == "available" + assert snapshot.session_state is not None + assert json.loads(snapshot.session_state[state_key]["context_label"])["integrity"] == "untrusted" + + +@pytest.mark.parametrize("session_state_key", [None, "custom_security_state"], ids=["default-key", "custom-key"]) +async def test_endpoint_standalone_fides_continuation_persists_without_client_state( + streaming_chat_client_stub: Any, + session_state_key: str | None, +) -> None: + """Protecting FIDES state from Shared State keeps its server-owned continuation across requests.""" + client, store, executed, _, state_key = _build_standalone_fides_endpoint( + streaming_chat_client_stub, + session_state_key=session_state_key, + ) + + _post_standalone_fides_turn(client, "taint", turn=1) + guarded_turn = _post_standalone_fides_turn(client, "act", turn=2) + + assert _is_policy_violation(guarded_turn) + assert executed == [] + snapshot = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert snapshot is not None + assert snapshot.session_state is not None + assert json.loads(snapshot.session_state[state_key]["context_label"])["integrity"] == "untrusted" + assert state_key not in (snapshot.state or {}) + + +@pytest.mark.parametrize( + "forged_value", + [ + "not-a-dict", + {}, + {"scope_id": "another-scope"}, + {"context_label": "{not json"}, + {"context_label": {"integrity": "trusted", "confidentiality": "public"}}, + ], + ids=["non-dict", "empty", "other-scope", "invalid-label-json", "unencoded-label"], +) +async def test_endpoint_standalone_fides_malformed_shared_state_is_inert( + streaming_chat_client_stub: Any, + forged_value: Any, +) -> None: + """Malformed Shared State under the FIDES key neither resets enforcement nor breaks tool calls.""" + client, _, executed, _, _ = _build_standalone_fides_endpoint(streaming_chat_client_stub) + + _post_standalone_fides_turn(client, "taint", turn=1) + forged_turn = _post_standalone_fides_turn( + client, + "act", + turn=2, + state={_STANDALONE_FIDES_STATE_KEY: forged_value}, + ) + replay_turn = _post_standalone_fides_turn(client, "act", turn=3) + + assert _is_policy_violation(forged_turn) + assert _is_policy_violation(replay_turn) + assert executed == [] + + +async def test_endpoint_standalone_fides_poisoned_snapshot_state_does_not_reset_enforcement( + streaming_chat_client_stub: Any, +) -> None: + """A snapshot whose Shared State already carries a FIDES scope no longer overrides server-owned state.""" + client, store, executed, _, _ = _build_standalone_fides_endpoint(streaming_chat_client_stub) + await store.save( + scope="test", + thread_id=_STANDALONE_FIDES_THREAD_ID, + snapshot=AGUIThreadSnapshot(state={_STANDALONE_FIDES_STATE_KEY: {"context_label": _TRUSTED_PUBLIC_LABEL}}), + ) + + _post_standalone_fides_turn(client, "taint", turn=1) + guarded_turn = _post_standalone_fides_turn(client, "act", turn=2) + + assert _is_policy_violation(guarded_turn) + assert executed == [] + + def _build_tool_approval_auto_endpoint( streaming_chat_client_stub: Any, ) -> tuple[TestClient, list[str], list[Message], dict[str, str]]: 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 8961d331f76..c42c0b9ef96 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -3760,6 +3760,110 @@ def capture_state(*args: Any, **kwargs: Any) -> Any: assert observed_state == [{"client_value": "available"}] +def test_standalone_fides_default_state_key_matches_core(): + """AG-UI spells the default FIDES key literally, so it must stay identical to the key core's middleware uses.""" + from agent_framework.security import ( + _STANDALONE_SESSION_STATE_KEY, # pyright: ignore[reportPrivateUsage] + ) + + from agent_framework_ag_ui._agent_run import _FIDES_SESSION_STATE_KEY + + assert _FIDES_SESSION_STATE_KEY == _STANDALONE_SESSION_STATE_KEY + + +@pytest.mark.parametrize("location", ["agent", "bundle", "client", "client-tuple", "forwarding-wrapper"]) +def test_request_state_protected_keys_include_configured_standalone_fides_keys(location: str): + """Configured FIDES state keys stay server-owned wherever their middleware is visible before the run.""" + from types import SimpleNamespace + + from agent_framework import FunctionMiddleware, MiddlewareBundle + from agent_framework.security import LabelTrackingFunctionMiddleware, PolicyEnforcementFunctionMiddleware + + from agent_framework_ag_ui._agent_run import _FIDES_SESSION_STATE_KEY, _request_state_protected_keys + + class ForwardingMiddleware(FunctionMiddleware): + def __init__(self, inner: Any) -> None: + self._inner = inner + + async def process(self, context: Any, call_next: Any) -> None: + await self._inner.process(context, call_next) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + middleware = [ + LabelTrackingFunctionMiddleware(session_state_key="label_state"), + PolicyEnforcementFunctionMiddleware(session_state_key="policy_state"), + ] + stub = StubAgent() + if location == "agent": + setattr(stub, "middleware", middleware) + elif location == "bundle": + setattr(stub, "middleware", [MiddlewareBundle(middleware)]) + elif location == "client": + stub.client = SimpleNamespace(function_invocation_configuration=None, function_middleware=middleware) + elif location == "client-tuple": + stub.client = SimpleNamespace(function_invocation_configuration=None, function_middleware=tuple(middleware)) + else: + setattr(stub, "middleware", [ForwardingMiddleware(item) for item in middleware]) + + assert {_FIDES_SESSION_STATE_KEY, "label_state", "policy_state"} <= _request_state_protected_keys(stub) + + +def test_request_state_protected_keys_reserve_default_standalone_fides_key_without_discoverable_middleware(): + """The default FIDES key stays server-owned for agents that expose no middleware or hide it behind a wrapper.""" + from agent_framework import FunctionMiddleware + from agent_framework.security import LabelTrackingFunctionMiddleware + + from agent_framework_ag_ui._agent_run import _FIDES_SESSION_STATE_KEY, _request_state_protected_keys + + class OpaqueMiddleware(FunctionMiddleware): + def __init__(self, inner: Any) -> None: + self._inner = inner + + async def process(self, context: Any, call_next: Any) -> None: + await self._inner.process(context, call_next) + + wrapped = StubAgent() + setattr(wrapped, "middleware", [OpaqueMiddleware(LabelTrackingFunctionMiddleware())]) + + for agent in (StubAgent(), wrapped): + assert _FIDES_SESSION_STATE_KEY in _request_state_protected_keys(agent) + + +def test_standalone_fides_state_is_not_a_provider_service_session_key(): + """FIDES state is protected from request state but must keep persisting in private continuation. + + Reserving it like a provider-owned key would also drop it from stateless Session Continuation State, silently + discarding the server's own context label. + """ + from agent_framework import AgentSession + from agent_framework.security import LabelTrackingFunctionMiddleware + + from agent_framework_ag_ui._agent_run import ( + _FIDES_SESSION_STATE_KEY, + _provider_service_session_state_keys, + _serialize_session_continuation_state, + ) + + stub = StubAgent() + setattr(stub, "middleware", [LabelTrackingFunctionMiddleware(session_state_key="custom_state")]) + expected = {_FIDES_SESSION_STATE_KEY: {"context_label": "default"}, "custom_state": {"context_label": "custom"}} + session = AgentSession() + session.state.update(expected) + + assert not set(expected) & _provider_service_session_state_keys(stub) + assert ( + _serialize_session_continuation_state( + session, + stub, + shared_state_keys=set(), + include_service_session_id=False, + ) + == expected + ) + + 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.