Python: A2UI (Agent-to-UI) support for the AG-UI adapter

[Ran Shemtov (ranst91)]

Motivation & Context

A2UI lets an agent generate rich, interactive UI (cards, forms, dashboards) that renders live in the client instead of plain text. It already shipped for .NET (#6494); this brings the same capability to the Python AG-UI adapter for parity. It reuses the shared, framework-agnostic ag-ui-a2ui-toolkit.

Description & Review Guide

• What are the major changes? Adds an in-package _a2ui module to agent-framework-ag-ui, delivering:

  • Progressive streaming — the render_a2ui sub-agent's argument deltas reach the wire as the surface builds, so the client paints it incrementally.
  • Error recovery — a validate-and-retry loop. An invalid surface never paints, and exhaustion produces a failure envelope instead of throwing.
  • Sub-agent based — generate_a2ui delegates UI design to a forced render_a2ui structured-output sub-agent, run through the toolkit's recovery loop.

  Plus two small bridge fixes needed for A2UI to behave in a real client: strip unanswered tool calls from replayed history before the planner call (A2UI surfaces persist as activities, so a surface action like a card button would otherwise fail the next turn), and skip the terminal MESSAGES_SNAPSHOT for A2UI runs so the streamed transcript order stays stable (mirrors the existing predictive-tool snapshot suppression). Also adds example agents and a unit suite.

• What is the impact of these changes? Purely additive. The toolkit is imported lazily behind an optional a2ui extra, so the base package is unchanged for anyone not using A2UI; the new a2ui_config parameters are optional with existing defaults. Verified end to end against a real model, not only fixtures.

• What do you want reviewers to focus on? The auto-injection gate in _agent_run.py (nullish precedence between the runtime injectA2UITool flag and a backend inject_a2ui_tool opt-in) and the streaming coalescing in _a2ui/_agent.py.

Related Issue

No dedicated Python issue; this is the Python counterpart of the .NET A2UI work in #6494.

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 is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change.

[Copilot]

Pull request overview

Adds Python-side A2UI (agent-to-UI) support to the AG-UI adapter for parity with the existing .NET implementation, enabling rich surface generation with progressive streaming and recovery semantics via the shared ag-ui-a2ui-toolkit.

Changes:

• Introduces a new agent_framework_ag_ui._a2ui module (state plumbing, context replay wrapper, A2UI generation wrapper, and auto-injection factory) with lazy imports so the base package remains usable without the optional toolkit.
• Updates the AG-UI hosting path to (a) stamp forwarded A2UI context into run options, (b) auto-inject the generate_a2ui tool when requested by the runtime, and (c) suppress the terminal snapshot for A2UI runs to preserve streamed ordering.
• Adds an A2UI-focused test suite and expands the examples server with A2UI demo endpoints.

Reviewed changes

Copilot reviewed 14 out of 15 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
python/uv.lock	Adds ag-ui-a2ui-toolkit and registers the new a2ui extra for agent-framework-ag-ui.
python/packages/ag-ui/tests/ag_ui/test_a2ui.py	New unit suite covering A2UI context plumbing, wrapper agents, streaming/non-streaming generation, and key regressions.
python/packages/ag-ui/pyproject.toml	Defines the a2ui optional dependency and configures mypy overrides for the toolkit.
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py	Adds a2ui_config parameter to the FastAPI endpoint helper and forwards it into config.
python/packages/ag-ui/agent_framework_ag_ui/_agent.py	Extends agent configuration to carry a2ui_config for runtime auto-injection.
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py	Implements A2UI auto-injection, AG-UI context stamping, and suppresses terminal snapshot for A2UI runs.
python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_state.py	New toolkit-free state/context utilities (slice building, stamping, stripping, history mapping).
python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_factory.py	New enable_a2ui helper and plan_a2ui_injection decision logic (Strands-parity rules).
python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_context_agent.py	New wrapper to replay forwarded A2UI catalog/guidelines into the prompt via a system message.
python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py	New A2UIAgent implementing streaming progressive paint + validate/retry and non-streaming tool execution.
python/packages/ag-ui/agent_framework_ag_ui/_a2ui/init.py	New lazy-export package initializer for A2UI-related symbols.
python/packages/ag-ui/agent_framework_ag_ui/init.py	Adds lazy exports for A2UI symbols at the package root.
python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py	Updates examples server to support OpenAI Chat Completions and adds A2UI demo endpoints.
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/a2ui_agents.py	Adds A2UI demo agents (dynamic schema, advanced zero-config, recovery, fixed schema).
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/init.py	Exposes new A2UI demo agents/config via the examples agents package.

Comments suppressed due to low confidence (1)

python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py:88

• After making python-dotenv optional, the .env load call also needs a guard; otherwise load_dotenv may be None and this will raise at runtime.

# Load the examples .env (OPENAI_API_KEY etc.) before constructing any chat client.
# override=True so a stale/empty exported OPENAI_API_KEY doesn't shadow the .env value.
load_dotenv(Path(__file__).resolve().parent.parent / ".env", override=True)

[Eduard van Valkenburg (eavanvalkenburg)]

Ran Shemtov (@ranst91) please use the defined PR template, and there are a number of merge conflicts

[Ran Shemtov (ranst91)]

Eduard van Valkenburg (@eavanvalkenburg) All comments (including about the PR body) seems to be addressed now. Lmk if there's anything else

[Evan Mattson (moonbox3)]

Hi Ran Shemtov (@ranst91), thanks for taking this on. The core of this PR looks solid. A2UIAgent hides a lot of behavior (progressive streaming, the validate-and-retry recovery loop, fragment coalescing) behind a small surface, and the test suite mostly drives it through run() with real doubles rather than mocks. That's great, thank you.

My main feedback is about how it wires into the run loop. The existing extension points in this package (predictive state via predict_state_config, HITL via require_confirmation) are declarative config that the core run loop consumes. A2UI instead swaps the agent object mid-run: _agent_run.py rebinds agent to the wrapper stack after server tools are collected but before the protected-state-key computation, approval resolution, and continuation-state serialization. Since the wrappers don't mirror context_providers or client, those reads see the wrong agent. Concretely: provider source_ids drop out of the protected keys (client shared state can then clobber provider namespaces) and approved-tool execution loses function middleware. I'd suggest expressing A2UI as config on AgentConfig and letting the core honor it, no object swap. That would also fix two related spots:

1. The terminal MESSAGES_SNAPSHOT skip matches the literal tool names "generate_a2ui"/"render_a2ui", so any user tool with that name silently loses its snapshot. Driving it off config, like _should_suppress_intermediate_snapshot already does for predictive tools, avoids that.
2. If the toolkit import fails, the run continues with the injected render_a2ui tool still advertised but with no executor, which produces exactly the unanswered-tool-call state your sanitizer exists to repair. Failing with a clear install message when A2UI was explicitly requested seems safer.

One smaller thing: AGUIContextAgent ports a .NET shape that no longer exists in this repo (the .NET package moved to the external AG-UI SDK), and the stamp/strip channel through additional_properties exists only to replicate ChatOptions.AdditionalProperties. Python already has ContextProvider for per-run instruction injection, and using it would also remove the footgun where standalone AGUIContextAgent use leaks ag_ui_context to the provider. Happy to discuss any of this, and thanks again, the streaming core itself is in great shape.

[Ran Shemtov (ranst91)]

Evan Mattson (@moonbox3) Thanks for the review. Comments should be addressed now. LMK if there's anything else

[Ran Shemtov (ranst91)]

Evan Mattson (@moonbox3) addressed the last round of comments. LMK if there's anything else