feat: Introduce schema bundling in A2uiSchemaManager and add a validator

Author: nan-yu

a2a_agents/python/a2ui_agent/pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
 ]
 
 [build-system]
-requires = ["hatchling"]
+requires = ["hatchling", "jsonschema"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]

a2a_agents/python/a2ui_agent/src/a2ui/core/validator.py

@@ -0,0 +1,32 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict
+import jsonschema
+
+
+class A2uiValidator:
+  """
+  Validates JSON instances against a bundled A2UI schema.
+  """
+
+  def __init__(self, schema: Dict[str, Any]):
+    self._validator = jsonschema.validators.validator_for(schema)(schema)
+
+  def validate(self, instance: Any) -> None:
+    """
+    Validates the given instance against the bundled schema.
+    Raises jsonschema.exceptions.ValidationError if validation fails.
+    """
+    self._validator.validate(instance)

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

@@ -17,9 +17,10 @@
 import logging
 import os
 import importlib.resources
-from typing import List, Dict, Any
-from .loader import A2uiSchemaLoader, PackageLoader, FileSystemLoader
+from typing import List, Dict, Any, Tuple, Set
 from ..inference_strategy import InferenceStrategy
+from .loader import A2uiSchemaLoader, PackageLoader, FileSystemLoader
+from a2ui.core.validator import A2uiValidator
 
 # Helper constants for schema management shared by build hook and runtime
 
@@ -30,6 +31,7 @@
 COMMON_TYPES_SCHEMA_KEY = "common_types"
 CATALOG_SCHEMA_KEY = "catalog"
 CATALOG_COMPONENTS_KEY = "components"
+CATALOG_STYLES_KEY = "styles"
 
 SPEC_VERSION_MAP = {
     "0.8": {
@@ -55,6 +57,16 @@ def __init__(self, version: str, custom_catalog_path: str = None):
     self.catalog_schema = None
     self.common_types_schema = None
     self._load_schemas(version, custom_catalog_path)
+    self._bundled_schema = self.bundle_schemas()
+    self._validator = A2uiValidator(self._bundled_schema)
+
+  @property
+  def bundled_schema(self):
+    return self._bundled_schema
+
+  @property
+  def validator(self):
+    return self._validator
 
   def _load_schemas(self, version: str, custom_catalog_path: str = None):
     """
@@ -220,6 +232,173 @@ def generate_system_prompt(
 
     return "\n\n".join(parts)
 
+  def bundle_schemas(self) -> Dict[str, Any]:
+    """
+    Bundles the loaded schemas into a single self-contained schema.
+    Returns:
+        A dictionary representing the bundled schema.
+    """
+    bundled = []
+    if self.version == "0.8":
+      bundled = self._bundle_0_8(
+          self.server_to_client_schema,
+          self.catalog_schema,
+          self.common_types_schema,
+      )
+    else:
+      # Default to 0.9+ behavior (using $defs)
+      bundled = self._bundle_0_9(
+          self.server_to_client_schema,
+          self.catalog_schema,
+          self.common_types_schema,
+      )
+    # LLM is instructed to generate a list of messages, so we wrap the bundled schema in an array.
+    return {
+        "type": "array",
+        "items": bundled,
+    }
+
+  def _inject_additional_properties(
+      self,
+      schema: Dict[str, Any],
+      source_properties: Dict[str, Any],
+      mapping: Dict[str, str] = None,
+  ) -> Tuple[Dict[str, Any], Set[str]]:
+    """
+    Recursively injects properties from source_properties into nodes with additionalProperties=True.
+    Args:
+        schema: The target schema to traverse and patch.
+        source_properties: A dictionary of top-level property groups (e.g., "components", "styles") from the source schema.
+    Returns:
+        A tuple containing:
+        - The patched schema.
+        - A set of keys from source_properties that were injected.
+    """
+    injected_keys = set()
+
+    def recursive_inject(obj):
+      if isinstance(obj, dict):
+        new_obj = {}
+        for k, v in obj.items():
+          if isinstance(v, dict) and v.get("additionalProperties") is True:
+            if k in source_properties:
+              injected_keys.add(k)
+              new_node = dict(v)
+              new_node["additionalProperties"] = False
+              new_node["properties"] = {
+                  **new_node.get("properties", {}),
+                  **source_properties[k],
+              }
+              new_obj[k] = new_node
+            else:  # No matching source group, keep as is but recurse children
+              new_obj[k] = recursive_inject(v)
+          else:  # Not a node with additionalProperties, recurse children
+            new_obj[k] = recursive_inject(v)
+        return new_obj
+      elif isinstance(obj, list):
+        return [recursive_inject(i) for i in obj]
+      return obj
+
+    return recursive_inject(schema), injected_keys
+
+  def _bundle_0_9(
+      self,
+      server_to_client: Dict[str, Any],
+      catalog: Dict[str, Any],
+      common: Dict[str, Any] = None,
+  ) -> Dict[str, Any]:
+    if not server_to_client:
+      return {}
+
+    bundled = copy.deepcopy(server_to_client)
+
+    # Collect source properties and merge schemas
+    source_properties = {}
+    schemas_to_merge = []
+    # merged in order so catalog overrides common.
+    if common:
+      schemas_to_merge.append(common)
+    if catalog:
+      schemas_to_merge.append(catalog)
+
+    if "$defs" not in bundled:
+      bundled["$defs"] = {}
+
+    for schema in schemas_to_merge:
+      source_properties.update(schema)
+
+      # Merge $defs
+      if "$defs" in schema:
+        bundled["$defs"].update(schema["$defs"])
+
+      # Merge other top-level properties
+      for k, v in schema.items():
+        if k not in ["$schema", "title", "$id", "description", "$defs"]:
+          bundled[k] = v
+
+    bundled["$id"] = "https://a2ui.dev/specification/0.9/server_to_client_bundled.json"
+    bundled["title"] = "Bundled A2UI Message Schema"
+    bundled["description"] = (
+        "A self-contained bundled schema including server messages, catalog components,"
+        " and common types."
+    )
+
+    bundled, injected_keys = self._inject_additional_properties(
+        bundled, source_properties
+    )
+
+    # Cleanup injected keys from bundled root to avoid duplication
+    for k in injected_keys:
+      if k in bundled:
+        del bundled[k]
+
+    # Recursively strip external file references
+    def rewrite_refs(obj):
+      if isinstance(obj, dict):
+        new_obj = {}
+        for k, v in obj.items():
+          if k == "$ref" and isinstance(v, str):
+            if ".json" in v:
+              ref_parts = v.split(".json")
+              fragment = ref_parts[-1]
+              if not fragment:
+                fragment = "#"
+              new_obj[k] = fragment
+            else:
+              new_obj[k] = v
+          else:
+            new_obj[k] = rewrite_refs(v)
+        return new_obj
+      elif isinstance(obj, list):
+        return [rewrite_refs(i) for i in obj]
+      return obj
+
+    return rewrite_refs(bundled)
+
+  def _bundle_0_8(
+      self,
+      server_to_client: Dict[str, Any],
+      catalog: Dict[str, Any],
+      common: Dict[str, Any] = None,
+  ) -> Dict[str, Any]:
+    if not server_to_client:
+      return {}
+
+    bundled = copy.deepcopy(server_to_client)
+
+    # Prepare catalog components and styles for injection
+    source_properties = {}
+    if catalog:
+      if CATALOG_COMPONENTS_KEY in catalog:
+        # Special mapping for v0.8: "components" -> "component"
+        source_properties["component"] = catalog[CATALOG_COMPONENTS_KEY]
+      if CATALOG_STYLES_KEY in catalog:
+        source_properties[CATALOG_STYLES_KEY] = catalog[CATALOG_STYLES_KEY]
+
+    bundled, _ = self._inject_additional_properties(bundled, source_properties)
+
+    return bundled
+
 
 def find_repo_root(start_path: str) -> str | None:
   """Finds the repository root by looking for the 'specification' directory."""

a2a_agents/python/a2ui_agent/tests/core/test_validator.py

@@ -0,0 +1,44 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import jsonschema
+from a2ui.core.validator import A2uiValidator
+
+
+@pytest.fixture
+def mock_bundled_schema():
+  return {
+      "type": "array",
+      "items": {
+          "$schema": "https://json-schema.org/draft/2020-12/schema",
+          "type": "object",
+          "properties": {"foo": {"type": "string"}},
+          "required": ["foo"],
+      },
+  }
+
+
+def test_validation_success(mock_bundled_schema):
+  validator = A2uiValidator(mock_bundled_schema)
+  instance = [{"foo": "bar"}]
+  validator.validate(instance)
+
+
+def test_validation_failure(mock_bundled_schema):
+  validator = A2uiValidator(mock_bundled_schema)
+  instance = [{"foo": 123}]  # Wrong type
+
+  with pytest.raises(jsonschema.exceptions.ValidationError):
+    validator.validate(instance)