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

[Eduard van Valkenburg (eavanvalkenburg)]

Motivation & Context

AG-UI merges client Shared State into AgentSession.state, minus a protected set of server-owned namespaces. That set is derived from optional attributes of the object AG-UI is handed (middleware, context_providers, service_session_state_keys), none of which are part of SupportsAgentRun.

#8714 added custom ToolApprovalMiddleware.source_id namespaces to that set, but only for agents that expose their middleware. A wrapper that does not expose it leaves the wrapped agent's custom namespace unprotected. That covers a plain SupportsAgentRun implementation, and a BaseAgent subclass that delegates run() while keeping its own empty middleware. Client Shared State, including Shared State replayed from a stored Thread Snapshot, could then supply or replace the standing approval state that the wrapped middleware reads.

This completes the protection from #8714 for wrapped agents.

Description & Review Guide

• What are the major changes?
  • _agent_run.py now protects tool approval state by shape as well as by name. A top-level client Shared State value holding a non-empty rules, queued_approval_requests or collected_approval_responses list is not merged into AgentSession.state, whatever its item encoding (ToolApprovalState loads more than JSON objects, for example key/value pair lists). A key that already holds approval state in restored server state, as a mapping or a ToolApprovalState object, cannot be replaced by client values, whatever their shape. A bounded warning names the ignored keys.
  • Tests: the Python: Protect custom AG-UI approval state namespaces #8714 matrix also covers a protocol-only wrapper and a BaseAgent-derived wrapper (request and snapshot state, trusted and untrusted). New tests cover client values overwriting a trusted rule (held as a mapping or as a ToolApprovalState object), rule encodings, forged queued and collected approvals, the compatibility boundary, a human approval round trip through the wrapper, and the helpers, including a pin of the field names to ToolApprovalState.
  • python/packages/ag-ui/AGENTS.md documents the rule.
• What is the impact of these changes?
  • Wrapped agents get the same protection as direct agents. No API change.
  • Behavior change for any agent: a Shared State key whose mapping holds a non-empty rules, queued_approval_requests or collected_approval_responses list is no longer mirrored into AgentSession.state. It stays in Shared State, snapshots and the model context. This is why the PR is marked [BREAKING].
  • Unchanged: the other protected namespaces, the snapshot storage format, and ordinary Shared State.
• What do you want reviewers to focus on?
  • The compatibility trade-off: any non-empty approval list counts, because narrower item-shape checks can be sidestepped by encodings that ToolApprovalState accepts.
  • The deliberate choice not to recompute protected keys at snapshot-save time: a legitimate client key such as {"policy": {"rules": [{"id": 1}]}} would otherwise be persisted as server-owned and then frozen against client updates.

Related Issue

N/A - no separate issue was filed. This completes the protection added in #8714 for wrapped agents.

AI Assistance

• No material AI assistance was used.
• This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.

AI assistance details: implementation, tests, the documentation note, and this PR description were prepared with GitHub Copilot.

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR links to an agreed issue with no competing open PR, or the Related Issue section documents a trivial-change or repository-automation exception.
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

[Copilot]

🟡 Changes recommended

Object-valued ToolApprovalState remains replaceable when middleware metadata is hidden.

1 open finding

• Protect object-valued ToolApprovalState from client replacement · New

What changed in this PR

Protects AG-UI tool-approval state when wrappers hide middleware metadata.

Changes:

• Detects approval authority by state shape.
• Protects restored server-owned approval namespaces.
• Adds wrapper scenarios, security tests, and guidance.

File	Description
python/​packages/​ag-ui/​agent_framework_ag_ui/​_agent_run.py	Adds shape-based approval-state protection.
python/​packages/​ag-ui/​tests/​ag_ui/​test_run.py	Tests protection helpers and field synchronization.
python/​packages/​ag-ui/​tests/​ag_ui/​test_endpoint.py	Tests wrapped-agent approval scenarios.
python/​packages/​ag-ui/​AGENTS.md	Documents the state-authority rule.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

[github-actions]

MAF Automated Review — Iteration 1

Result: No findings
Scope: full PR (1 commit(s)): 4535f4ae8514
Model: gpt-5.6-sol

Overview

The PR adds shape-based filtering so client Shared State cannot inject tool-approval rules, queued requests, or collected responses through middleware namespaces hidden by wrappers. Its strongest guardrails combine name-based and shape-based protection before the session merge, with coverage for request and snapshot origins, alternate encodings, forged approval content, trusted restored mappings, and human approval resume. After deduplication against existing feedback and reconciliation against the base behavior, no additional publishable issue remains.

Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
No publishable findings remained after source verification for this scope.