update key names, show supported catalogs in agent card (#405)

Author: wrenj

a2a_agents/python/a2ui_extension/README.md

@@ -2,6 +2,20 @@
 
 This is the Python implementation of the a2ui extension.
 
+## Running Tests
+
+1. Navigate to the a2ui_extension dir:
+
+   ```bash
+   cd a2a_agents/python/a2ui_extension
+   ```
+
+2. Run the tests
+
+   ```bash
+   uv run --with pytest pytest tests/test_extension.py
+   ```
+
 ## Disclaimer
 
 Important: The sample code provided is for demonstration purposes and illustrates the mechanics of A2UI and the Agent-to-Agent (A2A) protocol. When building production applications, it is critical to treat any agent operating outside of your direct control as a potentially untrusted entity.

a2a_agents/python/a2ui_extension/src/a2ui/a2ui_extension.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, List
 
 from a2a.server.agent_execution import RequestContext
 from a2a.types import AgentExtension, Part, DataPart
@@ -29,7 +29,7 @@
 SUPPORTED_CATALOG_IDS_KEY = "supportedCatalogIds"
 INLINE_CATALOGS_KEY = "inlineCatalogs"
 
-STANDARD_CATALOG_ID = "https://raw.githubusercontent.com/google/A2UI/refs/heads/main/specification/0.8/json/standard_catalog_definition.json"
+STANDARD_CATALOG_ID = "https://github.com/google/A2UI/blob/main/specification/0.8/json/standard_catalog_definition.json"
 
 def create_a2ui_part(a2ui_data: dict[str, Any]) -> Part:
     """Creates an A2A Part containing A2UI data.
@@ -80,20 +80,28 @@ def get_a2ui_datapart(part: Part) -> Optional[DataPart]:
     return None
 
 
+AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY = "supportedCatalogIds"
+AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY = "acceptsInlineCatalogs"
+
 def get_a2ui_agent_extension(
-    accepts_inline_custom_catalog: bool = False,
+    accepts_inline_catalogs: bool = False,
+    supported_catalog_ids: List[str] = [],
 ) -> AgentExtension:
     """Creates the A2UI AgentExtension configuration.
 
     Args:
-        accepts_inline_custom_catalog: Whether the agent accepts inline custom catalogs.
+        accepts_inline_catalogs: Whether the agent accepts inline custom catalogs.
+        supported_catalog_ids: All pre-defined catalogs the agent is known to support.
 
     Returns:
         The configured A2UI AgentExtension.
     """
     params = {}
-    if accepts_inline_custom_catalog:
-        params["acceptsInlineCustomCatalog"] = True  # Only set if not default of False
+    if accepts_inline_catalogs:
+        params[AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY] = True  # Only set if not default of False
+
+    if supported_catalog_ids:
+        params[AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY] = supported_catalog_ids
 
     return AgentExtension(
         uri=A2UI_EXTENSION_URI,

a2a_agents/python/a2ui_extension/tests/test_extension.py

@@ -16,7 +16,7 @@
 from a2a.server.agent_execution import RequestContext
 from a2a.types import DataPart, TextPart, Part
 from a2ui import a2ui_extension
-
+from a2ui.a2ui_extension import AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY, AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY
 from unittest.mock import MagicMock
 
 
@@ -64,12 +64,24 @@ def test_get_a2ui_agent_extension():
     assert agent_extension.params is None
 
 
-def test_get_a2ui_agent_extension_with_inline_custom_catalog():
+def test_get_a2ui_agent_extension_with_accepts_inline_catalogs():
+    accepts_inline_catalogs = True
+    agent_extension = a2ui_extension.get_a2ui_agent_extension(
+        accepts_inline_catalogs=accepts_inline_catalogs
+    )
+    assert agent_extension.uri == a2ui_extension.A2UI_EXTENSION_URI
+    assert agent_extension.params is not None
+    assert agent_extension.params.get(AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY) == accepts_inline_catalogs
+    
+    
+def test_get_a2ui_agent_extension_with_supported_catalog_ids():
+    supported_catalog_ids = ["a", "b", "c"]
     agent_extension = a2ui_extension.get_a2ui_agent_extension(
-        accepts_inline_custom_catalog=True
+        supported_catalog_ids=supported_catalog_ids
     )
     assert agent_extension.uri == a2ui_extension.A2UI_EXTENSION_URI
     assert agent_extension.params is not None
+    assert agent_extension.params.get(AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY) == supported_catalog_ids
 
 
 def test_try_activate_a2ui_extension():

a2a_agents/python/a2ui_extension/uv.lock

@@ -50,15 +50,15 @@ def main(host, port, subagent_urls):
 
         base_url = f"http://{host}:{port}"
 
-        orchestrator_agent = asyncio.run(OrchestratorAgent.build_agent(subagent_urls=subagent_urls))
-        agent_executor = OrchestratorAgentExecutor(base_url=base_url, agent=orchestrator_agent)
+        orchestrator_agent, agent_card = asyncio.run(OrchestratorAgent.build_agent(base_url=base_url, subagent_urls=subagent_urls))
+        agent_executor = OrchestratorAgentExecutor(agent=orchestrator_agent)
 
         request_handler = DefaultRequestHandler(
             agent_executor=agent_executor,
             task_store=InMemoryTaskStore(),
         )
         server = A2AStarletteApplication(
-            agent_card=agent_executor.get_agent_card(), http_handler=request_handler
+            agent_card=agent_card, http_handler=request_handler
         )
         import uvicorn
 

samples/agent/adk/orchestrator/__main__.py

@@ -30,15 +30,16 @@
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 from subagent_route_manager import SubagentRouteManager
-from a2ui.a2ui_extension import is_a2ui_part, A2UI_EXTENSION_URI
 from typing import override
 from a2a.types import TransportProtocol as A2ATransport
 
-logger = logging.getLogger(__name__)
 from a2a.client.middleware import ClientCallInterceptor
 from a2a.client.client import ClientConfig as A2AClientConfig
 from a2a.client.client_factory import ClientFactory as A2AClientFactory
-from a2ui.a2ui_extension import A2UI_CLIENT_CAPABILITIES_KEY
+from a2ui.a2ui_extension import is_a2ui_part, A2UI_CLIENT_CAPABILITIES_KEY, A2UI_EXTENSION_URI, AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY, AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY, get_a2ui_agent_extension
+from a2a.types import AgentCapabilities, AgentCard, AgentExtension
+
+logger = logging.getLogger(__name__)
 
 class A2UIMetadataInterceptor(ClientCallInterceptor):
     @override
@@ -115,18 +116,28 @@ async def programmtically_route_user_action_to_subagent(
         return None
 
     @classmethod
-    async def build_agent(cls, subagent_urls: List[str]) -> LlmAgent:
+    async def build_agent(cls, base_url: str, subagent_urls: List[str]) -> (LlmAgent, AgentCard):
         """Builds the LLM agent for the orchestrator_agent agent."""
 
         subagents = []
+        supported_catalog_ids = set()
+        skills = []
+        accepts_inline_catalogs = False
         for subagent_url in subagent_urls:
             async with httpx.AsyncClient() as httpx_client:
                 resolver = A2ACardResolver(
                     httpx_client=httpx_client,
                     base_url=subagent_url,
                 )
 
-                subagent_card =  await resolver.get_agent_card()
+                subagent_card =  await resolver.get_agent_card()                
+                for extension in subagent_card.capabilities.extensions or []:
+                    if extension.uri == A2UI_EXTENSION_URI and extension.params:
+                        supported_catalog_ids.update(extension.params.get(AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY) or [])
+                        accepts_inline_catalogs |= bool(extension.params.get(AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY))
+                
+                skills.extend(subagent_card.skills)
+                
                 logger.info('Successfully fetched public agent card:' + subagent_card.model_dump_json(indent=2, exclude_none=True))
 
                 # clean name for adk
@@ -172,7 +183,7 @@ async def build_agent(cls, subagent_urls: List[str]) -> LlmAgent:
                 logger.info(f'Created remote agent with description: {description}')
 
         LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gemini/gemini-2.5-flash")
-        return LlmAgent(
+        agent = LlmAgent(
             model=LiteLlm(model=LITELLM_MODEL),
             name="orchestrator_agent",
             description="An agent that orchestrates requests to multiple other agents",
@@ -186,3 +197,21 @@ async def build_agent(cls, subagent_urls: List[str]) -> LlmAgent:
             sub_agents=subagents,
             before_model_callback=cls.programmtically_route_user_action_to_subagent,
         )
+
+        agent_card = AgentCard(
+            name="Orchestrator Agent",
+            description="This agent orchestrates requests to multiple subagents.",
+            url=base_url,
+            version="1.0.0",
+            default_input_modes=OrchestratorAgent.SUPPORTED_CONTENT_TYPES,
+            default_output_modes=OrchestratorAgent.SUPPORTED_CONTENT_TYPES,
+            capabilities=AgentCapabilities(
+                streaming=True,
+                extensions=[get_a2ui_agent_extension(
+                    accepts_inline_catalogs=accepts_inline_catalogs,
+                    supported_catalog_ids=list(supported_catalog_ids))],
+            ),
+            skills=skills,
+        )
+
+        return agent, agent_card

samples/agent/adk/orchestrator/agent.py

@@ -30,8 +30,7 @@
     A2aAgentExecutorConfig,
     A2aAgentExecutor,
 )
-from a2a.types import AgentCapabilities, AgentCard, AgentExtension
-from a2ui.a2ui_extension import is_a2ui_part, try_activate_a2ui_extension, A2UI_EXTENSION_URI, STANDARD_CATALOG_ID, SUPPORTED_CATALOG_IDS_KEY, get_a2ui_agent_extension, A2UI_CLIENT_CAPABILITIES_KEY
+from a2ui.a2ui_extension import is_a2ui_part, try_activate_a2ui_extension, A2UI_EXTENSION_URI, STANDARD_CATALOG_ID, SUPPORTED_CATALOG_IDS_KEY, A2UI_CLIENT_CAPABILITIES_KEY
 from google.adk.a2a.converters import event_converter
 from a2a.server.events import Event as A2AEvent
 from google.adk.events.event import Event
@@ -48,9 +47,7 @@
 class OrchestratorAgentExecutor(A2aAgentExecutor):
     """Contact AgentExecutor Example."""
 
-    def __init__(self, base_url: str, agent: LlmAgent):
-        self._base_url = base_url
-
+    def __init__(self, agent: LlmAgent):
         config = A2aAgentExecutorConfig(
             gen_ai_part_converter=part_converters.convert_genai_part_to_a2a_part,
             a2a_part_converter=part_converters.convert_a2a_part_to_genai_part,
@@ -117,21 +114,6 @@ def convert_event_to_a2a_events_and_save_surface_id_to_subagent_name(
 
         return a2a_events
 
-    def get_agent_card(self) -> AgentCard:
-        return AgentCard(
-            name="Orchestrator Agent",
-            description="This agent orchestrates to multiple subagents to provide.",
-            url=self._base_url,
-            version="1.0.0",
-            default_input_modes=OrchestratorAgent.SUPPORTED_CONTENT_TYPES,
-            default_output_modes=OrchestratorAgent.SUPPORTED_CONTENT_TYPES,
-            capabilities=AgentCapabilities(
-                streaming=True,
-                extensions=[get_a2ui_agent_extension()],
-            ),
-            skills=[],
-        )
-
     @override
     async def _prepare_session(
         self,

samples/agent/adk/orchestrator/agent_executor.py

@@ -31,7 +31,7 @@
 
 logger = logging.getLogger(__name__)
 
-RIZZCHARTS_CATALOG_URI = "https://raw.githubusercontent.com/google/A2UI/refs/heads/main/a2a_agents/python/adk/samples/rizzcharts/rizzcharts_catalog_definition.json"
+RIZZCHARTS_CATALOG_URI = "https://github.com/google/A2UI/blob/main/samples/agent/adk/rizzcharts/rizzcharts_catalog_definition.json"
 
 class rizzchartsAgent:
     """An agent that runs an ecommerce dashboard"""

samples/agent/adk/rizzcharts/agent.py

@@ -84,7 +84,8 @@ def get_agent_card(self) -> AgentCard:
             default_output_modes=rizzchartsAgent.SUPPORTED_CONTENT_TYPES,
             capabilities=AgentCapabilities(
                 streaming=True,
-                extensions=[get_a2ui_agent_extension()],
+                extensions=[get_a2ui_agent_extension(
+                    supported_catalog_ids=[STANDARD_CATALOG_ID, RIZZCHARTS_CATALOG_URI])],
             ),
             skills=[
                 AgentSkill(

samples/agent/adk/rizzcharts/agent_executor.py

@@ -71,7 +71,7 @@ def load_a2ui_schema(self, client_ui_capabilities: Optional[dict[str, Any]]) ->
                 logger.info(f"Loading inline component catalog {inline_catalog_str[:200]}")
                 catalog_json = json.loads(inline_catalog_str)
             else:
-                raise ValueError("Client UI capabilities not provided")
+                raise ValueError("No supported catalogs found in client UI capabilities")
 
             logger.info(f"Loading A2UI schema at {self._a2ui_schema_path}")
             a2ui_schema = self.get_file_content(self._a2ui_schema_path)