Remove notion of "custom catalogs" from agent SDK

The basic catalog maintained by the A2UI team has no difference from
third-party catalogs.

This PR removes the notion of custom catalogs. Each catalog provided at
runtime should be independent and immutable. At build time, catalogs can
refer to components from other catalogs. They need to be bundled into a
free-standing one using the `tools/build_catalog/build_catalog.py`
script.

Fixes #650

Author: nan-yu

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

@@ -94,7 +94,7 @@ def get_a2ui_agent_extension(
   """Creates the A2UI AgentExtension configuration.
 
   Args:
-      accepts_inline_catalogs: Whether the agent accepts inline custom catalogs.
+      accepts_inline_catalogs: Whether the agent accepts inline catalogs.
       supported_catalog_ids: All pre-defined catalogs the agent is known to support.
 
   Returns:

a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/catalog.py

@@ -19,7 +19,7 @@
 from dataclasses import dataclass, field, replace
 from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
-from .constants import CATALOG_COMPONENTS_KEY, CATALOG_ID_KEY
+from .constants import CATALOG_COMPONENTS_KEY, CATALOG_ID_KEY, BASIC_CATALOG_NAME
 from referencing import Registry, Resource
 
 if TYPE_CHECKING:
@@ -28,13 +28,35 @@
 
 
 @dataclass
-class CustomCatalogConfig:
-  """Configuration for a custom component catalog."""
+class CatalogConfig:
+  """
+  Configuration for a catalog of components.
+
+  The catalog must be free standing and not reference any other catalogs, except
+  for the common types schema.
+
+  If a catalog references other catalogs, the references must be resolved before
+  loading the catalog by using `tools/build_catalog/build_catalog.py` script.
+
+  Attributes:
+    name: The name of the catalog.
+    catalog_path: The path to the catalog JSON file.
+    examples_path: The path to the examples directory.
+  """
 
   name: str
   catalog_path: str
   examples_path: Optional[str] = None
 
+  @classmethod
+  def for_basic_catalog(cls, examples_path: Optional[str] = None) -> "CatalogConfig":
+    """Returns a CatalogConfig for the basic catalog."""
+    return cls(
+        name=BASIC_CATALOG_NAME,
+        catalog_path="",  # basic catalog doesn't need a path from the user
+        examples_path=examples_path,
+    )
+
 
 @dataclass(frozen=True)
 class A2uiCatalog:
@@ -165,113 +187,6 @@ def load_examples(self, path: Optional[str], validate: bool = False) -> str:
       return ""
     return "\n\n".join(merged_examples)
 
-  @staticmethod
-  def resolve_schema(basic: Dict[str, Any], custom: Dict[str, Any]) -> Dict[str, Any]:
-    """Resolves references in custom catalog schema against the basic catalog.
-
-    Args:
-      basic: The basic catalog schema.
-      custom: The custom catalog schema.
-
-    Returns:
-      A new dictionary with references resolved.
-    """
-    result = copy.deepcopy(custom)
-
-    # Initialize registry with basic catalog and maybe others from basic's $id
-    registry = Registry()
-    if CATALOG_ID_KEY in basic:
-      basic_resource = Resource.from_contents(basic)
-      registry = registry.with_resource(basic[CATALOG_ID_KEY], basic_resource)
-
-    def resolve_ref(ref_uri: str) -> Any:
-      try:
-        resolver = registry.resolver()
-        resolved = resolver.lookup(ref_uri)
-        return resolved.contents
-      except Exception as e:
-        logging.warning("Could not resolve reference %s: %s", ref_uri, e)
-        return None
-
-    def merge_into(target: Dict[str, Any], source: Dict[str, Any]):
-      for key, value in source.items():
-        if key not in target:
-          target[key] = copy.deepcopy(value)
-
-    # Process components
-    if CATALOG_COMPONENTS_KEY in result:
-      comp_dict = result[CATALOG_COMPONENTS_KEY]
-      if "$ref" in comp_dict:
-        resolved = resolve_ref(comp_dict["$ref"])
-        if isinstance(resolved, dict):
-          merge_into(comp_dict, resolved)
-        del comp_dict["$ref"]
-
-    # Process functions
-    if "functions" in result:
-      func_dict = result["functions"]
-      if "$ref" in func_dict:
-        resolved = resolve_ref(func_dict["$ref"])
-        if isinstance(resolved, dict):
-          merge_into(func_dict, resolved)
-        del func_dict["$ref"]
-
-    # Process $defs
-    if "$defs" in result:
-      res_defs = result["$defs"]
-      if "$ref" in res_defs:
-        resolved = resolve_ref(res_defs["$ref"])
-        if isinstance(resolved, dict):
-          merge_into(res_defs, resolved)
-        del res_defs["$ref"]
-
-      for name in ["anyComponent", "anyFunction"]:
-        if name in res_defs:
-          target = res_defs[name]
-          one_of = target.get("oneOf", [])
-          new_one_of = []
-          for item in one_of:
-            if isinstance(item, dict) and "$ref" in item:
-              ref_uri = item["$ref"]
-              # Check if it points to basic collector
-              resolved = resolve_ref(ref_uri)
-              if isinstance(resolved, dict) and "oneOf" in resolved:
-                # Merge oneOf items and resolve transitive refs to components/functions
-                for sub_item in resolved["oneOf"]:
-                  if sub_item not in new_one_of:
-                    new_one_of.append(copy.deepcopy(sub_item))
-                    # Transitive merge: if sub_item is a ref to a component/function
-                    if isinstance(sub_item, dict) and "$ref" in sub_item:
-                      sub_ref = sub_item["$ref"]
-                      if (
-                          sub_ref.startswith("#/components/")
-                          and CATALOG_COMPONENTS_KEY in basic
-                      ):
-                        comp_name = sub_ref.split("/")[-1]
-                        if comp_name in basic[CATALOG_COMPONENTS_KEY]:
-                          if CATALOG_COMPONENTS_KEY not in result:
-                            result[CATALOG_COMPONENTS_KEY] = {}
-                          if comp_name not in result[CATALOG_COMPONENTS_KEY]:
-                            result[CATALOG_COMPONENTS_KEY][comp_name] = copy.deepcopy(
-                                basic[CATALOG_COMPONENTS_KEY][comp_name]
-                            )
-                      elif sub_ref.startswith("#/functions/") and "functions" in basic:
-                        func_name = sub_ref.split("/")[-1]
-                        if func_name in basic["functions"]:
-                          if "functions" not in result:
-                            result["functions"] = {}
-                          if func_name not in result["functions"]:
-                            result["functions"][func_name] = copy.deepcopy(
-                                basic["functions"][func_name]
-                            )
-              else:
-                new_one_of.append(item)
-            else:
-              new_one_of.append(item)
-          target["oneOf"] = new_one_of
-
-    return result
-
   def _validate_example(self, full_path: str, basename: str, content: str) -> bool:
     try:
       json_data = json.loads(content)

a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/manager.py

@@ -35,7 +35,7 @@
     BASIC_CATALOG_NAME,
     find_repo_root,
 )
-from .catalog import CustomCatalogConfig, A2uiCatalog
+from .catalog import CatalogConfig, A2uiCatalog
 from ...extension.a2ui_extension import INLINE_CATALOGS_KEY, SUPPORTED_CATALOG_IDS_KEY, get_a2ui_agent_extension
 from a2a.types import AgentExtension
 
@@ -118,30 +118,26 @@ class A2uiSchemaManager(InferenceStrategy):
   def __init__(
       self,
       version: str,
-      basic_examples_path: Optional[str] = None,
-      custom_catalogs: Optional[List[CustomCatalogConfig]] = None,
-      exclude_basic_catalog: bool = False,
+      catalogs: List[CatalogConfig] = [],
       accepts_inline_catalogs: bool = False,
-      schema_modifiers: List[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
+      schema_modifiers: List[Callable[[Dict[str, Any]], Dict[str, Any]]] = [],
   ):
     self._version = version
-    self._exclude_basic_catalog = exclude_basic_catalog
     self._accepts_inline_catalogs = accepts_inline_catalogs
 
     self._server_to_client_schema = None
     self._common_types_schema = None
-    self._supported_catalogs: Dict[str, A2uiCatalog] = {}
+    self._supported_catalogs: List[A2uiCatalog] = []
     self._catalog_example_paths: Dict[str, str] = {}
-    self._basic_catalog = None
     self._schema_modifiers = schema_modifiers
-    self._load_schemas(version, custom_catalogs, basic_examples_path)
+    self._load_schemas(version, catalogs)
 
   @property
   def accepts_inline_catalogs(self) -> bool:
     return self._accepts_inline_catalogs
 
   @property
-  def supported_catalogs(self) -> Dict[str, A2uiCatalog]:
+  def supported_catalogs(self) -> List[A2uiCatalog]:
     return self._supported_catalogs
 
   def _apply_modifiers(self, schema: Dict[str, Any]) -> Dict[str, Any]:
@@ -153,8 +149,7 @@ def _apply_modifiers(self, schema: Dict[str, Any]) -> Dict[str, Any]:
   def _load_schemas(
       self,
       version: str,
-      custom_catalogs: Optional[List[CustomCatalogConfig]] = None,
-      basic_examples_path: Optional[str] = None,
+      catalogs: List[CatalogConfig] = [],
   ):
     """Loads separate schema components and processes catalogs."""
     if version not in SPEC_VERSION_MAP:
@@ -172,11 +167,7 @@ def _load_schemas(
     )
 
     # Process basic catalog
-    basic_catalog_schema = self._apply_modifiers(
-        _load_basic_component(version, CATALOG_SCHEMA_KEY)
-    )
-    if not basic_catalog_schema:
-      basic_catalog_schema = {}
+    basic_catalog_schema = _load_basic_component(version, CATALOG_SCHEMA_KEY)
 
     # Ensure catalog id and schema url are set in the basic catalog schema
     if CATALOG_ID_KEY not in basic_catalog_schema:
@@ -190,44 +181,36 @@ def _load_schemas(
     if "$schema" not in basic_catalog_schema:
       basic_catalog_schema["$schema"] = "https://json-schema.org/draft/2020-12/schema"
 
-    self._basic_catalog = A2uiCatalog(
-        version=version,
-        name=BASIC_CATALOG_NAME,
-        catalog_schema=basic_catalog_schema,
-        s2c_schema=self._server_to_client_schema,
-        common_types_schema=self._common_types_schema,
-    )
-    if not self._exclude_basic_catalog:
-      self._supported_catalogs[self._basic_catalog.catalog_id] = self._basic_catalog
-      self._catalog_example_paths[self._basic_catalog.catalog_id] = basic_examples_path
-
-    # Process custom catalogs
-    if custom_catalogs:
-      for config in custom_catalogs:
-        custom_catalog_schema = self._apply_modifiers(
-            _load_from_path(config.catalog_path)
-        )
-        resolved_catalog_schema = A2uiCatalog.resolve_schema(
-            basic_catalog_schema, custom_catalog_schema
-        )
-        catalog = A2uiCatalog(
-            version=version,
-            name=config.name,
-            catalog_schema=self._apply_modifiers(resolved_catalog_schema),
-            s2c_schema=self._server_to_client_schema,
-            common_types_schema=self._common_types_schema,
-        )
-        self._supported_catalogs[catalog.catalog_id] = catalog
-        self._catalog_example_paths[catalog.catalog_id] = config.examples_path
+    # Process catalogs
+    if not catalogs:
+      # If no catalogs are provided, use the basic catalog
+      catalogs = [CatalogConfig.for_basic_catalog()]
+
+    for config in catalogs:
+      catalog_schema = (
+          basic_catalog_schema
+          if config.name == BASIC_CATALOG_NAME
+          else _load_from_path(config.catalog_path)
+      )
+      catalog_schema = self._apply_modifiers(catalog_schema)
+      catalog = A2uiCatalog(
+          version=version,
+          name=config.name,
+          catalog_schema=catalog_schema,
+          s2c_schema=self._server_to_client_schema,
+          common_types_schema=self._common_types_schema,
+      )
+      self._supported_catalogs.append(catalog)
+      self._catalog_example_paths[catalog.catalog_id] = config.examples_path
 
   def _determine_catalog(
       self, client_ui_capabilities: Optional[dict[str, Any]] = None
   ) -> A2uiCatalog:
     """Determines the catalog to use based on supported catalog IDs.
 
-    If neither inline catalogs nor supported catalog IDs are provided, the basic catalog is used.
+    If neither inline catalogs nor client-supported catalog IDs are provided, the first agent-supported catalog is used.
     If inline catalogs are provided, the first inline catalog is used.
-    If supported catalog IDs are provided, the first supported catalog that is recognized is used.
+    If client-supported catalog IDs are provided, the first one that is supported by the agent is used.
 
     Args:
       client_ui_capabilities: A dictionary of client UI capabilities.
@@ -236,16 +219,19 @@ def _determine_catalog(
       The A2uiCatalog to use to generate the schema string in the prompt.
 
     Raises:
-      ValueError: If both inline catalogs and supported catalog IDs are provided,
+      ValueError: If both inline catalogs and client-supported catalog IDs are provided,
         or if no supported catalog is recognized.
     """
+    if not self._supported_catalogs:
+      raise ValueError("No supported catalogs found.")  # This should not happen.
+
     if not client_ui_capabilities or not isinstance(client_ui_capabilities, dict):
-      return self._basic_catalog
+      return self._supported_catalogs[0]
 
     inline_catalogs: List[dict[str, Any]] = client_ui_capabilities.get(
         INLINE_CATALOGS_KEY, []
     )
-    supported_catalog_ids: List[str] = client_ui_capabilities.get(
+    client_supported_catalog_ids: List[str] = client_ui_capabilities.get(
         SUPPORTED_CATALOG_IDS_KEY, []
     )
 
@@ -255,37 +241,35 @@ def _determine_catalog(
           " capabilities. However, the agent does not accept inline catalogs."
       )
 
-    if inline_catalogs and supported_catalog_ids:
+    if inline_catalogs and client_supported_catalog_ids:
       raise ValueError(
           f"Both '{INLINE_CATALOGS_KEY}' and '{SUPPORTED_CATALOG_IDS_KEY}' "
           "are provided in client UI capabilities. Only one is allowed."
       )
 
     if inline_catalogs:
-      # Load the first custom inline catalog schema.
+      # Load the first inline catalog schema.
       inline_catalog_schema = inline_catalogs[0]
-      resolved_catalog_schema = A2uiCatalog.resolve_schema(
-          self._basic_catalog.catalog_schema, inline_catalog_schema
-      )
+      inline_catalog_schema = self._apply_modifiers(inline_catalog_schema)
       return A2uiCatalog(
           version=self._version,
           name=INLINE_CATALOG_NAME,
-          catalog_schema=resolved_catalog_schema,
+          catalog_schema=inline_catalog_schema,
           s2c_schema=self._server_to_client_schema,
           common_types_schema=self._common_types_schema,
       )
 
-    if not supported_catalog_ids:
-      return self._basic_catalog
+    if not client_supported_catalog_ids:
+      return self._supported_catalogs[0]
 
-    for scid in supported_catalog_ids:
-      if scid in self._supported_catalogs:
-        # Return the first supported catalog.
-        return self._supported_catalogs[scid]
+    agent_supported_catalogs = {c.catalog_id: c for c in self._supported_catalogs}
+    for cscid in client_supported_catalog_ids:
+      if cscid in agent_supported_catalogs:
+        return agent_supported_catalogs[cscid]
 
     raise ValueError(
-        "No supported catalog found on the agent side. Agent supported catalogs are:"
-        f" {list(self._supported_catalogs.keys())}"
+        "No client-supported catalog found on the agent side. Agent-supported catalogs"
+        f" are: {[c.catalog_id for c in self._supported_catalogs]}"
     )
 
   def get_effective_catalog(
@@ -339,5 +323,5 @@ def generate_system_prompt(
     return "\n\n".join(parts)
 
   def get_agent_extension(self) -> AgentExtension:
-    catalog_ids = self._supported_catalogs.keys()
-    return get_a2ui_agent_extension(supported_catalog_ids=list(catalog_ids))
+    catalog_ids = [c.catalog_id for c in self._supported_catalogs]
+    return get_a2ui_agent_extension(supported_catalog_ids=catalog_ids)

a2a_agents/python/a2ui_agent/tests/inference/test_modifiers.py

@@ -58,7 +58,7 @@ def test_manager_with_modifiers():
     assert "additionalProperties" not in manager._common_types_schema
 
     # basic catalog should also be modified
-    for catalog in manager._supported_catalogs.values():
+    for catalog in manager._supported_catalogs:
       assert "additionalProperties" not in catalog.catalog_schema