Refactor A2UI extension activation

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,