Treeshake common types schema

Author: jiahaog

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

@@ -102,7 +102,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 +132,54 @@ 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()
+    queue = []
+
+    # Initialize queue with refs from catalog_schema and s2c_schema
+    queue.extend(_collect_refs(self.catalog_schema))
+    queue.extend(_collect_refs(self.s2c_schema))
+
+    while queue:
+      ref = queue.pop(0)
+      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.
+        def_name = ref.split("/")[-1]
+        if (
+            def_name in self.common_types_schema["$defs"]
+            and def_name not in visited_defs
+        ):
+          visited_defs.add(def_name)
+          queue.extend(_collect_refs(self.common_types_schema["$defs"][def_name]))
+
+    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."""

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

@@ -209,3 +209,32 @@ def test_render_as_llm_instructions():
   assert '"catalog": "schema"' in schema_str
   assert '"catalogId": "id_basic"' in schema_str
   assert A2UI_SCHEMA_BLOCK_END in schema_str
+
+
+def test_with_pruned_components_prunes_common_types():
+  common_types = {
+      "$defs": {
+          "TypeForCompA": {"type": "string"},
+          "TypeForCompB": {"type": "number"},
+      }
+  }
+  catalog_schema = {
+      "catalogId": "basic",
+      "components": {
+          "CompA": {"$ref": "#/$defs/TypeForCompA"},
+          "CompB": {"$ref": "#/$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