Refactor A2UI extension activation (#926)

* Refactor A2UI extension activation

The agents may support multiple versions of A2UI extensions. Based on
client requested extension, the agent SDK should dynamically activate
one that is supported by the agent.

* fix a typo in the mcp-apps-calculator folder

* Fix the contact_multiple_surfaces sample

Author: nan-yu

agent_sdks/python/src/a2ui/a2a.py

@@ -16,11 +16,17 @@
 from typing import Any, Optional, List
 
 from a2a.server.agent_execution import RequestContext
-from a2a.types import AgentExtension, Part, DataPart, TextPart
+from a2a.types import (
+    AgentExtension,
+    AgentCard,
+    Part,
+    DataPart,
+    TextPart,
+)
 
 logger = logging.getLogger(__name__)
 
-A2UI_EXTENSION_URI = "https://a2ui.org/a2a-extension/a2ui/v0.8"
+A2UI_EXTENSION_BASE_URI = "https://a2ui.org/a2a-extension/a2ui"
 AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY = "supportedCatalogIds"
 AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY = "acceptsInlineCatalogs"
 
@@ -78,12 +84,14 @@ def get_a2ui_datapart(part: Part) -> Optional[DataPart]:
 
 
 def get_a2ui_agent_extension(
+    version: str,
     accepts_inline_catalogs: bool = False,
     supported_catalog_ids: List[str] = [],
 ) -> AgentExtension:
   """Creates the A2UI AgentExtension configuration.
 
   Args:
+      version: The version of the A2UI extension to use.
       accepts_inline_catalogs: Whether the agent accepts inline catalogs.
       supported_catalog_ids: All pre-defined catalogs the agent is known to support.
 
@@ -100,7 +108,7 @@ def get_a2ui_agent_extension(
     params[AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY] = supported_catalog_ids
 
   return AgentExtension(
-      uri=A2UI_EXTENSION_URI,
+      uri=f"{A2UI_EXTENSION_BASE_URI}/v{version}",
       description="Provides agent driven UI using the A2UI JSON format.",
       params=params if params else None,
   )
@@ -151,20 +159,70 @@ def parse_response_to_parts(
   return parts
 
 
-def try_activate_a2ui_extension(context: RequestContext) -> bool:
+def _agent_extensions(agent_card: AgentCard) -> List[str]:
+  """Returns the A2UI extension URIs supported by the agent."""
+  extensions = []
+  if (
+      agent_card
+      and hasattr(agent_card, "capabilities")
+      and agent_card.capabilities
+      and hasattr(agent_card.capabilities, "extensions")
+      and agent_card.capabilities.extensions
+  ):
+    for ext in agent_card.capabilities.extensions:
+      if ext.uri and ext.uri.startswith(A2UI_EXTENSION_BASE_URI):
+        extensions.append(ext.uri)
+  return extensions
+
+
+def _requested_a2ui_extensions(context: RequestContext) -> List[str]:
+  """Returns the A2UI extension URIs requested by the client."""
+  requested_extensions = []
+  if hasattr(context, "requested_extensions") and context.requested_extensions:
+    requested_extensions.extend([
+        ext
+        for ext in context.requested_extensions
+        if isinstance(ext, str) and ext.startswith(A2UI_EXTENSION_BASE_URI)
+    ])
+
+  if (
+      hasattr(context, "message")
+      and context.message
+      and hasattr(context.message, "extensions")
+      and context.message.extensions
+  ):
+    requested_extensions.extend([
+        ext
+        for ext in context.message.extensions
+        if isinstance(ext, str) and ext.startswith(A2UI_EXTENSION_BASE_URI)
+    ])
+
+  return requested_extensions
+
+
+def try_activate_a2ui_extension(
+    context: RequestContext, agent_card: AgentCard
+) -> Optional[str]:
   """Activates the A2UI extension if requested.
 
   Args:
       context: The request context to check.
+      agent_card: The agent card to check supported extensions.
 
   Returns:
-      True if activated, False otherwise.
+      The version string of the activated A2UI extension, or None if not activated.
   """
-  if A2UI_EXTENSION_URI in context.requested_extensions or (
-      context.message
-      and context.message.extensions
-      and A2UI_EXTENSION_URI in context.message.extensions
-  ):
-    context.add_activated_extension(A2UI_EXTENSION_URI)
-    return True
-  return False
+  requested_extensions = _requested_a2ui_extensions(context)
+  if not requested_extensions:
+    return None
+
+  agent_advertised_extensions = _agent_extensions(agent_card)
+  if not agent_advertised_extensions:
+    return None
+
+  for req_uri in requested_extensions:
+    if req_uri in agent_advertised_extensions:
+      context.add_activated_extension(req_uri)
+      return req_uri.replace(f"{A2UI_EXTENSION_BASE_URI}/v", "")
+
+  return None

agent_sdks/python/src/a2ui/adk/a2a_extension/send_a2ui_to_client_toolset.py

@@ -105,7 +105,6 @@ async def get_examples(ctx: ReadonlyContext) -> str:
 
 from a2a import types as a2a_types
 from a2ui.a2a import (
-    A2UI_EXTENSION_URI,
     create_a2ui_part,
     parse_response_to_parts,
 )

agent_sdks/python/tests/test_a2a.py

@@ -51,17 +51,19 @@ def test_non_a2ui_part():
 
 
 def test_get_a2ui_agent_extension():
-  agent_extension = get_a2ui_agent_extension()
-  assert agent_extension.uri == A2UI_EXTENSION_URI
+  version = "0.8"
+  agent_extension = get_a2ui_agent_extension(version)
+  assert agent_extension.uri == f"{A2UI_EXTENSION_BASE_URI}/v{version}"
   assert agent_extension.params is None
 
 
 def test_get_a2ui_agent_extension_with_accepts_inline_catalogs():
+  version = "0.8"
   accepts_inline_catalogs = True
   agent_extension = get_a2ui_agent_extension(
-      accepts_inline_catalogs=accepts_inline_catalogs
+      version, accepts_inline_catalogs=accepts_inline_catalogs
   )
-  assert agent_extension.uri == A2UI_EXTENSION_URI
+  assert agent_extension.uri == f"{A2UI_EXTENSION_BASE_URI}/v{version}"
   assert agent_extension.params is not None
   assert (
       agent_extension.params.get(AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY)
@@ -70,11 +72,12 @@ def test_get_a2ui_agent_extension_with_accepts_inline_catalogs():
 
 
 def test_get_a2ui_agent_extension_with_supported_catalog_ids():
+  version = "0.8"
   supported_catalog_ids = ["a", "b", "c"]
   agent_extension = get_a2ui_agent_extension(
-      supported_catalog_ids=supported_catalog_ids
+      version, supported_catalog_ids=supported_catalog_ids
   )
-  assert agent_extension.uri == A2UI_EXTENSION_URI
+  assert agent_extension.uri == f"{A2UI_EXTENSION_BASE_URI}/v{version}"
   assert agent_extension.params is not None
   assert (
       agent_extension.params.get(AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY)
@@ -84,15 +87,26 @@ def test_get_a2ui_agent_extension_with_supported_catalog_ids():
 
 def test_try_activate_a2ui_extension():
   context = MagicMock(spec=RequestContext)
-  context.requested_extensions = [A2UI_EXTENSION_URI]
+  uri = f"{A2UI_EXTENSION_BASE_URI}/v0.8"
+  context.requested_extensions = [uri]
 
-  assert try_activate_a2ui_extension(context)
-  context.add_activated_extension.assert_called_once_with(A2UI_EXTENSION_URI)
+  card = MagicMock()
+  ext = MagicMock()
+  ext.uri = uri
+  card.capabilities.extensions = [ext]
+
+  assert try_activate_a2ui_extension(context, card) == "0.8"
+  context.add_activated_extension.assert_called_once_with(uri)
 
 
 def test_try_activate_a2ui_extension_not_requested():
   context = MagicMock(spec=RequestContext)
   context.requested_extensions = []
 
-  assert not try_activate_a2ui_extension(context)
+  card = MagicMock()
+  ext = MagicMock()
+  ext.uri = f"{A2UI_EXTENSION_BASE_URI}/v0.8"
+  card.capabilities.extensions = [ext]
+
+  assert try_activate_a2ui_extension(context, card) is None
   context.add_activated_extension.assert_not_called()

samples/agent/adk/component_gallery/__main__.py

@@ -28,6 +28,7 @@
 from starlette.middleware.cors import CORSMiddleware
 from starlette.staticfiles import StaticFiles
 from dotenv import load_dotenv
+from a2ui.core.schema.constants import VERSION_0_8
 
 
 from agent_executor import ComponentGalleryExecutor
@@ -49,9 +50,12 @@
 @click.option("--port", default=10005)
 def main(host, port):
   try:
+    extensions = []
+    for v in [VERSION_0_8]:
+      extensions.append(get_a2ui_agent_extension(v))
     capabilities = AgentCapabilities(
         streaming=True,
-        extensions=[get_a2ui_agent_extension()],
+        extensions=extensions,
     )
 
     # Skill definition
@@ -76,7 +80,7 @@ def main(host, port):
         skills=[skill],
     )
 
-    agent_executor = ComponentGalleryExecutor(base_url=base_url)
+    agent_executor = ComponentGalleryExecutor(base_url=base_url, agent_card=agent_card)
 
     request_handler = DefaultRequestHandler(
         agent_executor=agent_executor,