[BREAKING] Python: Protect AG-UI tool approval state from client Shared State

python/packages/ag-ui/AGENTS.md

@@ -124,6 +124,15 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
   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.
+- `ToolApprovalMiddleware` names its session state with an application-chosen `source_id`, which a wrapper can hide
+  from the agent AG-UI is handed. Tool approval authority is therefore also protected by shape, not only by name: a
+  client Shared State value holding a non-empty `rules`, `queued_approval_requests` or `collected_approval_responses`
+  list is never merged into `AgentSession.state`, whatever its item encoding (`ToolApprovalState` loads more than
+  JSON objects), and a key that already holds that state in restored server state, as a mapping or a
+  `ToolApprovalState` object, cannot be replaced by client values. An ignored key stays in Shared State and the
+  model context; it is only not mirrored into the session.
+  Wrapper agents should still forward `middleware`, `context_providers` and `service_session_state_keys`, and
+  `_TOOL_APPROVAL_STATE_FIELDS` must stay in sync with `ToolApprovalState`.
 - 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

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -37,6 +37,7 @@
     Message,
     SupportsAgentRun,
     ToolApprovalMiddleware,
+    ToolApprovalState,
     WorkflowAgent,
 )
 from agent_framework._middleware import (
@@ -2582,7 +2583,11 @@ def _a2ui_existing_tool_names(agent: SupportsAgentRun, tools: list[Any] | None)
 
 
 def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]:
-    """Return session-state namespaces that client Shared State cannot own."""
+    """Return session-state namespaces that client Shared State cannot own.
+
+    Names come from optional attributes of ``agent``, which a wrapper may hide, so ``run_agent_stream`` also protects
+    tool approval state by shape; see :func:`_carries_tool_approval_authority`.
+    """
     context_providers = cast(list[Any], getattr(agent, "context_providers", []))
     agent_middleware = categorize_middleware(getattr(agent, "middleware", None))["agent"]
     return {
@@ -2596,6 +2601,41 @@ def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]:
     }
 
 
+# Field names ToolApprovalMiddleware reads from ``session.state[source_id]`` (``ToolApprovalState``). They are spelled
+# literally so this check does not depend on the installed core version; a test pins them to ``ToolApprovalState``.
+_TOOL_APPROVAL_STATE_FIELDS: tuple[str, ...] = ("rules", "queued_approval_requests", "collected_approval_responses")
+
+
+def _tool_approval_state_keys(session_state: Mapping[str, Any]) -> set[str]:
+    """Return trusted session-state keys holding ToolApprovalMiddleware state, whatever ``source_id`` named them.
+
+    The middleware reads either a mapping or a ``ToolApprovalState`` object, so both forms are recognized.
+    """
+    return {
+        key
+        for key, value in session_state.items()
+        if isinstance(value, ToolApprovalState)
+        or (
+            isinstance(value, Mapping)
+            and any(field in cast(Mapping[str, Any], value) for field in _TOOL_APPROVAL_STATE_FIELDS)
+        )
+    }
+
+
+def _carries_tool_approval_authority(value: Any) -> bool:
+    """Return whether client state would give ToolApprovalMiddleware approvals the server never recorded.
+
+    ``ToolApprovalMiddleware`` names its state with an application-chosen ``source_id``, so approval authority is
+    recognized by shape rather than by name. Any non-empty approval list counts whatever its item encoding, because
+    ``ToolApprovalState`` loads more than JSON objects (for example key/value pair lists), so matching the item
+    shape would leave gaps.
+    """
+    if not isinstance(value, Mapping):
+        return False
+    state = cast(Mapping[str, Any], value)
+    return any(isinstance(state.get(field), list) and bool(state[field]) for field in _TOOL_APPROVAL_STATE_FIELDS)
+
+
 def _fides_session_state_keys(agent: SupportsAgentRun) -> set[str]:
     """Return standalone FIDES session-state keys that client Shared State cannot supply.
 
@@ -3225,7 +3265,24 @@ async def _run_agent_stream(
         restore_service_session_id=config.use_service_session,
         excluded_state_keys=set() if config.use_service_session else _provider_service_session_state_keys(agent),
     )
-    protected_session_state_keys = _request_state_protected_keys(agent)
+    # A wrapper can hide ToolApprovalMiddleware from the agent AG-UI was handed, so its application-chosen source_id
+    # may be unknown here. Approval authority is therefore also protected by shape: namespaces the server already
+    # holds, and client values that would supply approvals.
+    named_protected_keys = _request_state_protected_keys(agent)
+    authority_state_keys = {key for key, value in flow.current_state.items() if _carries_tool_approval_authority(value)}
+    protected_session_state_keys = {
+        *named_protected_keys,
+        *_tool_approval_state_keys(session.state),
+        *authority_state_keys,
+    }
+    if ignored_authority_keys := authority_state_keys.difference(named_protected_keys):
+        shown_keys = sorted(ignored_authority_keys)[:5]
+        omitted = len(ignored_authority_keys) - len(shown_keys)
+        logger.warning(
+            "Ignoring client Shared State %s%s: it carries tool approval authority, which is server-owned.",
+            ", ".join(repr(key[:120]) for key in shown_keys),
+            f" and {omitted} more" if omitted else "",
+        )
     session.state.update(
         {
             key: copy.deepcopy(value)