Python: Protect custom AG-UI approval state namespaces

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

@@ -36,6 +36,7 @@
     MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
     Message,
     SupportsAgentRun,
+    ToolApprovalMiddleware,
     WorkflowAgent,
 )
 from agent_framework._middleware import (
@@ -2449,12 +2450,14 @@ 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."""
     context_providers = cast(list[Any], getattr(agent, "context_providers", []))
+    agent_middleware = categorize_middleware(getattr(agent, "middleware", None))["agent"]
     return {
         _TOOL_APPROVAL_STATE_KEY,
         InMemoryHistoryProvider.DEFAULT_SOURCE_ID,
         MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
         *(provider.source_id for provider in context_providers),
         *_provider_service_session_state_keys(agent),
+        *(middleware.source_id for middleware in agent_middleware if isinstance(middleware, ToolApprovalMiddleware)),
     }
 
 

python/packages/ag-ui/tests/ag_ui/test_endpoint.py

@@ -36,6 +36,7 @@
     InMemoryCheckpointStorage,
     InMemoryHistoryProvider,
     Message,
+    MiddlewareBundle,
     MiddlewareFailure,
     SessionContext,
     SupportsAgentRun,
@@ -3829,6 +3830,109 @@ async def stream_fn(
     return TestClient(app), executed, messages_received, state, wrapped_agent
 
 
+@pytest.mark.parametrize("middleware_kind", ["direct", "bundle", "wrapper"])
+@pytest.mark.parametrize("state_origin", ["request", "snapshot"])
+@pytest.mark.parametrize("trusted_approval", [False, True])
+async def test_endpoint_custom_approval_namespaces_remain_server_owned(
+    streaming_chat_client_stub: Any,
+    middleware_kind: str,
+    state_origin: str,
+    trusted_approval: bool,
+) -> None:
+    """Shared State cannot supply or replace approval rules, including after snapshot restore."""
+    executed: list[str] = []
+
+    def gated_tool() -> str:
+        executed.append("ran")
+        return "done"
+
+    async def stream_fn(
+        messages: list[Message],
+        options: dict[str, Any],
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        if executed:
+            yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")
+        else:
+            yield ChatResponseUpdate(
+                contents=[Content.from_function_call(call_id="call_gated", name="gated_tool", arguments="{}")],
+                role="assistant",
+            )
+
+    source_ids = ("custom_approvals", "other_approvals")
+    approval = ToolApprovalMiddleware(source_id="initial_approvals")
+    middleware = [approval, ToolApprovalMiddleware(source_id=source_ids[1])]
+    chat_client = streaming_chat_client_stub(stream_fn)
+    agent = Agent(
+        client=chat_client,
+        tools=[
+            FunctionTool(name="gated_tool", description="Test tool", func=gated_tool, approval_mode="always_require")
+        ],
+        middleware=[MiddlewareBundle(middleware)] if middleware_kind == "bundle" else middleware,
+    )
+
+    class ForwardingAgent:
+        def __getattr__(self, name: str) -> Any:
+            return getattr(agent, name)
+
+    wrapped = AgentFrameworkAgent(
+        agent=cast(SupportsAgentRun, ForwardingAgent()) if middleware_kind == "wrapper" else agent,
+        require_confirmation=False,
+    )
+    # Namespace discovery must use the live middleware, not construction-time configuration.
+    approval.source_id = source_ids[0]
+    rule = {"tool_name": "gated_tool"}
+    shared_state = {
+        **{source_id: {"rules": [] if trusted_approval else [rule]} for source_id in source_ids},
+        "client_value": "available",
+    }
+    store = InMemoryAGUIThreadSnapshotStore()
+    await store.save(
+        scope="test",
+        thread_id="thread-approval-state",
+        snapshot=AGUIThreadSnapshot(
+            state=shared_state if state_origin == "snapshot" else None,
+            session_state={source_id: {"rules": [rule]} for source_id in source_ids} if trusted_approval else None,
+        ),
+    )
+    app = FastAPI()
+    add_agent_framework_fastapi_endpoint(
+        app,
+        wrapped,
+        path="/approval",
+        snapshot_store=store,
+        snapshot_scope_resolver=lambda _request: "test",
+    )
+    payload: dict[str, Any] = {
+        "threadId": "thread-approval-state",
+        "runId": "run-approval-state",
+        "messages": [{"role": "user", "content": "Run the test tool"}],
+    }
+    if state_origin == "request":
+        payload["state"] = shared_state
+    response = TestClient(app).post("/approval", json=payload)
+
+    assert response.status_code == 200
+    events = _decode_sse_events(response)
+    assert not any(event["type"] == "RUN_ERROR" for event in events)
+    finished = next(event for event in events if event["type"] == "RUN_FINISHED")
+    assert executed == (["ran"] if trusted_approval else [])
+    if trusted_approval:
+        assert "outcome" not in finished
+    else:
+        assert finished["outcome"]["type"] == "interrupt"
+        assert len(finished["outcome"]["interrupts"]) == 1
+    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="thread-approval-state")
+    assert snapshot is not None
+    assert snapshot.session_state is not None
+    for source_id in source_ids:
+        assert [stored_rule["tool_name"] for stored_rule in snapshot.session_state[source_id]["rules"]] == (
+            ["gated_tool"] if trusted_approval else []
+        )
+
+
 def _build_tool_approval_auto_endpoint(
     streaming_chat_client_stub: Any,
 ) -> tuple[TestClient, list[str], list[Message], dict[str, str]]: