Python: A2UI (Agent-to-UI) support for the AG-UI adapter - #7423
Conversation
agent-framework-automation
Bot
added
the
python
There was a problem hiding this comment.
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._a2uimodule (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_a2uitool 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)
| from agent_framework.azure import AzureOpenAIChatClient | ||
| from agent_framework.openai import OpenAIChatCompletionClient | ||
| from dotenv import load_dotenv |
| # NOTE: requires the toolkit release that includes singular-`child` validation + | ||
| # `child_cycle` detection (ag-ui-protocol/ag-ui#1944); bump the floor to that | ||
| # published version once it ships (currently unreleased; 0.0.3 lacks it). |
| # The recovery-exhausted envelope builder is not part of the toolkit's public | ||
| # surface yet; import the shared implementation so the streaming twin and the | ||
| # synchronous loop cannot drift on the exhaustion envelope shape. | ||
| from ag_ui_a2ui_toolkit.recovery import _wrap_recovery_exhausted_envelope | ||
| from agent_framework import Content, FunctionTool, Message |
| _forwarded = input_data.get("forwarded_props") or input_data.get("forwardedProps") | ||
| if read_inject_a2ui_flag(_forwarded): |
|
@ranst91 please use the defined PR template, and there are a number of merge conflicts |
Adds an in-package _a2ui module to agent-framework-ag-ui delivering progressive-streaming, error-recovery, and sub-agent-based A2UI surface generation, reusing the shared ag-ui-a2ui-toolkit. Includes example agents, a unit suite, and two bridge fixes (strip unanswered tool calls from replayed history; suppress the terminal MESSAGES_SNAPSHOT for A2UI runs to keep streamed order stable). Signed-off-by: ran <ran@copilotkit.ai>
2d849e8 to
61ef22a
Compare
|
@eavanvalkenburg All comments (including about the PR body) seems to be addressed now. Lmk if there's anything else |
moonbox3
left a comment
moonbox3
left a comment
There was a problem hiding this comment.
Hi @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:
- 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.
- 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.
| # separately from forwardedProps at injection-decision time (not here). | ||
| a2ui_context_slice = build_ag_ui_context_slice(input_data.get("context")) | ||
| if a2ui_context_slice: | ||
| run_kwargs["options"] = stamp_context_slice(run_kwargs.get("options"), a2ui_context_slice) |
There was a problem hiding this comment.
Could we keep this adapter-private option off the raw-agent path unless A2UI is actually enabled? build_ag_ui_context_slice() returns a non-empty slice for any AG-UI context, but when injectA2UITool is unset or false no wrapper strips it. The raw OpenAI/Azure completion clients forward additional_properties as an unknown request option, so an existing non-A2UI request that supplies context can fail before the model call.
| pending, | ||
| stream=True, | ||
| session=pending_session, | ||
| tools=[*incoming_tools, generate_decl], |
There was a problem hiding this comment.
Could we handle the declaration-only A2UI call separately from ordinary executable calls? The FunctionInvocationLayer classifies the whole batch before execution, so when a model asks for generate_a2ui and an ordinary tool such as search together, the presence of generate_decl causes neither call to execute; the ordinary call is surfaced as user input and is not replayed in the wrapper's A2UI history. A common “look up data and render it” turn therefore silently skips the server-side operation.
| if target is not None: | ||
| gen_named[target] += frag | ||
| gen_all[target] += frag | ||
| elif name in ("", None) and cid in ("", None): |
There was a problem hiding this comment.
Could we associate nameless continuation fragments by the provider's tool-call index or call identity instead of one global active_gen? OpenAI-compatible Chat Completions can interleave argument deltas for parallel calls, and those continuation fragments arrive with empty id and name. After the second generate_a2ui opens, fragments for the first call are appended to the second, so two valid requests can become empty or concatenated JSON and generate surfaces with the wrong intent or an error envelope.
| try: | ||
| from ._a2ui import plan_a2ui_injection | ||
|
|
||
| existing_tool_names = [name for name in (getattr(t, "name", None) for t in (tools or [])) if name] |
There was a problem hiding this comment.
Could we include the agent's configured/default tool names in this duplicate check before deciding to inject? With no runtime tools, this list is empty even when the agent already has a default generate_a2ui; the wrapper then adds a second declaration and the core tool merge raises ValueError: Duplicate tool name 'generate_a2ui' before the provider call. This violates the documented user-prevails/no-double-injection path and turns backend opt-in into a request-time crash for an already-wired agent.
| Args: | ||
| inner_agent: The agent to wrap (any ``SupportsAgentRun``). | ||
| """ | ||
| self.inner_agent = inner_agent |
There was a problem hiding this comment.
Could we preserve the inner agent's provider namespaces when this wrapper is installed? The host computes protected session-state keys from the wrapped agent.context_providers; this wrapper exposes only identity, so client-supplied AG-UI state with a key such as tenant_auth is copied into the session and then supplied to the inner provider on the next call. That lets a caller overwrite provider-owned authorization or tenant state whenever A2UI auto-injection is enabled; delegating context_providers and the other protected-state metadata would preserve the existing boundary.
| Args: | ||
| inner_agent: The agent to wrap (any ``SupportsAgentRun``). | ||
| """ | ||
| self.inner_agent = inner_agent |
There was a problem hiding this comment.
Could we retain or delegate the inner client here? On an approval resume, the AG-UI runner executes statically available approved tools through its approval resolver, which builds the function middleware pipeline from getattr(agent, "client", None). Because this wrapper has no client, inner client-level authorization or audit middleware is skipped and an approved sensitive tool runs through an empty pipeline.
Reworks A2UI so it no longer swaps the agent object mid-run, and fixes the issues that swap caused. - Drive A2UI through a dedicated runner used only for the stream call; keep the original agent bound so protected-state-key computation, approval resolution, and continuation serialization still read its real context_providers and client (no more provider-namespace or approval middleware loss). - Hand the forwarded AG-UI context to the runner directly instead of stamping it onto run-option additional_properties. That channel leaked the slice to the provider SDK on any run carrying AG-UI context, including non-A2UI runs where nothing stripped it back. Removes the stamp/strip/read helpers and the dead .NET-shaped path. - Suppress the terminal MESSAGES_SNAPSHOT off whether A2UI actually drove the run, not the literal tool names, so an unrelated user tool named "generate_a2ui" keeps its snapshot. - Fail loud with an install hint when A2UI is requested but the toolkit isn't installed, instead of advertising render_a2ui with no executor. - Include the agent's own default tools in the no-double-injection check so an already-wired agent doesn't crash on a duplicate tool name. - Execute ordinary developer tools called in the same turn as generate_a2ui (the declaration-only tool poisons the inner batch invocation), so a "look up data then render it" turn no longer skips the backend call. - Attribute nameless streaming argument deltas by the provider tool-call index so interleaved parallel calls don't cross-contaminate; the OpenAI chat client preserves that index on the content. Adds tests for the mixed-batch execution, index-based fragment attribution, and the default-tool duplicate check. Signed-off-by: ran <ran@copilotkit.ai>
|
@moonbox3 Thanks for the review. Comments should be addressed now. LMK if there's anything else |
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
_a2uimodule toagent-framework-ag-ui, delivering:render_a2uisub-agent's argument deltas reach the wire as the surface builds, so the client paints it incrementally.generate_a2uidelegates UI design to a forcedrender_a2uistructured-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_SNAPSHOTfor 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
a2uiextra, so the base package is unchanged for anyone not using A2UI; the newa2ui_configparameters 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 runtimeinjectA2UIToolflag and a backendinject_a2ui_toolopt-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