Treeshake common types schema

agent_sdks/python/src/a2ui/core/schema/catalog.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import collections
 import copy
 import json
 import logging
@@ -102,7 +103,7 @@ def with_pruned_components(self, allowed_components: List[str]) -> "A2uiCatalog"
 
     # Allow all components if no allowed components are specified
     if not allowed_components:
-      return self
+      return self._with_pruned_common_types()
 
     if CATALOG_COMPONENTS_KEY in schema_copy and isinstance(
         schema_copy[CATALOG_COMPONENTS_KEY], dict
@@ -132,7 +133,58 @@ def with_pruned_components(self, allowed_components: List[str]) -> "A2uiCatalog"
 
         any_comp["oneOf"] = filtered_one_of
 
-    return replace(self, catalog_schema=schema_copy)
+    pruned_catalog = replace(self, catalog_schema=schema_copy)
+    return pruned_catalog._with_pruned_common_types()
+
+  def _with_pruned_common_types(self) -> "A2uiCatalog":
+    """Returns a new catalog with unused common types pruned from the schema."""
+    if not self.common_types_schema or "$defs" not in self.common_types_schema:
+      return self
+
+    def _collect_refs(obj: Any) -> set[str]:
+      refs = set()
+      if isinstance(obj, dict):
+        for k, v in obj.items():
+          if k == "$ref" and isinstance(v, str):
+            refs.add(v)
+          else:
+            refs.update(_collect_refs(v))
+      elif isinstance(obj, list):
+        for item in obj:
+          refs.update(_collect_refs(item))
+      return refs
+
+    visited_defs = set()
+    internal_refs_queue = collections.deque()
+
+    # Initialize queue with ONLY refs targeting common_types.json from external schemas
+    external_refs = _collect_refs(self.catalog_schema)
+    external_refs.update(_collect_refs(self.s2c_schema))
+
+    for ref in external_refs:
+      if ref.startswith("common_types.json#/$defs/"):
+        internal_refs_queue.append(ref.split("#/$defs/")[-1])
+
+    while internal_refs_queue:
+      def_name = internal_refs_queue.popleft()
+      if def_name in self.common_types_schema["$defs"] and def_name not in visited_defs:
+        visited_defs.add(def_name)
+
+        # Collect internal references (which just use #/$defs/)
+        internal_refs = _collect_refs(self.common_types_schema["$defs"][def_name])
+        for ref in internal_refs:
+          if ref.startswith("#/$defs/"):
+            # Note: This assumes a flat `$defs` namespace and no escaped
+            # slashes (~1) or tildes (~0) in the definition names as per RFC 6901.
+            internal_refs_queue.append(ref.split("#/$defs/")[-1])
+
+    new_common_types_schema = copy.deepcopy(self.common_types_schema)
+    all_defs = new_common_types_schema["$defs"]
+    new_common_types_schema["$defs"] = {
+        k: v for k, v in all_defs.items() if k in visited_defs
+    }
+
+    return replace(self, common_types_schema=new_common_types_schema)
 
   def render_as_llm_instructions(self) -> str:
     """Renders the catalog and schema as LLM instructions."""
@@ -144,7 +196,11 @@ def render_as_llm_instructions(self) -> str:
     )
     all_schemas.append(f"### Server To Client Schema:\n{server_client_str}")
 
-    if self.common_types_schema:
+    if (
+        self.common_types_schema
+        and "$defs" in self.common_types_schema
+        and self.common_types_schema["$defs"]
+    ):
       common_str = json.dumps(self.common_types_schema, indent=2)
       all_schemas.append(f"### Common Types Schema:\n{common_str}")
 

agent_sdks/python/tests/core/schema/test_catalog.py

@@ -193,7 +193,7 @@ def test_render_as_llm_instructions():
       version=VERSION_0_9,
       name=BASIC_CATALOG_NAME,
       s2c_schema={"s2c": "schema"},
-      common_types_schema={"common": "types"},
+      common_types_schema={"$defs": {"common": "types"}},
       catalog_schema={
           "$schema": "https://json-schema.org/draft/2020-12/schema",
           "catalog": "schema",
@@ -204,8 +204,87 @@ def test_render_as_llm_instructions():
   schema_str = catalog.render_as_llm_instructions()
   assert A2UI_SCHEMA_BLOCK_START in schema_str
   assert '### Server To Client Schema:\n{\n  "s2c": "schema"\n}' in schema_str
-  assert '### Common Types Schema:\n{\n  "common": "types"\n}' in schema_str
+  assert (
+      '### Common Types Schema:\n{\n  "$defs": {\n    "common": "types"\n  }\n}'
+      in schema_str
+  )
   assert "### Catalog Schema:" in schema_str
   assert '"catalog": "schema"' in schema_str
   assert '"catalogId": "id_basic"' in schema_str
   assert A2UI_SCHEMA_BLOCK_END in schema_str
+
+
+def test_render_as_llm_instructions_drops_empty_common_types():
+  # Test with empty common_types_schema
+  catalog_empty = A2uiCatalog(
+      version=VERSION_0_9,
+      name=BASIC_CATALOG_NAME,
+      s2c_schema={"s2c": "schema"},
+      common_types_schema={},
+      catalog_schema={
+          "$schema": "https://json-schema.org/draft/2020-12/schema",
+          "catalog": "schema",
+          "catalogId": "id_basic",
+      },
+  )
+  schema_str_empty = catalog_empty.render_as_llm_instructions()
+  assert "### Common Types Schema:" not in schema_str_empty
+
+  # Test with common_types_schema missing $defs
+  catalog_no_defs = A2uiCatalog(
+      version=VERSION_0_9,
+      name=BASIC_CATALOG_NAME,
+      s2c_schema={"s2c": "schema"},
+      common_types_schema={"something": "else"},
+      catalog_schema={
+          "$schema": "https://json-schema.org/draft/2020-12/schema",
+          "catalog": "schema",
+          "catalogId": "id_basic",
+      },
+  )
+  schema_str_no_defs = catalog_no_defs.render_as_llm_instructions()
+  assert "### Common Types Schema:" not in schema_str_no_defs
+
+  # Test with common_types_schema having empty $defs
+  catalog_empty_defs = A2uiCatalog(
+      version=VERSION_0_9,
+      name=BASIC_CATALOG_NAME,
+      s2c_schema={"s2c": "schema"},
+      common_types_schema={"$defs": {}},
+      catalog_schema={
+          "$schema": "https://json-schema.org/draft/2020-12/schema",
+          "catalog": "schema",
+          "catalogId": "id_basic",
+      },
+  )
+  schema_str_empty_defs = catalog_empty_defs.render_as_llm_instructions()
+  assert "### Common Types Schema:" not in schema_str_empty_defs
+
+
+def test_with_pruned_components_prunes_common_types():
+  common_types = {
+      "$defs": {
+          "TypeForCompA": {"type": "string"},
+          "TypeForCompB": {"type": "number"},
+      }
+  }
+  catalog_schema = {
+      "catalogId": "basic",
+      "components": {
+          "CompA": {"$ref": "common_types.json#/$defs/TypeForCompA"},
+          "CompB": {"$ref": "common_types.json#/$defs/TypeForCompB"},
+      },
+  }
+  catalog = A2uiCatalog(
+      version=VERSION_0_8,
+      name=BASIC_CATALOG_NAME,
+      s2c_schema={},
+      common_types_schema=common_types,
+      catalog_schema=catalog_schema,
+  )
+
+  pruned_catalog = catalog.with_pruned_components(["CompA"])
+  pruned_defs = pruned_catalog.common_types_schema["$defs"]
+
+  assert "TypeForCompA" in pruned_defs
+  assert "TypeForCompB" not in pruned_defs

agent_sdks/python/tests/core/schema/test_schema_manager.py

@@ -279,7 +279,7 @@ def joinpath_side_effect(path):
     content = '{"$schema": "https://json-schema.org/draft/2020-12/schema"}'
     if path == "common_types.json":
       content = (
-          '{"$schema": "https://json-schema.org/draft/2020-12/schema", "types":'
+          '{"$schema": "https://json-schema.org/draft/2020-12/schema", "$defs":'
           ' {"Common": {}}}'
       )
     elif "server_to_client" in path:
@@ -290,9 +290,9 @@ def joinpath_side_effect(path):
     elif "catalog" in path:
       content = (
           '{"$schema": "https://json-schema.org/draft/2020-12/schema", "catalogId":'
-          ' "basic", "components": {}}'
+          ' "basic", "components": {}, "$defs": {"test": {"$ref":'
+          ' "common_types.json#/$defs/Common"}}}'
       )
-
     mock_file.open.return_value.__enter__.return_value = io.StringIO(content)
     return mock_file
 
@@ -306,7 +306,7 @@ def joinpath_side_effect(path):
   prompt = manager.generate_system_prompt("Role", include_schema=True)
 
   assert "### Common Types Schema:" in prompt
-  assert '"types":{"Common":{}}' in prompt.replace(" ", "").replace("\n", "")
+  assert '"$defs":{"Common":{}}' in prompt.replace(" ", "").replace("\n", "")
 
 
 def test_generate_system_prompt_minimal_args(mock_importlib_resources):