A2ui extension selection

agent_sdks/python/src/a2ui/a2a.py

@@ -200,6 +200,25 @@ def _requested_a2ui_extensions(context: RequestContext) -> List[str]:
   return requested_extensions
 
 
+def _select_newest_a2ui_extension(
+    requested_extensions: List[str], agent_advertised_extensions: List[str]
+) -> Optional[str]:
+  """Selects the newest A2UI extension URI from the matched extensions."""
+  matched_extensions = [
+      uri for uri in requested_extensions if uri in agent_advertised_extensions
+  ]
+  if not matched_extensions:
+    return None
+
+  def _version_key(uri: str) -> tuple:
+    version_str = uri.replace(f"{A2UI_EXTENSION_BASE_URI}/v", "")
+    from packaging.version import parse as parse_version
+
+    return parse_version(version_str)
+
+  return max(matched_extensions, key=_version_key)
+
+
 def try_activate_a2ui_extension(
     context: RequestContext, agent_card: AgentCard
 ) -> Optional[str]:
@@ -220,9 +239,11 @@ def try_activate_a2ui_extension(
   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", "")
+  selected_uri = _select_newest_a2ui_extension(
+      requested_extensions, agent_advertised_extensions
+  )
+  if selected_uri:
+    context.add_activated_extension(selected_uri)
+    return selected_uri.replace(f"{A2UI_EXTENSION_BASE_URI}/v", "")
 
   return None

agent_sdks/python/tests/test_a2a.py

@@ -16,6 +16,7 @@
 from a2a.server.agent_execution import RequestContext
 from a2a.types import DataPart, TextPart, Part
 from a2ui.a2a import *
+from a2ui.a2a import _select_newest_a2ui_extension
 from unittest.mock import MagicMock
 
 
@@ -110,3 +111,49 @@ def test_try_activate_a2ui_extension_not_requested():
 
   assert try_activate_a2ui_extension(context, card) is None
   context.add_activated_extension.assert_not_called()
+
+
+def test_select_newest_a2ui_extension():
+  requested = [
+      f"{A2UI_EXTENSION_BASE_URI}/v0.1.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v1.2.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v0.8.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v1.10.0",
+  ]
+  advertised = [
+      f"{A2UI_EXTENSION_BASE_URI}/v0.1.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v1.2.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v1.10.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v2.0.0",
+  ]
+  # Should match 0.1.0, 1.2.0 and 1.10.0, and pick 1.10.0 as the newest
+  newest = _select_newest_a2ui_extension(requested, advertised)
+  assert newest == f"{A2UI_EXTENSION_BASE_URI}/v1.10.0"
+
+
+def test_select_newest_a2ui_extension_no_match():
+  requested = [f"{A2UI_EXTENSION_BASE_URI}/v0.1.0"]
+  advertised = [f"{A2UI_EXTENSION_BASE_URI}/v1.2.0"]
+  assert _select_newest_a2ui_extension(requested, advertised) is None
+
+
+def test_try_activate_a2ui_extension_multiple_versions():
+  context = MagicMock(spec=RequestContext)
+  context.requested_extensions = [
+      f"{A2UI_EXTENSION_BASE_URI}/v0.8.0",
+      f"{A2UI_EXTENSION_BASE_URI}/v1.0.0",
+  ]
+
+  card = MagicMock()
+  ext1 = MagicMock()
+  ext1.uri = f"{A2UI_EXTENSION_BASE_URI}/v0.8.0"
+  ext2 = MagicMock()
+  ext2.uri = f"{A2UI_EXTENSION_BASE_URI}/v1.0.0"
+  ext3 = MagicMock()
+  ext3.uri = f"{A2UI_EXTENSION_BASE_URI}/v1.2.0"
+  card.capabilities.extensions = [ext1, ext2, ext3]
+
+  assert try_activate_a2ui_extension(context, card) == "1.0.0"
+  context.add_activated_extension.assert_called_once_with(
+      f"{A2UI_EXTENSION_BASE_URI}/v1.0.0"
+  )