diff --git a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py index 9be38154a33..5369913ba9f 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py @@ -3,6 +3,7 @@ """AG-UI protocol integration for Agent Framework.""" import importlib.metadata +from typing import TYPE_CHECKING, Any from ._agent import AgentFrameworkAgent from ._client import AGUIChatClient @@ -22,6 +23,9 @@ from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata from ._workflow import AgentFrameworkWorkflow, WorkflowFactory +if TYPE_CHECKING: + from ._a2ui import A2UIAgent, AGUIContextAgent, enable_a2ui, plan_a2ui_injection + try: __version__ = importlib.metadata.version(__name__) except importlib.metadata.PackageNotFoundError: @@ -53,4 +57,21 @@ "DEFAULT_TAGS", "state_update", "__version__", + # A2UI (lazy — require ag-ui-a2ui-toolkit) + "A2UIAgent", + "AGUIContextAgent", + "enable_a2ui", + "plan_a2ui_injection", ] + +# A2UI symbols are loaded lazily so importing this package does not require +# ag-ui-a2ui-toolkit unless A2UI is actually used. +_A2UI_LAZY = {"A2UIAgent", "AGUIContextAgent", "enable_a2ui", "plan_a2ui_injection"} + + +def __getattr__(name: str) -> Any: + if name in _A2UI_LAZY: + import importlib + + return getattr(importlib.import_module("._a2ui", __name__), name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/__init__.py new file mode 100644 index 00000000000..0306c8a3e13 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/__init__.py @@ -0,0 +1,59 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""A2UI (agent-generated UI) support for the AG-UI Agent Framework adapter. + +Toolkit-dependent exports (``AGUIContextAgent`` and, later, ``A2UIAgent`` / +``get_a2ui_tools`` / ``plan_a2ui_injection``) are loaded lazily via PEP 562 so this +package — and the toolkit-free context-plumbing helpers in :mod:`._state` it +re-exports — can be imported by the hosting loop even when ``ag-ui-a2ui-toolkit`` is +not installed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +# Toolkit-free helpers — safe to import eagerly (no ag_ui_a2ui_toolkit dependency). +from ._state import ( + A2UI_CONTEXT_KEY, + A2UI_SCHEMA_CONTEXT_DESCRIPTION, + build_ag_ui_context_slice, + read_agent_state, + read_inject_a2ui_flag, + stamp_context_slice, +) + +if TYPE_CHECKING: + from ._agent import A2UIAgent + from ._context_agent import AGUIContextAgent + from ._factory import enable_a2ui, plan_a2ui_injection + +__all__ = [ + "A2UI_CONTEXT_KEY", + "A2UI_SCHEMA_CONTEXT_DESCRIPTION", + "A2UIAgent", + "AGUIContextAgent", + "build_ag_ui_context_slice", + "enable_a2ui", + "plan_a2ui_injection", + "read_agent_state", + "read_inject_a2ui_flag", + "stamp_context_slice", +] + +# Lazy (toolkit-dependent) exports: module -> symbols defined there. +_LAZY: dict[str, tuple[str, ...]] = { + "_agent": ("A2UIAgent",), + "_context_agent": ("AGUIContextAgent",), + "_factory": ("enable_a2ui", "plan_a2ui_injection"), +} + + +def __getattr__(name: str) -> Any: + for module, symbols in _LAZY.items(): + if name in symbols: + import importlib + + mod = importlib.import_module(f".{module}", __name__) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py new file mode 100644 index 00000000000..617828a0640 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py @@ -0,0 +1,781 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""A2UIAgent — wraps an agent with A2UI surface-generation support. + +Port of the .NET ``A2UIAgent`` (microsoft/agent-framework#6494). Every run gets a +``generate_a2ui`` tool that delegates UI generation to a render sub-agent and returns +a validated A2UI operations envelope as its tool result, run through the shared +toolkit's validate→retry recovery loop. + +Two paths mirror the .NET adapter: + +* **Non-streaming** advertises a REAL ``generate_a2ui`` tool whose body runs the + toolkit's synchronous ``run_a2ui_generation_with_recovery``; ordinary automatic + function invocation executes it. +* **Streaming** advertises ``generate_a2ui`` as a declaration-only tool (``func=None``) + so the planner's call surfaces on the update stream instead of being auto-invoked. + An agent-level planner-rounds loop then runs the render sub-agent with a streaming + chat call, forwards its per-chunk ``render_a2ui`` argument fragments (progressive + paint), balances the forwarded call with a ``{"status": "rendered"}`` tool result, + feeds the envelope back to the planner, and continues — the same wire shape the + LangGraph adapters produce. MAF-python is async-native, so the streaming loop runs + on one event loop with no worker-thread/queue bridge. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import AsyncIterator +from typing import Any, cast + +from ag_ui_a2ui_toolkit import ( + BASIC_CATALOG_ID, + GENERATE_A2UI_ARG_DESCRIPTIONS, + MAX_A2UI_ATTEMPTS, + RENDER_A2UI_TOOL_DEF, + A2UIToolParams, + augment_prompt_with_validation_errors, + build_a2ui_envelope, + prepare_a2ui_request, + resolve_a2ui_catalog, + resolve_a2ui_tool_params, + run_a2ui_generation_with_recovery, + validate_a2ui_components, + wrap_error_envelope, +) +from agent_framework import Content, FunctionTool, Message + +from ._state import read_agent_state, strip_context_slice, to_history_messages + +logger = logging.getLogger(__name__) + + +def _recovery_exhausted_envelope(max_attempts: int, attempts: list[dict[str, Any]]) -> str: + """Build the surface envelope shown when recovery attempts are exhausted. + + Prefers the toolkit's shared implementation so the streaming twin and the + synchronous loop cannot drift on the exhaustion-envelope shape. That helper is + private (``_wrap_recovery_exhausted_envelope``) and not part of the toolkit's + public surface, so a patch release could rename or promote it; fall back to the + public :func:`wrap_error_envelope` rather than failing hard. + """ + try: + from ag_ui_a2ui_toolkit.recovery import _wrap_recovery_exhausted_envelope + except ImportError: + return cast( + str, + wrap_error_envelope(f"A2UI generation failed after {max_attempts} attempt(s)."), + ) + return cast(str, _wrap_recovery_exhausted_envelope(max_attempts, attempts)) + + +# Inner render_a2ui tool name (canonical, from the toolkit's tool def). +RENDER_A2UI_TOOL_NAME: str = RENDER_A2UI_TOOL_DEF["function"]["name"] +_RENDER_PARAMETERS: dict[str, Any] = RENDER_A2UI_TOOL_DEF["function"]["parameters"] +_RENDER_DESCRIPTION: str = RENDER_A2UI_TOOL_DEF["function"]["description"] + +# Bare acknowledgement returned as the inner render_a2ui tool result; the painted +# surface rides the streamed arguments, so the result only has to balance the call. +RENDER_ACKNOWLEDGEMENT = '{"status": "rendered"}' + +# Cap on planner rounds (model turn -> generation -> result fed back) per run, +# guarding against a planner that keeps requesting surfaces without terminating. +MAX_PLANNER_ROUNDS = 8 + + +def _generate_tool_schema() -> dict[str, Any]: + """JSON schema for the planner-facing generate_a2ui tool (string args).""" + return { + "type": "object", + "properties": { + "intent": {"type": "string", "description": GENERATE_A2UI_ARG_DESCRIPTIONS["intent"]}, + "target_surface_id": { + "type": "string", + "description": GENERATE_A2UI_ARG_DESCRIPTIONS["target_surface_id"], + }, + "changes": {"type": "string", "description": GENERATE_A2UI_ARG_DESCRIPTIONS["changes"]}, + }, + } + + +def _as_tool_list(tools: Any) -> list[Any]: + if tools is None: + return [] + if isinstance(tools, (list, tuple)): + return list(tools) + return [tools] + + +def _string_arg(args: dict[str, Any], name: str) -> str | None: + value = args.get(name) + return value if isinstance(value, str) else None + + +def _normalize_catalog(catalog: dict[str, Any] | None) -> dict[str, Any] | None: + """Coerce a catalog into the shape ``validate_a2ui_components`` expects. + + The validator reads ``catalog["components"]`` as a NAME→schema MAPPING + (``catalog_components.get(ctype)``). Frontend catalogs — and the schema the + A2UI middleware forwards in ``RunAgentInput.context`` — list components as an + ARRAY of ``{"name": ..., ...}`` definitions instead. Re-key the array by + component name so catalog-membership / required-prop checks run instead of + crashing. Returns ``None`` (structural validation only) when there is no usable + component mapping — matching the toolkit's ``resolve_a2ui_catalog``, which + surfaces no component schema for the forwarded-catalog path. + """ + if not isinstance(catalog, dict): + return None + comps = catalog.get("components") + if isinstance(comps, dict): + return catalog + if isinstance(comps, list): + mapping = {c["name"]: c for c in comps if isinstance(c, dict) and isinstance(c.get("name"), str) and c["name"]} + return {**catalog, "components": mapping} if mapping else None + return None + + +def _resolve_catalog(state: dict[str, Any], explicit_catalog: dict[str, Any] | None) -> dict[str, Any] | None: + """Resolve the structured catalog used for validation. + + An explicit ``params.catalog`` wins (nullish — it is the host's deliberate + choice). Otherwise the forwarded A2UI schema (``state["ag-ui"]["a2ui_schema"]``, + a JSON string or dict carrying the component catalog) is used so catalog-membership + and required-prop checks run, not just structural ones. The result is normalized to + the validator's name→schema mapping shape (see :func:`_normalize_catalog`). Warns and + degrades to no catalog (structural checks only) on empty/unparseable/non-object + schema rather than failing the run. + """ + if explicit_catalog is not None: + return _normalize_catalog(explicit_catalog) + schema = (state.get("ag-ui") or {}).get("a2ui_schema") + if schema is None or schema == "": + return None + if isinstance(schema, dict): + return _normalize_catalog(schema) + if isinstance(schema, str): + try: + parsed = json.loads(schema) + except (ValueError, TypeError): + logger.warning("A2UI schema context is not valid JSON; validating without a catalog.") + return None + if isinstance(parsed, dict): + return _normalize_catalog(parsed) + logger.warning("A2UI schema context did not parse to an object; validating without a catalog.") + return None + logger.warning("A2UI schema context has unexpected type %s; validating without a catalog.", type(schema)) + return None + + +def _effective_default_catalog_id(configured: str, state: dict[str, Any]) -> str: + """Prefer the frontend-forwarded catalog id when the host configured none. + + ``resolve_a2ui_tool_params`` fills an unset ``default_catalog_id`` with the basic + catalog. Zero-config demos rely on the catalog the client registered (forwarded in + ``RunAgentInput.context``), so when the configured id is the basic fallback and a + forwarded catalog id is available, bind generated surfaces to that instead. An + explicitly-configured non-basic id always wins. + """ + if configured and configured != BASIC_CATALOG_ID: + return configured + resolved = resolve_a2ui_catalog(state) + forwarded_id = resolved[1] if resolved else None + return forwarded_id or configured + + +def _first_parsable_object(*candidates: str) -> dict[str, Any] | None: + """Return the first candidate string that parses to a JSON object. + + Streaming tool-call arguments must be reassembled in a way robust to how the + provider emits them: + + * MAF's Chat-Completions client emits argument DELTAS (``name=""``, + ``call_id=""`` after the first fragment) for progressive visibility PLUS a + FINAL COALESCED fragment repeating the tool name with the FULL arguments + (verified against gpt-4o: the client yields the args twice, so the all-fragment + concat is exactly ``2x`` the name-bearing concat). Concatenating only the + NAME-BEARING fragments (opening ``""`` + coalesced full) yields complete JSON; + concatenating ALL fragments would append the deltas AND the coalesced copy and + double the JSON into an unparseable string. + * A deltas-only provider (no coalesced fragment) yields the full JSON only when + ALL fragments are concatenated. + + Passing ``(name_bearing_concat, all_fragment_concat)`` is robust to both: the + name-bearing concat wins when a coalesced fragment is present (no doubling), and + the all-fragment concat is the fallback otherwise. Empty / non-object / unparsable + candidates are skipped. + """ + for cand in candidates: + if not cand: + continue + try: + obj = json.loads(cand) + except (ValueError, TypeError): + continue + if isinstance(obj, dict): + return obj + return None + + +def _sanitize_unanswered_tool_calls(messages: list[Any]) -> list[Any]: + """Drop function-call contents whose call id has no matching function-result, and + drop any message left with no contents. + + A2UI surfaces are rendered as ACTIVITIES, so the runtime persists the assistant's + ``generate_a2ui`` / ``render_a2ui`` tool calls but NOT their tool results (those + became the rendered surface, not a tool message). On a later turn — e.g. a surface + action ("Book") — the replayed history then contains an assistant message whose + tool calls are unanswered, which OpenAI-family providers reject at the planner call + ("400: an assistant message with 'tool_calls' must be followed by tool messages + responding to each 'tool_call_id'"). An unanswered tool call cannot be sent to the + provider, so strip the dangling calls before the planner runs (keeping any text and + balanced call/result pairs). The surface itself already rendered client-side. + """ + answered: set[str] = set() + for m in messages: + for c in getattr(m, "contents", None) or []: + if getattr(c, "type", None) == "function_result": + cid = getattr(c, "call_id", None) + if cid: + answered.add(cid) + + out: list[Any] = [] + for m in messages: + contents = getattr(m, "contents", None) + if not contents: + out.append(m) + continue + kept = [ + c + for c in contents + if not (getattr(c, "type", None) == "function_call" and getattr(c, "call_id", None) not in answered) + ] + if len(kept) == len(contents): + out.append(m) + elif kept: + out.append(Message(role=getattr(m, "role", "assistant"), contents=kept)) + # else: message had only unanswered calls -> drop it entirely + return out + + +def _is_recoverable_error(exc: BaseException) -> bool: + """Classify a sub-agent error. Programmer errors and cancellation rethrow; + everything else is a recoverable attempt failure (retry). + + Mirrors the .NET/TS classify: rethrow cancellation, TypeError, NameError, and any + non-Exception BaseException (SystemExit/KeyboardInterrupt). No undici exemption is + needed — Python transports do not raise TypeError for transient network failures. + """ + if isinstance(exc, asyncio.CancelledError): + return False + if not isinstance(exc, Exception): + return False + if isinstance(exc, (TypeError, NameError)): + return False + return True + + +class A2UIAgent: + """Wraps an agent so each run can generate A2UI surfaces. + + Args: + inner_agent: The agent to wrap (any ``SupportsAgentRun``). + subagent_chat_client: A raw chat client (NO automatic function invocation) + used to run the UI-generation sub-agent; the adapter reads the forced + ``render_a2ui`` call's arguments directly. + params: Behavior knobs (model is not needed here — the sub-agent client is + supplied directly); defaults are filled per the shared toolkit rules. + """ + + def __init__( + self, + inner_agent: Any, + subagent_chat_client: Any, + params: A2UIToolParams | None = None, + ) -> None: + self.inner_agent = inner_agent + self.subagent_chat_client = subagent_chat_client + self.params = resolve_a2ui_tool_params(params or {}) + self.id = getattr(inner_agent, "id", None) + self.name = getattr(inner_agent, "name", None) + self.description = getattr(inner_agent, "description", None) + + # -- public run ------------------------------------------------------- + + def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any: + """Run the wrapped agent with A2UI support. + + Returns an async iterator (stream) or an awaitable (non-stream), matching + the inner agent's protocol. + """ + if stream: + return self._run_streaming(messages, **kwargs) + return self._run_non_streaming(messages, **kwargs) + + # -- non-streaming ---------------------------------------------------- + + async def _run_non_streaming(self, messages: Any = None, **kwargs: Any) -> Any: + """Advertise a real generate_a2ui tool and let auto-invocation run it.""" + from agent_framework import normalize_messages + + normalized = _sanitize_unanswered_tool_calls(normalize_messages(messages)) + options = kwargs.get("options") + tools = _as_tool_list(kwargs.pop("tools", None)) + generate_tool = self._build_generate_tool(normalized, options) + merged = [*[t for t in tools if getattr(t, "name", None) != self.params["tool_name"]], generate_tool] + # Strip the adapter-private ag-ui context slice from options before it reaches + # the planner's chat client (additional_properties is forwarded to the provider + # SDK, which rejects the unknown key). The slice was already read above. + kwargs["options"] = strip_context_slice(options) + return await self.inner_agent.run(normalized, stream=False, tools=merged, **kwargs) + + def _build_generate_tool(self, conversation: list[Any], options: Any) -> FunctionTool: + """Build an EXECUTABLE generate_a2ui tool whose body runs the recovery loop.""" + state = {"ag-ui": read_agent_state(options)} + history = to_history_messages(conversation) + catalog = _resolve_catalog(state, self.params["catalog"]) + + async def _generate_a2ui( + intent: str | None = None, + target_surface_id: str | None = None, + changes: str | None = None, + ) -> str: + prep = prepare_a2ui_request( + intent=intent, + target_surface_id=target_surface_id, + changes=changes, + messages=history, + state=state, + guidelines=self.params["guidelines"], + ) + prep_error = prep.get("error") + if prep_error: + return cast(str, wrap_error_envelope(prep_error)) + + def _run_recovery() -> dict[str, Any]: + # The toolkit recovery loop is synchronous and calls invoke_subagent + # once per attempt. Drive every attempt on ONE dedicated event loop + # (created here, on the executor thread) rather than a fresh + # asyncio.run per attempt: a per-attempt loop closes after attempt 1, + # and a retry then fails with "Event loop is closed" because the + # sub-agent chat client's async transport stays bound to the first + # (now-closed) loop. + worker_loop = asyncio.new_event_loop() + # Set the worker loop as this executor thread's current loop so a + # sub-agent client that reaches for asyncio.get_event_loop() binds to it. + asyncio.set_event_loop(worker_loop) + try: + + def _invoke(prompt: str, attempt: int) -> dict[str, Any] | None: + # Mirror the streaming path's error handling: a recoverable + # (transient) sub-agent error becomes a failed attempt the + # toolkit loop retries (the toolkit does NOT catch invoke_subagent + # exceptions itself); a programmer error / cancellation propagates. + try: + return worker_loop.run_until_complete(self._invoke_render_subagent(prompt, conversation)) + except BaseException as exc: # noqa: BLE001 — reclassified below + if not _is_recoverable_error(exc): + raise + logger.warning( + "A2UI render sub-agent failed (non-streaming attempt %d, recoverable): %s", + attempt, + exc, + ) + return None + + recovery_result: dict[str, Any] = run_a2ui_generation_with_recovery( + base_prompt=prep.get("prompt", ""), + invoke_subagent=_invoke, + build_envelope=lambda args: build_a2ui_envelope( + args=args, + is_update=prep.get("is_update", False), + target_surface_id=target_surface_id, + prior=prep.get("prior"), + default_surface_id=self.params["default_surface_id"], + default_catalog_id=_effective_default_catalog_id(self.params["default_catalog_id"], state), + ), + catalog=catalog, + config=self.params["recovery"], + on_attempt=self.params["on_a2ui_attempt"], + ) + return recovery_result + finally: + asyncio.set_event_loop(None) + worker_loop.close() + + # Run the whole synchronous recovery loop in an executor so we never nest + # event loops on the agent's running loop. + loop = asyncio.get_running_loop() + result = await loop.run_in_executor(None, _run_recovery) + return cast(str, result["envelope"]) + + return FunctionTool( + name=self.params["tool_name"], + description=self.params["tool_description"], + func=_generate_a2ui, + input_model=_generate_tool_schema(), + ) + + async def _invoke_render_subagent(self, prompt: str, conversation: list[Any]) -> dict[str, Any] | None: + """Run the render sub-agent (forced render_a2ui) and return its args, or None.""" + response = await self.subagent_chat_client.get_response( + self._subagent_messages(prompt, conversation), + stream=False, + options=self._subagent_options(), + ) + return self._read_render_args(response) + + # -- streaming -------------------------------------------------------- + + async def _run_streaming(self, messages: Any = None, **kwargs: Any) -> Any: + from agent_framework import AgentResponseUpdate, normalize_messages + + history = _sanitize_unanswered_tool_calls(normalize_messages(messages)) + session = kwargs.pop("session", None) + options = kwargs.get("options") + incoming_tools = [ + t for t in _as_tool_list(kwargs.pop("tools", None)) if getattr(t, "name", None) != self.params["tool_name"] + ] + state = {"ag-ui": read_agent_state(options)} + # Strip the adapter-private ag-ui context slice from options before it reaches + # the planner's chat client on any round (or the closing turn): the slice was + # read into `state` above, and additional_properties is forwarded to the + # provider SDK, which rejects the unknown key. All inner_agent.run calls below + # pass **kwargs, so cleaning it once here covers every planner turn. + kwargs["options"] = strip_context_slice(options) + generate_decl = FunctionTool( + name=self.params["tool_name"], + description=self.params["tool_description"], + func=None, + input_model=_generate_tool_schema(), + ) + + pending: list[Any] = history + pending_session = session + for _round in range(1, MAX_PLANNER_ROUNDS + 1): + text_contents: list[Content] = [] + # Coalesce the planner's generate_a2ui call(s) by call id. A single call is + # streamed as MANY function_call fragments: an opening fragment carries + # id + name, argument-delta fragments carry name=""/call_id="", and MAF adds + # a final coalesced fragment repeating id + name + the FULL args. Treating + # each fragment as its own call would run the sub-agent repeatedly and emit + # duplicate tool results for one call id (an unbalanced assistant/tool + # sequence the provider rejects on the next turn). For each call we keep BOTH + # the name-bearing concat (opening + coalesced) and the all-fragment concat + # (name-bearing + deltas), then pick the first that parses — robust whether or + # not the provider emits a coalesced fragment, without doubling the args. + gen_order: list[str] = [] + gen_named: dict[str, str] = {} + gen_all: dict[str, str] = {} + active_gen: str | None = None # call id whose name="" deltas we attribute + async for update in self.inner_agent.run( + pending, + stream=True, + session=pending_session, + tools=[*incoming_tools, generate_decl], + **kwargs, + ): + for content in update.contents: + ctype = getattr(content, "type", None) + if ctype == "text": + text_contents.append(content) + continue + if ctype != "function_call": + continue + name = getattr(content, "name", None) + cid = getattr(content, "call_id", None) + raw = getattr(content, "arguments", None) + frag = raw if isinstance(raw, str) else (json.dumps(raw) if isinstance(raw, dict) else "") + if name == self.params["tool_name"]: + # Opening or coalesced fragment of a generate_a2ui call. + if cid: + if cid not in gen_named: + gen_order.append(cid) + gen_named[cid] = "" + gen_all[cid] = "" + active_gen = cid + target = cid or active_gen + if target is not None: + gen_named[target] += frag + gen_all[target] += frag + elif name in ("", None) and cid in ("", None): + # Continuation delta — belongs to the active generate call (if any). + # A named non-generate tool call clears active_gen below, so dev-tool + # deltas are never misattributed to generate. + if active_gen is not None: + gen_all[active_gen] += frag + else: + # A different named tool call opened — generate is no longer active. + active_gen = None + yield update + + if not gen_order: + return + + # One coalesced function_call per distinct generate_a2ui call id, selecting + # the first parsable of (name-bearing concat, all-fragment concat). + generate_calls = [] + for cid in gen_order: + obj = _first_parsable_object(gen_named[cid], gen_all[cid]) + args_str = json.dumps(obj) if obj is not None else "" + generate_calls.append( + Content.from_function_call(call_id=cid, name=self.params["tool_name"], arguments=args_str) + ) + assistant_contents = [*text_contents, *generate_calls] + + results: list[Content] = [] + for call in generate_calls: + box: list[Any] = [None] + async for update in self._run_generate_streaming(call, history, state, box): + yield update + results.append(Content.from_function_result(call_id=call.call_id or "", result=box[0])) + + # Surface the generate_a2ui tool results on the wire and feed back. + yield AgentResponseUpdate(role="tool", contents=results) + history = [ + *history, + Message(role="assistant", contents=assistant_contents), + Message(role="tool", contents=results), + ] + pending = history + pending_session = None + + # Planner kept requesting generations to the cap. Give it one final turn to + # consume the last result and narrate, with the generate tool withheld so it + # cannot request another surface (otherwise the run ends on an unanswered + # tool result with no closing assistant message). + async for update in self.inner_agent.run( + history, + stream=True, + session=None, + tools=incoming_tools, + **kwargs, + ): + yield update + + async def _run_generate_streaming( + self, + call: Any, + conversation: list[Any], + state: dict[str, Any], + box: list[Any], + ) -> AsyncIterator[Any]: + """Run one generate_a2ui invocation with the streaming validate-retry loop. + + Forwards the render sub-agent's per-chunk updates (progressive paint) and + deposits the final envelope into ``box[0]``. + """ + from agent_framework import AgentResponseUpdate + + args = call.parse_arguments() or {} + intent = _string_arg(args, "intent") + target_surface_id = _string_arg(args, "target_surface_id") + changes = _string_arg(args, "changes") + + prep = prepare_a2ui_request( + intent=intent, + target_surface_id=target_surface_id, + changes=changes, + messages=to_history_messages(conversation), + state=state, + guidelines=self.params["guidelines"], + ) + prep_error = prep.get("error") + if prep_error: + box[0] = wrap_error_envelope(prep_error) + return + + catalog = _resolve_catalog(state, self.params["catalog"]) + max_attempts = (self.params["recovery"] or {}).get("maxAttempts", MAX_A2UI_ATTEMPTS) + attempts: list[dict[str, Any]] = [] + last_errors: list[dict[str, str]] = [] + for attempt in range(1, max_attempts + 1): + prompt = augment_prompt_with_validation_errors(prep.get("prompt", ""), last_errors) + updates: list[Any] = [] + # Track the live render call id so a mid-stream failure can still close it + # — an unanswered (unbalanced) tool call is a protocol violation and the + # retry would open a second one on top. + live_render_call_id: str | None = None + try: + async for cru in self.subagent_chat_client.get_response( + self._subagent_messages(prompt, conversation), + stream=True, + options=self._subagent_options(), + ): + updates.append(cru) + if live_render_call_id is None: + live_render_call_id = self._first_render_call_id(cru) + # Forward each per-chunk update so the hosting loop paints the + # render_a2ui argument fragments progressively. + yield AgentResponseUpdate( + contents=cru.contents, + role=getattr(cru, "role", None), + message_id=getattr(cru, "message_id", None), + response_id=getattr(cru, "response_id", None), + raw_representation=cru, + ) + render_call_id, render_args = self._extract_render_args(updates) + except BaseException as exc: + # Classify BEFORE yielding anything. A non-recoverable exception includes + # GeneratorExit / CancelledError raised into this async generator when the + # consumer closes it (client disconnect / aclose): yielding during + # GeneratorExit raises "async generator ignored GeneratorExit", masking a + # clean cancellation. Re-raise those immediately, without a balancing yield. + if not _is_recoverable_error(exc): + raise + # Recoverable provider error mid-stream: close the live render call with a + # balancing tool result so the retry does not open a second one on top and + # the next turn's history stays balanced. + if live_render_call_id is not None: + yield AgentResponseUpdate( + role="tool", + contents=[ + Content.from_function_result(call_id=live_render_call_id, result=RENDER_ACKNOWLEDGEMENT) + ], + ) + logger.warning("A2UI render sub-agent failed (attempt %d, recoverable): %s", attempt, exc) + err = { + "code": "subagent_error", + "path": "render_a2ui", + "message": str(exc) or type(exc).__name__, + } + record = {"attempt": attempt, "ok": False, "errors": [err]} + attempts.append(record) + self._notify_attempt(record) + last_errors = [err] + continue + + # Balance the forwarded render_a2ui call with a tool result so the next + # turn's history is valid (an unanswered tool call is rejected on replay). + balance_call_id = render_call_id or live_render_call_id + if balance_call_id is not None: + yield AgentResponseUpdate( + role="tool", + contents=[Content.from_function_result(call_id=balance_call_id, result=RENDER_ACKNOWLEDGEMENT)], + ) + + raw_components = (render_args or {}).get("components") + components = raw_components if isinstance(raw_components, list) else [] + raw_data = (render_args or {}).get("data") + data = raw_data if isinstance(raw_data, dict) else {} + result = validate_a2ui_components(components=components, data=data, catalog=catalog) + record = {"attempt": attempt, "ok": result["valid"], "errors": result["errors"]} + attempts.append(record) + self._notify_attempt(record) + + if result["valid"] and render_args is not None: + box[0] = build_a2ui_envelope( + args=render_args, + is_update=prep.get("is_update", False), + target_surface_id=target_surface_id, + prior=prep.get("prior"), + default_surface_id=self.params["default_surface_id"], + default_catalog_id=_effective_default_catalog_id(self.params["default_catalog_id"], state), + ) + return + last_errors = result["errors"] + + box[0] = _recovery_exhausted_envelope(max_attempts, attempts) + + # -- sub-agent helpers ------------------------------------------------ + + def _subagent_messages(self, prompt: str, conversation: list[Any]) -> list[Any]: + return [Message(role="system", contents=[Content.from_text(text=prompt)]), *conversation] + + def _subagent_options(self) -> dict[str, Any]: + render_decl = FunctionTool( + name=RENDER_A2UI_TOOL_NAME, + description=_RENDER_DESCRIPTION, + func=None, + input_model=_RENDER_PARAMETERS, + ) + return { + "tools": [render_decl], + "tool_choice": {"mode": "required", "required_function_name": RENDER_A2UI_TOOL_NAME}, + } + + def _first_render_call_id(self, update: Any) -> str | None: + for content in getattr(update, "contents", None) or []: + if ( + getattr(content, "type", None) == "function_call" + and getattr(content, "name", None) == RENDER_A2UI_TOOL_NAME + and getattr(content, "call_id", None) + ): + return cast("str | None", content.call_id) + return None + + def _extract_render_args(self, updates: list[Any]) -> tuple[str | None, dict[str, Any] | None]: + """Reassemble the render_a2ui call id + arguments from streamed updates. + + MAF's ``ChatResponse.from_updates`` coalesces text content but NOT function-call + argument fragments, so a coalesced render call carries only a partial fragment. + We concatenate the render_a2ui argument deltas ourselves — the same reassembly + the AG-UI wire performs — and parse the joined JSON. Returns + ``(call_id, args)``; ``args`` is ``None`` when the sub-agent produced no render + call or the accumulated arguments are not valid JSON (→ a failed attempt the + recovery loop retries). + """ + call_id: str | None = None + named_buf = "" # name == render_a2ui fragments (opening + coalesced-final) + all_buf = "" # every function_call fragment (deltas + coalesced) + parsed_obj: dict[str, Any] | None = None + for upd in updates: + for content in getattr(upd, "contents", None) or []: + if getattr(content, "type", None) != "function_call": + continue + # The sub-agent is forced (tool_choice) to call ONLY render_a2ui, so + # every function_call fragment belongs to that one call — the + # name=="" argument-delta fragments included. + name = getattr(content, "name", None) + if call_id is None and name == RENDER_A2UI_TOOL_NAME and getattr(content, "call_id", None): + call_id = content.call_id + raw = getattr(content, "arguments", None) + if isinstance(raw, str): + all_buf += raw + if name == RENDER_A2UI_TOOL_NAME: + named_buf += raw + elif isinstance(raw, dict): + # A provider that already coalesced this fragment to a dict. + parsed_obj = raw + if parsed_obj is not None: + return call_id, parsed_obj + # Prefer the name-bearing concat (opening + coalesced full); fall back to the + # all-fragment concat for a deltas-only provider. See _first_parsable_object. + obj = _first_parsable_object(named_buf, all_buf) + if obj is None and (named_buf or all_buf): + logger.warning( + "A2UI render_a2ui arguments did not parse as a JSON object " + "(named=%d chars, all=%d chars); treating attempt as no surface.", + len(named_buf), + len(all_buf), + ) + return call_id, obj + + def _notify_attempt(self, record: dict[str, Any]) -> None: + """Invoke the host ``on_a2ui_attempt`` callback, isolating its failures. + + The callback is host-supplied observability; a throwing callback must not abort + an in-flight run (the surface has already streamed to the client by this point). + """ + callback = self.params["on_a2ui_attempt"] + if callback is None: + return + try: + callback(record) + except Exception as exc: # noqa: BLE001 — host callback must not break the run + logger.warning("A2UI on_a2ui_attempt callback raised (ignored): %s", exc) + + def _find_render_call(self, response: Any) -> Any: + for message in getattr(response, "messages", None) or []: + for content in getattr(message, "contents", None) or []: + if ( + getattr(content, "type", None) == "function_call" + and getattr(content, "name", None) == RENDER_A2UI_TOOL_NAME + ): + return content + return None + + def _read_render_args(self, response: Any) -> dict[str, Any] | None: + call = self._find_render_call(response) + return (call.parse_arguments() or {}) if call is not None else None diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_context_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_context_agent.py new file mode 100644 index 00000000000..f1cd5f40e04 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_context_agent.py @@ -0,0 +1,62 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AGUIContextAgent — surface forwarded AG-UI context to the model. + +Port of the .NET ``AGUIContextAgent``. The AG-UI hosting layer stamps the ag-ui +context slice (component catalog schema + usage guidelines) onto the run options' +``additional_properties`` (see :mod:`._state`), where the model never sees it. This +wrapper renders that slice as a system message via the toolkit's +``build_context_prompt`` and prepends it, so a plain chat agent can render A2UI +surfaces with no further setup. +""" + +from __future__ import annotations + +from typing import Any + +from ag_ui_a2ui_toolkit import build_context_prompt +from agent_framework import Content, Message, normalize_messages + +from ._state import read_agent_state + + +class AGUIContextAgent: + """Wraps an agent, prepending forwarded AG-UI context as a system message. + + Transparent pass-through: ``run`` returns whatever the inner agent's ``run`` + returns (an awaitable for non-streaming, an async iterable for streaming), so it + composes with the AG-UI hosting loop and with ``A2UIAgent`` without caring which + mode is in play. + """ + + def __init__(self, inner_agent: Any) -> None: + """Initialize the wrapper. + + Args: + inner_agent: The agent to wrap (any ``SupportsAgentRun``). + """ + self.inner_agent = inner_agent + # Mirror identity so the wrapper is indistinguishable to the hosting layer. + self.id = getattr(inner_agent, "id", None) + self.name = getattr(inner_agent, "name", None) + self.description = getattr(inner_agent, "description", None) + + def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any: + """Prepend the context system message and delegate to the inner agent.""" + return self.inner_agent.run( + self._with_context_prompt(messages, kwargs.get("options")), + stream=stream, + **kwargs, + ) + + @staticmethod + def _with_context_prompt(messages: Any, options: Any) -> list[Message]: + normalized = normalize_messages(messages) + slice_ = read_agent_state(options) + if not slice_: + return normalized + prompt = build_context_prompt({"ag-ui": slice_}) + if not prompt: + return normalized + system = Message(role="system", contents=[Content.from_text(text=prompt)]) + return [system, *normalized] diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_factory.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_factory.py new file mode 100644 index 00000000000..8457b802eb7 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_factory.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Convenience wiring for A2UI on MAF-python. + +MAF does not own a tool registry the adapter can mutate per run, and (like the .NET +sibling) auto-injection of ``render_a2ui`` happens CLIENT-side via the AG-UI +a2ui-middleware (``injectA2UITool`` puts ``render_a2ui`` into ``input.tools``). The +server-side opt-in is therefore explicit: the developer wraps their agent. This +module provides that wrapping as a one-liner. + +Progressive-paint note: surface generation must run at the AGENT level +(:class:`~._agent.A2UIAgent`) so the render sub-agent's argument fragments reach the +wire incrementally. A plain executable ``generate_a2ui`` tool added to an agent's +tool list would invoke the sub-agent hidden inside the tool body, producing the +"5s blank then bulk paint" regression — hence the wrapper, not a bare tool. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from ag_ui_a2ui_toolkit import GENERATE_A2UI_TOOL_NAME, A2UIToolParams + +from ._agent import RENDER_A2UI_TOOL_NAME, A2UIAgent +from ._context_agent import AGUIContextAgent +from ._state import read_inject_a2ui_flag + +logger = logging.getLogger(__name__) + + +def enable_a2ui( + inner_agent: Any, + subagent_chat_client: Any, + params: A2UIToolParams | None = None, +) -> AGUIContextAgent: + """Wrap ``inner_agent`` with full A2UI support. + + Composes the two wrappers in the order .NET uses: an :class:`A2UIAgent` (the + ``generate_a2ui`` streaming loop) wrapped by an :class:`AGUIContextAgent` (which + replays the forwarded component catalog into the prompt). The result is a drop-in + ``SupportsAgentRun`` to hand to ``AgentFrameworkAgent`` / + ``add_agent_framework_fastapi_endpoint``. + + Args: + inner_agent: The planner agent to wrap. + subagent_chat_client: Raw chat client for the render sub-agent. Prefer a + Chat-Completions client (``OpenAIChatCompletionClient``) so the balancing + ``render_a2ui`` tool result replays cleanly on the next turn. + params: Optional A2UI behaviour knobs (guidelines, recovery, catalog, ...). + + Returns: + The wrapped agent. + """ + return AGUIContextAgent(A2UIAgent(inner_agent, subagent_chat_client, params)) + + +def plan_a2ui_injection( + *, + agent: Any, + forwarded_props: Any, + existing_tool_names: list[str], + config: dict[str, Any] | None = None, + log: logging.Logger | None = None, +) -> dict[str, Any] | None: + """Decide whether to auto-inject A2UI for this run (Strands-parity contract). + + CopilotKit's runtime composes ``forwardedProps``; the AG-UI a2ui-middleware sets + ``injectA2UITool`` there. When true (or a string naming the injected render tool to + drop), the framework auto-wraps the agent with :func:`enable_a2ui`; when false/unset + the developer wires A2UI themselves. Rules: + + 1. Off unless ``forwardedProps.injectA2UITool`` is truthy OR a backend + ``config["inject_a2ui_tool"]`` override is — nullish fallback, so an explicit + runtime ``false`` beats a backend opt-in. + 2. USER PREVAILS — never double-inject when ``generate_a2ui`` is already wired, or + when the agent is already A2UI-wrapped. + 3. The render sub-agent client is inferred from the planner agent's ``.client``; if + none can be inferred, warn and skip (wire ``enable_a2ui`` explicitly instead). + + Returns ``{"agent", "drop_tool_names"}`` (the wrapped agent + the middleware-injected + render tool to strip from the planner's tool list) or ``None`` when no injection. + """ + log = log or logger + config = config or {} + + flag: Any = read_inject_a2ui_flag(forwarded_props) + if flag is None: + flag = config.get("inject_a2ui_tool") + if not flag: + return None + + # USER PREVAILS: explicit dev wiring or an already-wrapped agent wins. + if GENERATE_A2UI_TOOL_NAME in existing_tool_names: + return None + if isinstance(agent, (A2UIAgent, AGUIContextAgent)): + return None + + subagent_client = getattr(agent, "client", None) + if subagent_client is None: + log.warning( + "injectA2UITool is set but no chat client could be inferred from the agent; " + "skipping A2UI auto-injection. Wire enable_a2ui() explicitly." + ) + return None + + render_tool_name = flag if isinstance(flag, str) else RENDER_A2UI_TOOL_NAME + params: A2UIToolParams = { + "catalog": config.get("catalog"), + "guidelines": config.get("guidelines"), + "recovery": config.get("recovery"), + "default_catalog_id": config.get("default_catalog_id"), + "default_surface_id": config.get("default_surface_id"), + } + return { + "agent": enable_a2ui(agent, subagent_client, params), + "drop_tool_names": [render_tool_name], + } diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_state.py new file mode 100644 index 00000000000..1f2a2e16e2b --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_state.py @@ -0,0 +1,195 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AG-UI context plumbing for A2UI. + +The AG-UI A2UI middleware injects the component catalog and usage guidelines as +``RunAgentInput.context`` entries, and the ``injectA2UITool`` flag via +``forwardedProps``. MAF's AG-UI hosting (``run_agent_stream``) does not forward +``context`` to the agent by default, so this module builds the ``ag-ui`` state +slice the toolkit expects and stamps it onto the run's +``options.additional_properties`` under :data:`A2UI_CONTEXT_KEY` — the analog of +.NET ``MapAGUI`` stamping ``ag_ui_context`` onto ``ChatOptions.AdditionalProperties``. + +Wrapper agents (:class:`~.._context_agent.AGUIContextAgent`, ``A2UIAgent``) read it +back with :func:`read_agent_state` and feed it to the toolkit's +``build_context_prompt`` / ``prepare_a2ui_request``. +""" + +from __future__ import annotations + +import json +from typing import Any + +# Additional-properties key under which the ag-ui context slice is stamped onto the +# run options. Mirrors .NET's "ag_ui_context" AdditionalProperties key. +A2UI_CONTEXT_KEY = "ag_ui_context" + +# MUST stay byte-identical to the A2UI middleware's exported +# A2UI_SCHEMA_CONTEXT_DESCRIPTION (middlewares/a2ui-middleware/src/index.ts) and to +# the langgraph-python adapter's copy. The match below is exact-equality; any drift +# silently routes the schema into the generic context section instead of +# ``a2ui_schema``, defeating the catalog-aware validation path. +A2UI_SCHEMA_CONTEXT_DESCRIPTION = ( + "A2UI Component Schema — available components for generating UI surfaces. " + "Use these component names and properties when creating A2UI operations." +) + + +def _entry_desc_value(entry: Any) -> tuple[str, Any]: + """Read (description, value) from a context entry that may be a dict or object.""" + if isinstance(entry, dict): + return entry.get("description", "") or "", entry.get("value") + return getattr(entry, "description", "") or "", getattr(entry, "value", None) + + +def build_ag_ui_context_slice(context: list[Any] | None) -> dict[str, Any]: + """Build the ``ag-ui`` context slice from AG-UI ``context`` entries. + + Splits the A2UI schema context entry (matched by exact description) into + ``a2ui_schema`` and routes the remaining entries to ``context``, mirroring the + langgraph-python adapter. This slice is catalog/guidelines ONLY — it is what the + toolkit's ``build_context_prompt`` consumes. + + The auto-inject ENABLEMENT flag is deliberately NOT part of this slice: enablement + is sourced from ``forwardedProps`` (see :func:`read_inject_a2ui_flag`), a separate + concern from the catalog context. Returns an empty dict when there is no A2UI + context, so callers can skip stamping for non-A2UI runs. + """ + schema_value: Any = None + regular_context: list[Any] = [] + for entry in context or []: + desc, value = _entry_desc_value(entry) + if desc == A2UI_SCHEMA_CONTEXT_DESCRIPTION: + schema_value = value + else: + regular_context.append(entry) + + slice_: dict[str, Any] = {} + if regular_context: + slice_["context"] = regular_context + if schema_value is not None: + slice_["a2ui_schema"] = schema_value + return slice_ + + +def read_inject_a2ui_flag(forwarded_props: dict[str, Any] | None) -> Any: + """Read the A2UI auto-inject enablement flag from ``forwardedProps``. + + Enablement comes from ``forwardedProps.injectA2UITool`` (set by the AG-UI + a2ui-middleware), NOT from ``context``. Returns the RAW value — ``True``/``False``, + or a string naming the injected render tool to drop (Strands-parity) — or ``None`` + when unset so a backend opt-in can take over with nullish fallback. Callers gate on + truthiness; auto-injection preserves a string value as the render-tool name. MAF + does not snake-mangle ``forwardedProps`` keys (unlike langgraph), so the camelCase + form is canonical; the snake form is accepted for safety. + """ + forwarded = forwarded_props if isinstance(forwarded_props, dict) else {} + if "injectA2UITool" in forwarded: + return forwarded["injectA2UITool"] + if "inject_a2ui_tool" in forwarded: + return forwarded["inject_a2ui_tool"] + return None + + +def _role_value(msg: Any) -> str | None: + role = getattr(msg, "role", None) + return getattr(role, "value", role) if role is not None else None + + +def to_history_messages(messages: list[Any]) -> list[dict[str, Any]]: + """Map MAF ``Message`` objects onto the toolkit's history shape. + + The toolkit's ``find_prior_surface`` / ``prepare_a2ui_request`` walk a list of + ``{"role", "content"}`` entries, where a tool message's content is the JSON + string of its function-result payload (the A2UI operations envelope for prior + renders). Mirrors .NET ``ToHistoryMessage``. Duck-typed (no agent_framework + import) so this module stays dependency-light. + """ + out: list[dict[str, Any]] = [] + for msg in messages: + role = _role_value(msg) + content: Any = getattr(msg, "text", None) + if not content and role == "tool": + # A tool message carries its payload as function_result content; the + # prior-surface walker needs the raw JSON string. + for c in getattr(msg, "contents", None) or []: + if getattr(c, "type", None) == "function_result": + result = getattr(c, "result", None) + content = result if isinstance(result, str) else _safe_json(result) + if content: + break + out.append({"role": role, "content": content}) + return out + + +def _safe_json(value: Any) -> str: + try: + return json.dumps(value, default=str) + except (TypeError, ValueError): + return str(value) + + +def _additional_properties(options: Any) -> dict[str, Any] | None: + """Read ``additional_properties`` from a ChatOptions object or a plain dict.""" + if options is None: + return None + if isinstance(options, dict): + ap = options.get("additional_properties") + return ap if isinstance(ap, dict) else None + ap = getattr(options, "additional_properties", None) + return ap if isinstance(ap, dict) else None + + +def read_agent_state(options: Any) -> dict[str, Any]: + """Read the ag-ui context slice stamped on the run options. + + Accepts a ChatOptions object or a plain options dict (the shape + ``run_agent_stream`` passes). Returns the slice, or an empty dict when absent. + """ + ap = _additional_properties(options) + if not ap: + return {} + slice_ = ap.get(A2UI_CONTEXT_KEY) + return slice_ if isinstance(slice_, dict) else {} + + +def stamp_context_slice(options: Any, slice_: dict[str, Any]) -> dict[str, Any]: + """Stamp the ag-ui context slice onto run options' additional_properties. + + Normalizes ``options`` to a dict (creating one if needed) and returns it, so the + caller can assign it back into ``run_kwargs["options"]``. No-op-safe: an empty + slice still returns valid options. + """ + opts: dict[str, Any] = dict(options) if isinstance(options, dict) else {} + ap = opts.get("additional_properties") + ap = dict(ap) if isinstance(ap, dict) else {} + ap[A2UI_CONTEXT_KEY] = slice_ + opts["additional_properties"] = ap + return opts + + +def strip_context_slice(options: Any) -> Any: + """Return ``options`` with the ag-ui context slice removed from + ``additional_properties``. + + The slice is adapter-private run state; ``additional_properties`` is forwarded to + the provider SDK by the underlying chat client, which rejects the unknown key. + Wrapper agents therefore read the slice with :func:`read_agent_state` and then + strip it via this helper BEFORE delegating to the inner planner agent, so it never + reaches ``responses.create`` / ``chat.completions.create``. If removing the slice + empties ``additional_properties``, the key is dropped entirely. Non-dict options are + returned unchanged (best effort — only the dict shape the hosting loop passes carries + the stamped slice). + """ + if not isinstance(options, dict): + return options + ap = options.get("additional_properties") + if not isinstance(ap, dict) or A2UI_CONTEXT_KEY not in ap: + return options + opts = dict(options) + new_ap = {k: v for k, v in ap.items() if k != A2UI_CONTEXT_KEY} + if new_ap: + opts["additional_properties"] = new_ap + else: + opts.pop("additional_properties", None) + return opts diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 09766b7e1aa..e92e38e3d09 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -25,6 +25,7 @@ def __init__( use_service_session: bool = False, require_confirmation: bool = True, snapshot_store: AGUIThreadSnapshotStore | None = None, + a2ui_config: dict[str, Any] | None = None, ): """Initialize agent configuration. @@ -38,12 +39,16 @@ def __init__( emitted; tools remain gated server-side. snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless endpoint setup also provides an explicit Snapshot Scope resolver. + require_confirmation: Whether predictive updates require user confirmation before applying + a2ui_config: Optional backend A2UI config consumed by auto-injection + (``forwardedProps.injectA2UITool``). See ``plan_a2ui_injection``. """ self.state_schema = self._normalize_state_schema(state_schema) self.predict_state_config = predict_state_config or {} self.use_service_session = use_service_session self.require_confirmation = require_confirmation self.snapshot_store = snapshot_store + self.a2ui_config = a2ui_config @staticmethod def _normalize_state_schema(state_schema: Any | None) -> dict[str, Any]: @@ -90,6 +95,7 @@ def __init__( require_confirmation: bool = True, use_service_session: bool = False, snapshot_store: AGUIThreadSnapshotStore | None = None, + a2ui_config: dict[str, Any] | None = None, ): """Initialize the AG-UI compatible agent wrapper. @@ -106,6 +112,7 @@ def __init__( use_service_session: Whether the agent session is service-managed snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless endpoint setup also provides an explicit Snapshot Scope resolver. + a2ui_config: Optional backend A2UI config consumed by auto-injection. """ self.agent = agent self.name = name or getattr(agent, "name", "agent") @@ -117,6 +124,7 @@ def __init__( use_service_session=use_service_session, require_confirmation=require_confirmation, snapshot_store=snapshot_store, + a2ui_config=a2ui_config, ) # Server-side Approval State. Populated when approval requests are emitted diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index ab074ae15cd..5e32658f542 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -49,6 +49,7 @@ from agent_framework._types import ResponseStream from agent_framework.exceptions import AgentInvalidResponseException +from ._a2ui._state import build_ag_ui_context_slice, read_inject_a2ui_flag, stamp_context_slice from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id from ._message_adapters import normalize_agui_input_messages from ._orchestration._predictive_state import PredictiveStateHandler @@ -2017,6 +2018,41 @@ async def run_agent_stream( register_additional_client_tools(agent, client_tools) tools = merge_tools(server_tools, client_tools) + # A2UI auto-injection: CopilotKit's runtime composes forwardedProps; the AG-UI + # a2ui-middleware sets injectA2UITool there. When set, auto-wrap the agent with A2UI + # surface generation (inferring the render sub-agent from the planner's chat client) + # and strip the middleware-injected render tool from the planner's list. When unset, + # the developer wires A2UI themselves. A backend config["inject_a2ui_tool"] opt-in + # enables A2UI without the runtime flag (nullish fallback: an explicit runtime false + # still wins). plan_a2ui_injection is imported lazily so this hosting path stays + # importable without the optional ag-ui-a2ui-toolkit. + _forwarded = input_data.get("forwarded_props") or input_data.get("forwardedProps") + _a2ui_config = getattr(config, "a2ui_config", None) + _a2ui_flag = read_inject_a2ui_flag(_forwarded) + if _a2ui_flag is None and _a2ui_config: + _a2ui_flag = _a2ui_config.get("inject_a2ui_tool") + if _a2ui_flag: + 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] + plan = plan_a2ui_injection( + agent=agent, + forwarded_props=_forwarded, + existing_tool_names=existing_tool_names, + config=_a2ui_config, + ) + if plan is not None: + agent = plan["agent"] + drop = set(plan["drop_tool_names"]) + if tools: + tools = [t for t in tools if getattr(t, "name", None) not in drop] + except ImportError as exc: + logger.warning( + "injectA2UITool is set but A2UI support is unavailable (install the [a2ui] extra): %s", + exc, + ) + # Create session (with service session support) if config.use_service_session: supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") @@ -2061,6 +2097,17 @@ async def run_agent_stream( if safe_metadata: run_kwargs["options"] = {"metadata": safe_metadata, "store": True} + # Forward the AG-UI context (A2UI component catalog + guidelines) onto the run + # options' additional_properties. MAF otherwise drops input_data["context"]; A2UI + # wrapper agents read this slice back via ag_ui_a2ui_toolkit-free helpers (no + # toolkit import here). Mirrors .NET MapAGUI stamping ag_ui_context onto + # ChatOptions.AdditionalProperties. Only stamped when there is A2UI context, so + # non-A2UI runs are unaffected. The auto-inject ENABLEMENT flag is sourced + # 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) + # Resolve approval responses (execute approved tools, replace approvals with results) # This must happen before running the agent so it sees the tool results tools_for_execution = tools if tools is not None else server_tools @@ -2389,7 +2436,17 @@ async def run_agent_stream( last_result = flow.tool_results[-1] last_call_id = last_result.get("toolCallId") last_tool_name = flow.get_tool_name(last_call_id) - if not _should_suppress_intermediate_snapshot( + # A2UI surfaces stream as activities in emission order (tool card -> surface -> + # narration). A terminal MessagesSnapshotEvent makes the client re-render from + # the reconciled message list, which drops that order — the injected + # generate_a2ui tool card re-positions BELOW the surface and text. Other AG-UI + # frameworks emit no terminal snapshot here, so skip it for A2UI runs; the next + # turn's history is still reconstructable from the streamed events. + tool_names = {(tc.get("function") or {}).get("name") for tc in flow.tool_calls_by_id.values()} + a2ui_run = "generate_a2ui" in tool_names or "render_a2ui" in tool_names + if a2ui_run: + logger.info("Suppressing terminal MessagesSnapshotEvent for A2UI run to preserve streamed message order.") + if not a2ui_run and not _should_suppress_intermediate_snapshot( last_tool_name, predict_state_config, config.require_confirmation ): yield snapshot_event diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py index 5486802d647..3c7e4956b00 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -91,6 +91,7 @@ def add_agent_framework_fastapi_endpoint( snapshot_store: AGUIThreadSnapshotStore | None = None, snapshot_scope_resolver: SnapshotScopeResolver | None = None, keepalive_seconds: float | None = 15, + a2ui_config: dict[str, Any] | None = None, ) -> None: """Add an AG-UI endpoint to a FastAPI app. @@ -116,6 +117,10 @@ def add_agent_framework_fastapi_endpoint( keepalive_seconds: Endpoint SSE keepalive interval in seconds. Defaults to 15. Positive values emit fixed SSE comments while the stream is open. None disables keepalive and preserves the non-keepalive response path. Keepalive comments are transport traffic and do not change AG-UI events. + a2ui_config: Optional backend A2UI config used when the runtime auto-injects + the surface-generation tool (``forwardedProps.injectA2UITool``). Keys: + ``inject_a2ui_tool`` (backend opt-in override), ``default_catalog_id``, + ``catalog``, ``guidelines``, ``recovery``, ``default_surface_id``. """ _validate_keepalive_seconds(keepalive_seconds) @@ -132,6 +137,7 @@ def add_agent_framework_fastapi_endpoint( state_schema=state_schema, predict_state_config=predict_state_config, snapshot_store=snapshot_store, + a2ui_config=a2ui_config, ) else: raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.") diff --git a/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py index 767dddca9cb..c2dd12984be 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py +++ b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py @@ -2,6 +2,13 @@ """Example agents for AG-UI demonstration.""" +from .a2ui_agents import ( + A2UI_DEMO_CONFIG, + a2ui_advanced_agent, + a2ui_dynamic_schema_agent, + a2ui_fixed_schema_agent, + a2ui_recovery_agent, +) from .document_writer_agent import document_writer_agent from .human_in_the_loop_agent import human_in_the_loop_agent from .recipe_agent import recipe_agent @@ -14,6 +21,11 @@ from .weather_agent import weather_agent __all__ = [ + "A2UI_DEMO_CONFIG", + "a2ui_advanced_agent", + "a2ui_dynamic_schema_agent", + "a2ui_fixed_schema_agent", + "a2ui_recovery_agent", "document_writer_agent", "human_in_the_loop_agent", "recipe_agent", diff --git a/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/a2ui_agents.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/a2ui_agents.py new file mode 100644 index 00000000000..4e52741cb05 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/a2ui_agents.py @@ -0,0 +1,261 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""A2UI (agent-generated UI) example agents. + +Four dojo demos mirroring the .NET / AWS Strands / LangGraph A2UI samples: + +* ``a2ui_dynamic_schema_agent`` — happy path: a valid surface is generated from the + frontend-supplied component catalog, with a backend composition guide. +* ``a2ui_advanced_agent`` — ZERO-CONFIG: no backend catalog or composition guide. The + catalog schema and ``injectA2UITool`` arrive on ``forwardedProps`` from the frontend; + the adapter auto-injects ``generate_a2ui`` and the render sub-agent composes against + the forwarded catalog. Proves the easy-devex path (a plain agent + a client catalog). +* ``a2ui_recovery_agent`` — exercises the validate→retry recovery loop when the render + sub-agent first emits an invalid surface; deterministic under aimock fixtures. +* ``a2ui_fixed_schema_agent`` — direct-tool (NO subagent/recovery): plain backend tools + (``search_flights`` / ``search_hotels``) return a pre-authored ``a2ui_operations`` + envelope as a JSON STRING; the runtime A2UI middleware detects it in the tool result + and paints. Only the DATA changes per call. Works even when the runtime forwards + ``injectA2UITool`` because + ``A2UIAgent`` passes non-``generate_a2ui`` tool calls (and their results) straight + through to the inner function-invoker, so the fixed tool's envelope still paints. + +The first three are PLAIN agents with NO a2ui wiring. The dojo's copilotkit route sends +``forwardedProps.injectA2UITool``; the adapter auto-injects ``generate_a2ui`` and infers +the render sub-agent from the agent's chat client. Pass :data:`A2UI_DEMO_CONFIG` as the +endpoint's ``a2ui_config`` so auto-injected surfaces reference the dojo catalog (the +zero-config ``a2ui_advanced`` demo passes no config). + +Use a Chat-Completions client (``OpenAIChatCompletionClient``) for these agents: it +streams ``render_a2ui`` argument deltas (progressive paint) and replays the balancing +tool result cleanly, where the Responses-API client (``OpenAIChatClient``) +buffers/rejects. +""" + +from typing import Any + +from agent_framework import Agent, tool + +# The dojo registers its dynamic component catalog (HotelCard, ProductCard, +# TeamMemberCard, StarRating, Row) under this id; auto-injected surfaces must +# reference it so the renderer can resolve their components. +DOJO_CATALOG_ID = "https://a2ui.org/demos/dojo/dynamic_catalog.json" + +# Teaches the render sub-agent how to compose the dojo catalog's components. +# Mirrors the LangGraph / Strands dynamic-schema COMPOSITION_GUIDE so a real +# model (not just the aimock fixtures) can produce valid surfaces. +COMPOSITION_GUIDE = """ +## Available Pre-made Components + +You have card components plus a Row container. Use Row as the root with structural +children to repeat a card per item. + +### Row +Layout container. Repeat a card template via structural children: + {"id":"root","component":"Row","children":{"componentId":"card","path":"/items"}} + +### HotelCard +Props: name, location, rating (number 0-5), pricePerNight, action + +### ProductCard +Props: name, price, rating (number 0-5), description (optional), action + +### TeamMemberCard +Props: name, role, department (optional), email (optional), action + +## RULES +- Root is ALWAYS a Row with structural children: {"componentId":"","path":"/items"} +- ALWAYS include the referenced card component in the components array. +- Inside templates use RELATIVE paths (no leading slash): {"path":"name"}. +- Always provide data in the "data" argument as {"items":[...]}. +- Pick the card type that best matches the request; generate 3-4 realistic items. +""" + +#: Backend A2UI config for the dojo demos, passed as ``a2ui_config`` to +#: ``add_agent_framework_fastapi_endpoint``. Consumed by auto-injection. +A2UI_DEMO_CONFIG: dict[str, Any] = { + "default_catalog_id": DOJO_CATALOG_ID, + "guidelines": {"composition_guide": COMPOSITION_GUIDE}, +} + +_SYSTEM_PROMPT = ( + "You are a helpful assistant that creates rich visual UI on the fly. When the user " + "asks for visual content (product comparisons, dashboards, team rosters, lists, " + "cards, etc.), use the generate_a2ui tool to create a dynamic A2UI surface. " + "IMPORTANT: after calling the tool, do NOT repeat the data in your text response — " + "the tool renders UI automatically. Just confirm what was rendered." +) + + +def a2ui_dynamic_schema_agent(client: Any) -> Agent[Any]: + """Plain agent for the A2UI dynamic-schema demo (auto-injected generate_a2ui).""" + return Agent[Any](name="a2ui_dynamic_schema", instructions=_SYSTEM_PROMPT, client=client) + + +def a2ui_advanced_agent(client: Any) -> Agent[Any]: + """Plain agent for the ZERO-CONFIG A2UI demo. + + Identical to the dynamic-schema agent but its endpoint passes NO ``a2ui_config``: + the catalog schema arrives on ``forwardedProps`` and the render sub-agent composes + against it with only the toolkit's built-in generation/design guidelines. This is + the "a plain agent plus a client catalog is enough" path. + """ + return Agent[Any](name="a2ui_advanced", instructions=_SYSTEM_PROMPT, client=client) + + +def a2ui_recovery_agent(client: Any) -> Agent[Any]: + """Plain agent for the A2UI recovery demo (auto-injected generate_a2ui).""" + return Agent[Any](name="a2ui_recovery", instructions=_SYSTEM_PROMPT, client=client) + + +# --------------------------------------------------------------------------- # +# Fixed-schema (direct-tool) demo +# --------------------------------------------------------------------------- # + +# Custom fixed catalog id the dojo's a2ui_fixed_schema page registers (StarRating + +# HotelCard). The component TREE is authored ahead of time here; only the DATA +# (the hotel list) changes per call. Mirrors the LangGraph fixed-schema demo. +FIXED_CATALOG_ID = "https://a2ui.org/demos/dojo/fixed_catalog.json" +FLIGHT_SURFACE_ID = "flight-search-results" +HOTEL_SURFACE_ID = "hotel-search-results" + +# Pre-authored A2UI v0.9 component array: a Row repeating a FlightCard per /flights item. +FLIGHT_SCHEMA: list[dict[str, Any]] = [ + { + "id": "root", + "component": "Row", + "children": {"componentId": "flight-card", "path": "/flights"}, + "gap": 16, + }, + { + "id": "flight-card", + "component": "FlightCard", + "airline": {"path": "airline"}, + "airlineLogo": {"path": "airlineLogo"}, + "flightNumber": {"path": "flightNumber"}, + "origin": {"path": "origin"}, + "destination": {"path": "destination"}, + "date": {"path": "date"}, + "departureTime": {"path": "departureTime"}, + "arrivalTime": {"path": "arrivalTime"}, + "duration": {"path": "duration"}, + "status": {"path": "status"}, + "price": {"path": "price"}, + "action": { + "event": { + "name": "book_flight", + "context": { + "flightNumber": {"path": "flightNumber"}, + "origin": {"path": "origin"}, + "destination": {"path": "destination"}, + "price": {"path": "price"}, + }, + } + }, + }, +] + +# Pre-authored A2UI v0.9 component array: a Row repeating a HotelCard per /hotels item. +HOTEL_SCHEMA: list[dict[str, Any]] = [ + { + "id": "root", + "component": "Row", + "children": {"componentId": "hotel-card", "path": "/hotels"}, + "gap": 16, + }, + { + "id": "hotel-card", + "component": "HotelCard", + "name": {"path": "name"}, + "location": {"path": "location"}, + "rating": {"path": "rating"}, + "pricePerNight": {"path": "price"}, + "action": { + "event": { + "name": "book_hotel", + "context": {"hotelName": {"path": "name"}, "price": {"path": "price"}}, + } + }, + }, +] + +_FIXED_SYSTEM_PROMPT = ( + "You are a helpful travel assistant. When the user asks about flights, call the " + "search_flights tool; when they ask about hotels, call the search_hotels tool. " + "Provide 3-4 realistic results. " + "IMPORTANT: after calling the tool, do NOT repeat the data in your text response — " + "the tool renders rich UI automatically. Just say something brief like 'Here are " + "your results.'" +) + + +@tool +def search_flights(flights: list[dict[str, Any]]) -> str: + """Search for flights and display the results as rich cards. + + Each flight must have: id, airline, airlineLogo (a logo URL), flightNumber, + origin, destination, date, departureTime, arrivalTime, duration, status, and + price (e.g. "$289"). Generate 3-4 realistic flight results. + + Returns: + The A2UI operations envelope (JSON string) the runtime middleware paints. + """ + from ag_ui_a2ui_toolkit import ( + create_surface, + update_components, + update_data_model, + wrap_as_operations_envelope, + ) + + return wrap_as_operations_envelope( + [ + create_surface(FLIGHT_SURFACE_ID, FIXED_CATALOG_ID), + update_components(FLIGHT_SURFACE_ID, FLIGHT_SCHEMA), + update_data_model(FLIGHT_SURFACE_ID, {"flights": flights}), + ] + ) + + +@tool +def search_hotels(hotels: list[dict[str, Any]]) -> str: + """Search for hotels and display the results as rich cards with star ratings. + + Each hotel must have: id, name (e.g. "The Plaza"), location (e.g. "Midtown + Manhattan, NYC"), rating (float 0-5, e.g. 4.5), and price (per night, e.g. "$350"). + Generate 3-4 realistic hotel results. + + Returns: + The A2UI operations envelope (JSON string) the runtime middleware paints. + """ + # Import the toolkit lazily so the examples package imports without the optional + # ag-ui-a2ui-toolkit dependency unless the fixed-schema demo is actually used. + from ag_ui_a2ui_toolkit import ( + create_surface, + update_components, + update_data_model, + wrap_as_operations_envelope, + ) + + return wrap_as_operations_envelope( + [ + create_surface(HOTEL_SURFACE_ID, FIXED_CATALOG_ID), + update_components(HOTEL_SURFACE_ID, HOTEL_SCHEMA), + update_data_model(HOTEL_SURFACE_ID, {"hotels": hotels}), + ] + ) + + +def a2ui_fixed_schema_agent(client: Any) -> Agent[Any]: + """Plain agent for the fixed-schema (direct-tool) A2UI demo. + + Wires its OWN ``search_flights`` / ``search_hotels`` backend tools whose results are + complete ``a2ui_operations`` envelopes — no ``generate_a2ui``, no render sub-agent, + no recovery loop. The pre-authored :data:`FLIGHT_SCHEMA` / :data:`HOTEL_SCHEMA` bind + to the frontend's fixed catalog; only the data varies per call. + """ + return Agent[Any]( + name="a2ui_fixed_schema", + instructions=_FIXED_SYSTEM_PROMPT, + client=client, + tools=[search_flights, search_hotels], + ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py index 6425f00a29a..82526f2ed0e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py +++ b/python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py @@ -6,16 +6,25 @@ import logging import os +from pathlib import Path from typing import Any, cast import uvicorn from agent_framework import ChatOptions from agent_framework._clients import SupportsChatGetResponse from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint +from agent_framework.azure import AzureOpenAIChatClient from agent_framework.openai import OpenAIChatCompletionClient from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware +from ..agents.a2ui_agents import ( + A2UI_DEMO_CONFIG, + a2ui_advanced_agent, + a2ui_dynamic_schema_agent, + a2ui_fixed_schema_agent, + a2ui_recovery_agent, +) from ..agents.document_writer_agent import document_writer_agent from ..agents.human_in_the_loop_agent import human_in_the_loop_agent from ..agents.recipe_agent import recipe_agent @@ -74,15 +83,36 @@ allow_headers=["*"], ) +# Load the examples .env (OPENAI_API_KEY etc.) before constructing any chat client. +# python-dotenv is an examples-only convenience; if it isn't installed, fall back to +# whatever the shell already exported. override=True so a stale/empty exported +# OPENAI_API_KEY doesn't shadow the .env value. +try: + from dotenv import load_dotenv +except ImportError: + load_dotenv = None # type: ignore[assignment] +if load_dotenv is not None: + load_dotenv(Path(__file__).resolve().parent.parent / ".env", override=True) + # Create a shared chat client for all agents # You can use different chat clients for different agents if needed -# Set CHAT_CLIENT=anthropic to use Anthropic, defaults to Azure OpenAI -client: SupportsChatGetResponse[ChatOptions] = cast( - SupportsChatGetResponse[ChatOptions], - AnthropicClient() - if AnthropicClient is not None and os.getenv("CHAT_CLIENT", "").lower() == "anthropic" - else OpenAIChatCompletionClient(), -) +# Set CHAT_CLIENT=anthropic for Anthropic, CHAT_CLIENT=openai (or just an +# OPENAI_API_KEY) for OpenAI Chat-Completions; otherwise defaults to Azure OpenAI. +_chat_client_kind = os.getenv("CHAT_CLIENT", "").lower() + + +def _default_shared_client() -> SupportsChatGetResponse[ChatOptions]: + if AnthropicClient is not None and _chat_client_kind == "anthropic": + return cast(SupportsChatGetResponse[ChatOptions], AnthropicClient()) + if _chat_client_kind == "openai" or (not _chat_client_kind and os.getenv("OPENAI_API_KEY")): + return cast( + SupportsChatGetResponse[ChatOptions], + OpenAIChatCompletionClient(model=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o")), + ) + return cast(SupportsChatGetResponse[ChatOptions], AzureOpenAIChatClient()) + + +client: SupportsChatGetResponse[ChatOptions] = _default_shared_client() # Agentic Chat - basic chat agent add_agent_framework_fastapi_endpoint( @@ -150,6 +180,54 @@ path="/deterministic_state", ) +# --- A2UI (agent-generated UI) demos --------------------------------------- +# A2UI surface streaming needs a Chat-Completions client: it emits render_a2ui +# argument deltas per chunk (progressive paint) and replays the balancing tool +# result cleanly, where the Responses/Azure path buffers. Use a dedicated +# OpenAIChatCompletionClient for these endpoints when OPENAI_API_KEY is set; otherwise fall +# back to the shared client with a warning (streaming may not paint incrementally). +if os.getenv("OPENAI_API_KEY"): + a2ui_client: SupportsChatGetResponse[ChatOptions] = cast( + SupportsChatGetResponse[ChatOptions], + OpenAIChatCompletionClient(model=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o")), + ) +else: + logging.getLogger(__name__).warning( + "OPENAI_API_KEY not set; A2UI demos fall back to the shared client and may not " + "stream render_a2ui deltas incrementally. Set it in the examples .env." + ) + a2ui_client = client + +# Dynamic schema — subagent generates a surface against the dojo catalog. +add_agent_framework_fastapi_endpoint( + app=app, + agent=a2ui_dynamic_schema_agent(a2ui_client), + path="/a2ui_dynamic_schema", + a2ui_config=A2UI_DEMO_CONFIG, +) + +# Advanced — ZERO-CONFIG: no backend catalog/guide; catalog arrives on forwardedProps. +add_agent_framework_fastapi_endpoint( + app=app, + agent=a2ui_advanced_agent(a2ui_client), + path="/a2ui_advanced", +) + +# Recovery — validate->retry loop; structural validation drives regeneration. +add_agent_framework_fastapi_endpoint( + app=app, + agent=a2ui_recovery_agent(a2ui_client), + path="/a2ui_recovery", + a2ui_config=A2UI_DEMO_CONFIG, +) + +# Fixed schema — direct backend tool returns a pre-authored a2ui_operations envelope. +add_agent_framework_fastapi_endpoint( + app=app, + agent=a2ui_fixed_schema_agent(a2ui_client), + path="/a2ui_fixed_schema", +) + def main(): """Run the server.""" diff --git a/python/packages/ag-ui/pyproject.toml b/python/packages/ag-ui/pyproject.toml index b943b92dab6..c4e331e12ec 100644 --- a/python/packages/ag-ui/pyproject.toml +++ b/python/packages/ag-ui/pyproject.toml @@ -34,6 +34,13 @@ dev = [ "pytest==9.1.1", "httpx==0.28.1", ] +# A2UI (agent-generated UI) support. Optional: the agent_framework_ag_ui._a2ui +# modules import this lazily, so the base package works without it. +# Floor 0.0.4 is the first release with singular-`child` validation and +# `child_cycle` detection (ag-ui-protocol/ag-ui#1944). +a2ui = [ + "ag-ui-a2ui-toolkit>=0.0.4", +] [dependency-groups] test = [ @@ -69,6 +76,12 @@ warn_return_any = true warn_unused_configs = true disallow_untyped_defs = false +# ag-ui-a2ui-toolkit is a runtime-only dep that does not yet ship a py.typed marker. +[[tool.mypy.overrides]] +module = ["ag_ui_a2ui_toolkit", "ag_ui_a2ui_toolkit.*"] +ignore_missing_imports = true +follow_untyped_imports = true + [tool.pyright] include = ["agent_framework_ag_ui"] exclude = ["tests", "tests/ag_ui", "examples"] diff --git a/python/packages/ag-ui/tests/ag_ui/test_a2ui.py b/python/packages/ag-ui/tests/ag_ui/test_a2ui.py new file mode 100644 index 00000000000..cdf81585842 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_a2ui.py @@ -0,0 +1,852 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for A2UI (agent-generated UI) support in the AG-UI adapter. + +Covers the context plumbing (_state), the AGUIContextAgent wrapper, and the +A2UIAgent streaming/non-streaming loop including recovery, catalog-driven +validation, error classification, and mid-stream render-call balancing. + +Skipped wholesale when ag-ui-a2ui-toolkit is not installed (A2UI is an optional +extra; the base package does not require it). +""" + +import asyncio +import json + +import pytest + +pytest.importorskip("ag_ui_a2ui_toolkit") + +from agent_framework import AgentResponseUpdate, ChatResponse, ChatResponseUpdate, Content, Message # noqa: E402 + +from agent_framework_ag_ui._a2ui import ( # noqa: E402 + A2UI_SCHEMA_CONTEXT_DESCRIPTION, + A2UIAgent, + AGUIContextAgent, + build_ag_ui_context_slice, + enable_a2ui, + plan_a2ui_injection, + read_agent_state, + read_inject_a2ui_flag, + stamp_context_slice, +) +from agent_framework_ag_ui._a2ui._state import to_history_messages # noqa: E402 + +# --------------------------------------------------------------------------- # +# Test doubles +# --------------------------------------------------------------------------- # + + +def _role(msg): + r = getattr(msg, "role", None) + return getattr(r, "value", r) + + +class _GenerateOnceInner: + """Planner that requests one generate_a2ui surface, then narrates.""" + + id = name = description = "planner" + + def __init__(self): + self.calls = 0 + self.last_tools = None + + def run(self, messages, *, stream=False, session=None, tools=None, **kwargs): + self.calls += 1 + self.last_tools = [getattr(t, "name", None) for t in (tools or [])] + n = self.calls + + async def gen(): + if n == 1: + yield AgentResponseUpdate( + role="assistant", + contents=[ + Content.from_function_call( + call_id="g1", name="generate_a2ui", arguments=json.dumps({"intent": "create"}) + ) + ], + ) + else: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="Done.")]) + + return gen() + + +class _RenderSub: + """Render sub-agent that streams render_a2ui args as fragments.""" + + def __init__(self, components=None, fragments=None): + self._components = components or [{"id": "root", "component": "Card"}] + self._fragments = fragments + + def get_response(self, messages, *, stream=False, options=None): + full = {"surfaceId": "s1", "components": self._components} + text = json.dumps(full) + frags = self._fragments if self._fragments is not None else [text[:8], text[8:20], text[20:]] + + async def gen(): + for frag in frags: + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="r1", name="render_a2ui", arguments=frag)], + ) + + return gen() + + +async def _drive(agent, options=None): + """Run a streaming A2UIAgent and classify the yielded content.""" + kinds = [] + async for update in agent.run("make a card", stream=True, options=options): + for c in update.contents: + t = getattr(c, "type", None) + if t == "function_call": + kinds.append(("call", c.name, c.arguments)) + elif t == "function_result": + kinds.append(("result", getattr(c, "call_id", None), c.result)) + elif t == "text": + kinds.append(("text", c.text)) + return kinds + + +def _generate_envelope(kinds): + for kind, call_id, result in (k for k in kinds if k[0] == "result"): + if call_id == "g1": + return json.loads(result) + return None + + +# --------------------------------------------------------------------------- # +# _state: context slice + enablement flag +# --------------------------------------------------------------------------- # + + +def test_context_slice_is_catalog_only(): + ctx = [ + {"description": A2UI_SCHEMA_CONTEXT_DESCRIPTION, "value": "CAT"}, + {"description": "Generation Guidelines", "value": "be terse"}, + ] + slice_ = build_ag_ui_context_slice(ctx) + assert slice_["a2ui_schema"] == "CAT" + assert slice_["context"] == [{"description": "Generation Guidelines", "value": "be terse"}] + assert "inject_a2ui_tool" not in slice_ # enablement is NOT bundled into context + + +def test_inject_flag_sourced_from_forwarded_props_only(): + assert read_inject_a2ui_flag({"injectA2UITool": True}) is True + assert read_inject_a2ui_flag({"injectA2UITool": False}) is False + assert read_inject_a2ui_flag({"inject_a2ui_tool": True}) is True # snake fallback + assert read_inject_a2ui_flag({}) is None # unset -> None (nullish) + assert read_inject_a2ui_flag(None) is None + # A context entry never enables injection. + assert read_inject_a2ui_flag({"context": [{"description": "x", "value": "y"}]}) is None + + +def test_stamp_and_read_round_trip(): + slice_ = {"a2ui_schema": "CAT"} + opts = stamp_context_slice({"metadata": {"x": "y"}, "store": True}, slice_) + assert opts["metadata"] == {"x": "y"} and opts["store"] is True + assert read_agent_state(opts) == slice_ + # object-attr form + obj = type("O", (), {"additional_properties": {"ag_ui_context": slice_}})() + assert read_agent_state(obj) == slice_ + assert read_agent_state(None) == {} + + +def test_no_a2ui_context_yields_empty_slice(): + assert build_ag_ui_context_slice(None) == {} + assert build_ag_ui_context_slice([{"description": "x", "value": "y"}]) == { + "context": [{"description": "x", "value": "y"}] + } + + +def test_to_history_messages_extracts_tool_result_payload(): + envelope = '{"a2ui_operations": []}' + messages = [ + Message(role="user", contents=[Content.from_text(text="hi")]), + Message(role="tool", contents=[Content.from_function_result(call_id="c", result=envelope)]), + ] + history = to_history_messages(messages) + assert history[0] == {"role": "user", "content": "hi"} + assert history[1]["role"] == "tool" + assert history[1]["content"] == envelope + + +# --------------------------------------------------------------------------- # +# AGUIContextAgent +# --------------------------------------------------------------------------- # + + +def test_context_agent_prepends_catalog_system_message(): + slice_ = build_ag_ui_context_slice( + [{"description": A2UI_SCHEMA_CONTEXT_DESCRIPTION, "value": '{"components":{"Card":{}}}'}] + ) + opts = stamp_context_slice(None, slice_) + + class Inner: + id = name = description = "x" + + def run(self, messages, *, stream=False, **kwargs): + return (stream, messages) + + stream, msgs = AGUIContextAgent(Inner()).run("hi", stream=True, options=opts) + assert stream is True + assert _role(msgs[0]) == "system" + assert "Available Components" in msgs[0].text and "Card" in msgs[0].text + assert _role(msgs[-1]) == "user" + + +def test_context_agent_passthrough_without_context(): + class Inner: + id = name = description = "x" + + def run(self, messages, *, stream=False, **kwargs): + return messages + + msgs = AGUIContextAgent(Inner()).run("hi", stream=False, options=None) + assert [_role(m) for m in msgs] == ["user"] + + +# --------------------------------------------------------------------------- # +# A2UIAgent streaming loop +# --------------------------------------------------------------------------- # + + +def test_streaming_progressive_paint_balancing_and_envelope(): + kinds = asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), _RenderSub()))) + + # >= 3 incremental render_a2ui arg fragments forwarded (progressive paint) + render_frags = [k for k in kinds if k[0] == "call" and k[1] == "render_a2ui"] + assert len(render_frags) >= 3 + + # balancing render result emitted + assert any(k[0] == "result" and k[1] == "r1" and "rendered" in str(k[2]) for k in kinds) + + # generate_a2ui envelope fed back with operations + env = _generate_envelope(kinds) + assert env is not None and "a2ui_operations" in env + assert any("createSurface" in op for op in env["a2ui_operations"]) + + # closing turn narrated + assert any(k[0] == "text" and "Done" in k[1] for k in kinds) + + +def test_streaming_loop_ends_when_planner_stops_requesting(): + # Normal termination: planner requests once, then narrates -> loop ends after + # the planner's no-generate turn (the tool is still advertised on that turn; + # withholding only applies to the round-cap closing turn, below). + inner = _GenerateOnceInner() + asyncio.run(_drive(A2UIAgent(inner, _RenderSub()))) + assert inner.calls == 2 + + +def test_streaming_closing_turn_withholds_generate_tool_at_round_cap(): + from agent_framework_ag_ui._a2ui._agent import MAX_PLANNER_ROUNDS + + class AlwaysGenerateInner: + id = name = description = "planner" + + def __init__(self): + self.tools_per_call = [] + + def run(self, messages, *, stream=False, session=None, tools=None, **kwargs): + self.tools_per_call.append([getattr(t, "name", None) for t in (tools or [])]) + + async def gen(): + yield AgentResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="g", name="generate_a2ui", arguments="{}")], + ) + + return gen() + + inner = AlwaysGenerateInner() + asyncio.run(_drive(A2UIAgent(inner, _RenderSub()))) + # MAX_PLANNER_ROUNDS planner rounds (tool advertised) + 1 closing turn (withheld). + assert len(inner.tools_per_call) == MAX_PLANNER_ROUNDS + 1 + assert all("generate_a2ui" in t for t in inner.tools_per_call[:MAX_PLANNER_ROUNDS]) + assert "generate_a2ui" not in inner.tools_per_call[-1] + + +def test_streaming_recovery_exhaustion(): + # Always-invalid (no root) -> recovery_exhausted after default 3 attempts. + sub = _RenderSub(components=[{"id": "x", "component": "Card"}]) + env = _generate_envelope(asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), sub)))) + assert env["code"] == "a2ui_recovery_exhausted" + assert len(env["attempts"]) == 3 + + +def test_streaming_retry_then_success(): + class RetrySub: + def __init__(self): + self.n = 0 + + def get_response(self, messages, *, stream=False, options=None): + self.n += 1 + comps = [{"id": "x", "component": "Card"}] if self.n == 1 else [{"id": "root", "component": "Card"}] + + async def gen(): + yield ChatResponseUpdate( + role="assistant", + contents=[ + Content.from_function_call( + call_id=f"r{self.n}", + name="render_a2ui", + arguments=json.dumps({"surfaceId": "s", "components": comps}), + ) + ], + ) + + return gen() + + env = _generate_envelope(asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), RetrySub())))) + assert "a2ui_operations" in env + + +# --------------------------------------------------------------------------- # +# P3: catalog-driven validation + error classification +# --------------------------------------------------------------------------- # + + +def test_forwarded_schema_drives_validation_catalog(): + # Card requires "text"; render omits it -> invalid via the forwarded catalog. + catalog = {"components": {"Card": {"required": ["text"]}}} + opts = stamp_context_slice( + None, + build_ag_ui_context_slice([{"description": A2UI_SCHEMA_CONTEXT_DESCRIPTION, "value": json.dumps(catalog)}]), + ) + sub = _RenderSub(components=[{"id": "root", "component": "Card"}]) + env = _generate_envelope(asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), sub), options=opts))) + assert env["code"] == "a2ui_recovery_exhausted" + errors = [e for attempt in env["attempts"] for e in attempt["errors"]] + assert any(e["code"] == "missing_required_prop" for e in errors) + + +def test_recoverable_subagent_error_retries(): + class RaisingSub: + def get_response(self, messages, *, stream=False, options=None): + async def gen(): + raise ValueError("transient") + yield # pragma: no cover + + return gen() + + env = _generate_envelope(asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), RaisingSub())))) + assert env["code"] == "a2ui_recovery_exhausted" + assert all(e["code"] == "subagent_error" for a in env["attempts"] for e in a["errors"]) + + +def test_programmer_error_is_rethrown(): + class TypeErrorSub: + def get_response(self, messages, *, stream=False, options=None): + async def gen(): + raise TypeError("bug") + yield # pragma: no cover + + return gen() + + with pytest.raises(TypeError): + asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), TypeErrorSub()))) + + +def test_mid_stream_death_balances_render_call(): + class PartialThenDieSub: + def get_response(self, messages, *, stream=False, options=None): + async def gen(): + yield ChatResponseUpdate( + role="assistant", + contents=[ + Content.from_function_call(call_id="rX", name="render_a2ui", arguments='{"surfaceId":"s","comp') + ], + ) + raise ValueError("died mid-stream") + + return gen() + + kinds = asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), PartialThenDieSub()))) + assert any(k[0] == "result" and k[1] == "rX" and "rendered" in str(k[2]) for k in kinds) + + +# --------------------------------------------------------------------------- # +# Non-streaming tool body + enable_a2ui +# --------------------------------------------------------------------------- # + + +def test_non_streaming_tool_body_runs_recovery(): + class NonStreamSub: + async def get_response(self, messages, *, stream=False, options=None): + return ChatResponse( + messages=[ + Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="r", + name="render_a2ui", + arguments=json.dumps( + {"surfaceId": "s", "components": [{"id": "root", "component": "Card"}]} + ), + ) + ], + ) + ] + ) + + agent = A2UIAgent( + inner_agent=type("I", (), {"id": "i", "name": "n", "description": "d"})(), subagent_chat_client=NonStreamSub() + ) + tool = agent._build_generate_tool([Message(role="user", contents=[Content.from_text(text="card")])], None) + assert tool.func is not None # executable on the non-streaming path + env = json.loads(asyncio.run(tool.func(intent="create"))) + assert "a2ui_operations" in env + + +def test_enable_a2ui_composition_order(): + inner = type("I", (), {"id": "i", "name": "n", "description": "d"})() + wrapped = enable_a2ui(inner, object()) + assert type(wrapped).__name__ == "AGUIContextAgent" + assert type(wrapped.inner_agent).__name__ == "A2UIAgent" + + +# --------------------------------------------------------------------------- # +# P5: auto-injection (plan_a2ui_injection) +# --------------------------------------------------------------------------- # + + +class _AgentWithClient: + id = name = description = "planner" + client = object() # inferable render sub-agent client + + +def test_plan_injection_off_without_flag(): + assert plan_a2ui_injection(agent=_AgentWithClient(), forwarded_props=None, existing_tool_names=[]) is None + assert plan_a2ui_injection(agent=_AgentWithClient(), forwarded_props={}, existing_tool_names=[]) is None + assert ( + plan_a2ui_injection(agent=_AgentWithClient(), forwarded_props={"injectA2UITool": False}, existing_tool_names=[]) + is None + ) + + +def test_plan_injection_wraps_and_drops_render_tool(): + plan = plan_a2ui_injection( + agent=_AgentWithClient(), + forwarded_props={"injectA2UITool": True}, + existing_tool_names=["some_other_tool"], + ) + assert plan is not None + assert type(plan["agent"]).__name__ == "AGUIContextAgent" + assert plan["drop_tool_names"] == ["render_a2ui"] + + +def test_plan_injection_string_flag_names_render_tool_to_drop(): + plan = plan_a2ui_injection( + agent=_AgentWithClient(), + forwarded_props={"injectA2UITool": "render_custom"}, + existing_tool_names=[], + ) + assert plan is not None and plan["drop_tool_names"] == ["render_custom"] + + +def test_plan_injection_user_prevails_when_generate_already_wired(): + plan = plan_a2ui_injection( + agent=_AgentWithClient(), + forwarded_props={"injectA2UITool": True}, + existing_tool_names=["generate_a2ui"], + ) + assert plan is None + + +def test_plan_injection_skips_already_wrapped_agent(): + wrapped = enable_a2ui(_AgentWithClient(), object()) + assert plan_a2ui_injection(agent=wrapped, forwarded_props={"injectA2UITool": True}, existing_tool_names=[]) is None + + +def test_plan_injection_skips_when_no_client_inferable(): + agent = type("NoClient", (), {"id": "i", "name": "n", "description": "d"})() + assert plan_a2ui_injection(agent=agent, forwarded_props={"injectA2UITool": True}, existing_tool_names=[]) is None + + +def test_plan_injection_runtime_false_beats_backend_opt_in(): + # Nullish fallback: explicit runtime false disables even when backend config opts in. + assert ( + plan_a2ui_injection( + agent=_AgentWithClient(), + forwarded_props={"injectA2UITool": False}, + existing_tool_names=[], + config={"inject_a2ui_tool": True}, + ) + is None + ) + # Backend opt-in applies when runtime is unset. + plan = plan_a2ui_injection( + agent=_AgentWithClient(), + forwarded_props={}, + existing_tool_names=[], + config={"inject_a2ui_tool": True}, + ) + assert plan is not None + + +def test_run_agent_stream_auto_wraps_and_drops_render_tool(stub_agent): + """Integration: run_agent_stream auto-wraps the agent and strips render_a2ui.""" + from agent_framework_ag_ui._agent import AgentConfig + from agent_framework_ag_ui._agent_run import run_agent_stream + + # Planner that narrates (no generate call) so the A2UIAgent loop ends after one + # inner turn — no render sub-agent call needed for this assertion. + agent = stub_agent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="hi there")], role="assistant")], + client=object(), # inferable render sub-agent client + ) + + input_data = { + "messages": [{"role": "user", "content": "make a card"}], + # The a2ui-middleware injects render_a2ui into the tool list. + "tools": [ + { + "name": "render_a2ui", + "description": "render", + "parameters": {"type": "object", "properties": {}}, + } + ], + "context": [{"description": A2UI_SCHEMA_CONTEXT_DESCRIPTION, "value": '{"components":{}}'}], + "forwardedProps": {"injectA2UITool": True}, + } + + async def _consume(): + async for _ in run_agent_stream(input_data, agent, AgentConfig()): + pass + + asyncio.run(_consume()) + + received = [getattr(t, "name", None) for t in (agent.tools_received or [])] + assert "generate_a2ui" in received # auto-injected + assert "render_a2ui" not in received # middleware-injected render tool dropped + + +# --------------------------------------------------------------------------- # +# Regression: real-LLM-only bugs (masked by pre-coalesced test doubles / aimock) +# --------------------------------------------------------------------------- # + + +class _FragmentedGenerateInner: + """Planner streaming ONE generate_a2ui call the way MAF's Chat-Completions client + actually does (verified empirically against gpt-4o): + + 1. an OPENING fragment carrying id + name + EMPTY args, + 2. argument-delta fragments with name="" and call_id="" (streaming visibility), + 3. a FINAL COALESCED fragment repeating id + name + the FULL args. + + The coalescer must select the coalesced full args BY NAME. Accumulating every + fragment name-agnostically would concatenate the deltas AND the coalesced copy, + doubling the JSON into an unparseable string (the client emits the args twice: + ``buf_all == 2 * buf_named``). It must also not treat each fragment as its own + call (that emitted duplicate tool results for one call id -> a 400 on replay). + """ + + id = name = description = "planner" + + def __init__(self): + self.calls = 0 + + def run(self, messages, *, stream=False, session=None, tools=None, **kwargs): + self.calls += 1 + n = self.calls + # intent=update + a target surface with no prior render. The OUTCOME (an + # "update target not found" error envelope) depends on BOTH the intent and the + # target_surface_id being correctly coalesced. If the args are doubled/lost, the + # tool defaults to intent=create and paints a normal surface instead — so this + # test fails under the name-agnostic-doubling mutation (mutation-effective). + full = '{"intent":"update","target_surface_id":"ghost-surface"}' + deltas = ['{"intent":"upda', 'te","target_surface_id":', '"ghost-surface"}'] + + async def gen(): + if n == 1: + yield AgentResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="g1", name="generate_a2ui", arguments="")], + ) + for piece in deltas: + yield AgentResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="", name="", arguments=piece)], + ) + # Final coalesced fragment (MAF re-emits id + name + FULL args). + yield AgentResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="g1", name="generate_a2ui", arguments=full)], + ) + else: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="Done.")]) + + return gen() + + +class _RenderSubMAF: + """Render sub-agent streaming render_a2ui the way MAF's Chat-Completions client does: + an opening fragment (id + name + empty args), name="" argument deltas, then a final + coalesced fragment (id + name + FULL args). Guards ``_extract_render_args`` against a + name-agnostic rewrite that would double the arguments (deltas + coalesced) into + invalid JSON.""" + + def __init__(self, components=None): + self._components = components or [{"id": "root", "component": "Card"}] + + def get_response(self, messages, *, stream=False, options=None): + full = json.dumps({"surfaceId": "s1", "components": self._components}) + mid = len(full) // 2 + + async def gen(): + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="r1", name="render_a2ui", arguments="")], + ) + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="", name="", arguments=full[:mid])], + ) + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="", name="", arguments=full[mid:])], + ) + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="r1", name="render_a2ui", arguments=full)], + ) + + return gen() + + +def test_streaming_coalesces_fragmented_generate_call(): + kinds = asyncio.run(_drive(A2UIAgent(_FragmentedGenerateInner(), _RenderSub()))) + # Exactly ONE generate_a2ui result despite the call arriving as many fragments + + # a coalesced copy (no duplicate tool results for one call id). + gen_results = [k for k in kinds if k[0] == "result" and k[1] == "g1"] + assert len(gen_results) == 1 + env = _generate_envelope(kinds) + assert env is not None + # The coalesced args (intent=update, target_surface_id=ghost-surface) drive an + # "update target not found" error envelope. This assertion depends on the args + # being correctly coalesced: doubling/losing them defaults to create + a painted + # surface (no error), which fails here. + assert "error" in env + assert "ghost-surface" in env["error"] + assert "a2ui_operations" not in env + + +def test_streaming_render_selects_coalesced_args_not_doubled(): + # MAF emits render args as deltas + a coalesced copy; _extract_render_args must + # select the coalesced full args by name. Name-agnostic accumulation would join + # deltas + coalesced into doubled, unparseable JSON -> None -> no surface. + kinds = asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), _RenderSubMAF()))) + env = _generate_envelope(kinds) + assert env is not None + uc = next(o["updateComponents"] for o in env["a2ui_operations"] if "updateComponents" in o) + assert any(c.get("id") == "root" for c in uc["components"]) + + +def test_forwarded_list_catalog_normalized_not_crash(): + # The A2UI middleware forwards components as an ARRAY; the validator wants a + # name->schema mapping. _resolve_catalog must normalize instead of crashing. + schema = json.dumps({"catalogId": "cat", "components": [{"name": "Card"}]}) + opts = stamp_context_slice( + None, build_ag_ui_context_slice([{"description": A2UI_SCHEMA_CONTEXT_DESCRIPTION, "value": schema}]) + ) + sub = _RenderSub(components=[{"id": "root", "component": "Card"}]) + env = _generate_envelope(asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), sub), options=opts))) + assert env is not None # did not raise AttributeError on the list-shaped catalog + + +def test_forwarded_catalog_id_binds_surface_when_unconfigured(): + # Zero-config (advanced): no backend default_catalog_id -> bind the surface to the + # catalog id the client forwarded, not the basic fallback. + schema = json.dumps({"catalogId": "https://x/custom_catalog.json", "components": [{"name": "Card"}]}) + opts = stamp_context_slice( + None, build_ag_ui_context_slice([{"description": A2UI_SCHEMA_CONTEXT_DESCRIPTION, "value": schema}]) + ) + sub = _RenderSub(components=[{"id": "root", "component": "Card"}]) + env = _generate_envelope(asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), sub), options=opts))) + assert env["a2ui_operations"][0]["createSurface"]["catalogId"] == "https://x/custom_catalog.json" + + +def test_strip_context_slice_removes_key_and_prunes_empty(): + from agent_framework_ag_ui._a2ui._state import A2UI_CONTEXT_KEY, strip_context_slice + + opts = {"additional_properties": {A2UI_CONTEXT_KEY: {"a2ui_schema": "x"}, "keep": 1}, "metadata": {}} + out = strip_context_slice(opts) + assert A2UI_CONTEXT_KEY not in out["additional_properties"] + assert out["additional_properties"]["keep"] == 1 + # When removing the slice empties additional_properties, drop the key entirely. + out2 = strip_context_slice({"additional_properties": {A2UI_CONTEXT_KEY: {}}}) + assert "additional_properties" not in out2 + assert strip_context_slice(None) is None + + +def test_planner_options_do_not_leak_context_slice(): + # additional_properties is forwarded to the provider SDK, which rejects the + # adapter-private ag-ui context slice; A2UIAgent must strip it before the planner. + from agent_framework_ag_ui._a2ui._state import A2UI_CONTEXT_KEY + + opts = stamp_context_slice(None, {"a2ui_schema": "x"}) + + class _RecordInner: + id = name = description = "planner" + + def __init__(self): + self.seen: list = [] + + def run(self, messages, *, stream=False, session=None, tools=None, **kwargs): + self.seen.append(kwargs.get("options")) + + async def gen(): + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="Done.")]) + + return gen() + + inner = _RecordInner() + asyncio.run(_drive(A2UIAgent(inner, _RenderSub()), options=opts)) + seen = inner.seen[0] + ap = seen.get("additional_properties", {}) if isinstance(seen, dict) else {} + assert A2UI_CONTEXT_KEY not in ap + + +# Robustness to a DELTAS-ONLY provider (no coalesced-final fragment): the all-fragment +# concat fallback must reassemble the full args. Guards against a regression that only +# handles MAF's deltas+coalesced pattern (the reviewers' feared case). + + +class _RenderSubDeltasOnly: + """Render sub-agent that streams args as deltas ONLY (name="" after the first + fragment) with NO final coalesced fragment — the hypothetical provider the + name-bearing concat alone cannot reassemble.""" + + def __init__(self, components=None): + self._components = components or [{"id": "root", "component": "Card"}] + + def get_response(self, messages, *, stream=False, options=None): + full = json.dumps({"surfaceId": "s1", "components": self._components}) + third = max(1, len(full) // 3) + pieces = [full[:third], full[third : 2 * third], full[2 * third :]] + + async def gen(): + # Opening fragment: name + id, empty args. + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="r1", name="render_a2ui", arguments="")], + ) + # Argument deltas: name="" / call_id="", no coalesced final. + for p in pieces: + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="", name="", arguments=p)], + ) + + return gen() + + +class _DeltasOnlyGenerateInner: + """Planner streaming generate_a2ui as deltas ONLY (no coalesced final).""" + + id = name = description = "planner" + + def __init__(self): + self.calls = 0 + + def run(self, messages, *, stream=False, session=None, tools=None, **kwargs): + self.calls += 1 + n = self.calls + + async def gen(): + if n == 1: + yield AgentResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="g1", name="generate_a2ui", arguments="")], + ) + # Deltas-only (no coalesced final); the all-fragment concat fallback must + # reassemble intent=update + target_surface_id=ghost-surface. + for p in ('{"intent":"upda', 'te","target_surface_id":', '"ghost-surface"}'): + yield AgentResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="", name="", arguments=p)], + ) + else: + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text(text="Done.")]) + + return gen() + + +def test_streaming_render_reassembles_deltas_only_provider(): + kinds = asyncio.run(_drive(A2UIAgent(_GenerateOnceInner(), _RenderSubDeltasOnly()))) + env = _generate_envelope(kinds) + assert env is not None # all-fragment concat fallback reassembled the full args + uc = next(o["updateComponents"] for o in env["a2ui_operations"] if "updateComponents" in o) + assert any(c.get("id") == "root" for c in uc["components"]) + + +def test_streaming_generate_reassembles_deltas_only_provider(): + kinds = asyncio.run(_drive(A2UIAgent(_DeltasOnlyGenerateInner(), _RenderSub()))) + gen_results = [k for k in kinds if k[0] == "result" and k[1] == "g1"] + assert len(gen_results) == 1 + env = _generate_envelope(kinds) + assert env is not None + # Fallback all-fragment concat reassembled intent=update + ghost target -> the + # "update target not found" error envelope (depends on the coalesced args). + assert "error" in env + assert "ghost-surface" in env["error"] + + +def test_streaming_aclose_mid_render_does_not_raise_generator_exit(): + # Closing the async generator mid-render (client disconnect / aclose) throws + # GeneratorExit into the streaming loop. The mid-stream except must re-raise it + # WITHOUT yielding a balancing tool result — yielding during GeneratorExit raises + # "async generator ignored GeneratorExit". Red-green guard for that fix. + sub = _RenderSub(fragments=['{"surfaceId":"s1","comp', 'onents":[{"id":"root",', '"component":"Card"}]}']) + agent = A2UIAgent(_GenerateOnceInner(), sub) + + async def run(): + agen = agent.run("x", stream=True) + async for update in agen: + if any( + getattr(c, "type", None) == "function_call" and getattr(c, "name", None) == "render_a2ui" + for c in update.contents + ): + await agen.aclose() # must complete cleanly, not raise RuntimeError + return True + return False + + assert asyncio.run(run()) is True + + +def test_sanitize_unanswered_tool_calls_strips_dangling_a2ui_calls(): + # A2UI surfaces persist as activities, so the assistant's generate_a2ui/render_a2ui + # tool calls come back on a later turn WITHOUT tool results. Those dangling calls + # must be stripped before the planner call (OpenAI rejects unanswered tool_calls), + # while assistant text and balanced call/result pairs (log_a2ui_event) are kept. + from agent_framework_ag_ui._a2ui._agent import _sanitize_unanswered_tool_calls + + msgs = [ + Message(role="user", contents=[Content.from_text(text="hi")]), + Message( + role="assistant", + contents=[ + Content.from_text(text="Here is your UI."), + Content.from_function_call(call_id="g1", name="generate_a2ui", arguments="{}"), + Content.from_function_call(call_id="r1", name="render_a2ui", arguments="{}"), + ], + ), + Message( + role="assistant", contents=[Content.from_function_call(call_id="e1", name="log_a2ui_event", arguments="{}")] + ), + Message(role="tool", contents=[Content.from_function_result(call_id="e1", result="{}")]), + ] + out = _sanitize_unanswered_tool_calls(msgs) + call_ids = [ + getattr(c, "call_id", None) + for m in out + for c in (getattr(m, "contents", None) or []) + if getattr(c, "type", None) == "function_call" + ] + assert "g1" not in call_ids and "r1" not in call_ids # dangling a2ui calls dropped + assert "e1" in call_ids # balanced log_a2ui_event kept + texts = [c.text for m in out for c in (getattr(m, "contents", None) or []) if getattr(c, "type", None) == "text"] + assert "Here is your UI." in texts # assistant narration preserved diff --git a/python/uv.lock b/python/uv.lock index 64bf0b02598..92c963b7a76 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -104,6 +104,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6a/00/b08f23b7d7e1e14ce01419a467b583edbb93c6cdb8654e54a9cc579cd61f/addict-2.4.0-py3-none-any.whl", hash = "sha256:249bb56bbfd3cdc2a004ea0ff4c2b6ddc84d53bc2194761636eb314d5cfa5dfc", size = 3832, upload-time = "2020-11-21T16:21:29.588Z" }, ] +[[package]] +name = "ag-ui-a2ui-toolkit" +version = "0.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/ce/85f3960a83d962e5690bc0f27a3baf3bf1602edc2b0603085928c964ea14/ag_ui_a2ui_toolkit-0.0.4.tar.gz", hash = "sha256:172e2724e53df8173685a3fb896a6e5175eea06e1dc166c715db110ba4beba76", size = 18960, upload-time = "2026-06-17T13:34:28.695Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/7a/acf85b01cd996bd011b71e181fd9f3daff5396fc3b7d78ba9445bfc08ecf/ag_ui_a2ui_toolkit-0.0.4-py3-none-any.whl", hash = "sha256:236fc511e1ec2399bcda0c14a109b3fb0a0c3e3988c18ef1918745b1c1535e30", size = 21315, upload-time = "2026-06-17T13:34:29.505Z" }, +] + [[package]] name = "ag-ui-protocol" version = "0.1.19" @@ -209,6 +218,9 @@ dependencies = [ ] [package.optional-dependencies] +a2ui = [ + { name = "ag-ui-a2ui-toolkit", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] dev = [ { name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pytest", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -221,6 +233,7 @@ test = [ [package.metadata] requires-dist = [ + { name = "ag-ui-a2ui-toolkit", marker = "extra == 'a2ui'", specifier = ">=0.0.4" }, { name = "ag-ui-protocol", specifier = ">=0.1.19,<0.2" }, { name = "agent-framework-core", editable = "packages/core" }, { name = "fastapi", specifier = ">=0.121.0,<0.140.0" }, @@ -229,7 +242,7 @@ requires-dist = [ { name = "sse-starlette", specifier = ">=3.4.5,<4" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.30.0,<1" }, ] -provides-extras = ["dev"] +provides-extras = ["a2ui", "dev"] [package.metadata.requires-dev] test = [{ name = "agent-framework-orchestrations", editable = "packages/orchestrations" }]