feat: Add a A2uiValidator to validate loaded examples

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/inference/inference_strategy.py

@@ -28,6 +28,7 @@ def generate_system_prompt(
       allowed_components: List[str] = [],
       include_schema: bool = False,
       include_examples: bool = False,
+      validate_examples: bool = False,
   ) -> str:
     """
     Generates a system prompt for all LLM requests.
@@ -40,6 +41,7 @@ def generate_system_prompt(
       allowed_components: List of allowed components.
       include_schema: Whether to include the schema.
       include_examples: Whether to include examples.
+      validate_examples: Whether to validate examples.
 
     Returns:
       The system prompt.

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

@@ -16,12 +16,15 @@
 import json
 import logging
 import os
-from dataclasses import dataclass, replace
-from typing import Any, Dict, List, Optional
+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 referencing import Registry, Resource
 
+if TYPE_CHECKING:
+  from .validator import A2uiValidator
+
 
 @dataclass
 class CustomCatalogConfig:
@@ -56,6 +59,12 @@ def catalog_id(self) -> str:
       raise ValueError(f"Catalog '{self.name}' missing catalogId")
     return self.catalog_schema[CATALOG_ID_KEY]
 
+  @property
+  def validator(self) -> "A2uiValidator":
+    from .validator import A2uiValidator
+
+    return A2uiValidator(self)
+
   def with_pruned_components(self, allowed_components: List[str]) -> "A2uiCatalog":
     """Returns a new catalog with only allowed components.
 
@@ -123,7 +132,7 @@ def render_as_llm_instructions(self) -> str:
 
     return "\n\n".join(all_schemas)
 
-  def load_examples(self, path: Optional[str]) -> str:
+  def load_examples(self, path: Optional[str], validate: bool = False) -> str:
     """Loads and validates examples from a directory."""
     if not path or not os.path.isdir(path):
       if path:
@@ -138,6 +147,8 @@ def load_examples(self, path: Optional[str]) -> str:
         try:
           with open(full_path, "r", encoding="utf-8") as f:
             content = f.read()
+            if validate and not self._validate_example(full_path, basename, content):
+              continue
             merged_examples.append(
                 f"---BEGIN {basename}---\n{content}\n---END {basename}---"
             )
@@ -253,3 +264,12 @@ def merge_into(target: Dict[str, Any], source: Dict[str, Any]):
           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)
+      self.validator.validate(json_data)
+    except Exception as e:
+      logging.warning(f"Failed to validate example {full_path}: {e}")
+      return False
+    return True

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

@@ -20,6 +20,7 @@
 CATALOG_SCHEMA_KEY = "catalog"
 CATALOG_COMPONENTS_KEY = "components"
 CATALOG_ID_KEY = "catalogId"
+CATALOG_STYLES_KEY = "styles"
 
 BASE_SCHEMA_URL = "https://a2ui.org/"
 BASIC_CATALOG_NAME = "basic"

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

@@ -283,10 +283,12 @@ def get_effective_catalog(
     pruned_catalog = catalog.with_pruned_components(allowed_components)
     return pruned_catalog
 
-  def load_examples(self, catalog: A2uiCatalog) -> str:
+  def load_examples(self, catalog: A2uiCatalog, validate: bool = False) -> str:
     """Loads examples for a catalog."""
     if catalog.catalog_id in self._catalog_example_paths:
-      return catalog.load_examples(self._catalog_example_paths[catalog.catalog_id])
+      return catalog.load_examples(
+          self._catalog_example_paths[catalog.catalog_id], validate=validate
+      )
     return ""
 
   def generate_system_prompt(
@@ -298,6 +300,7 @@ def generate_system_prompt(
       allowed_components: List[str] = [],
       include_schema: bool = False,
       include_examples: bool = False,
+      validate_examples: bool = False,
   ) -> str:
     """Assembles the final system instruction for the LLM."""
     parts = [role_description]
@@ -314,7 +317,7 @@ def generate_system_prompt(
       parts.append(final_catalog.render_as_llm_instructions())
 
     if include_examples:
-      examples_str = self.load_examples(final_catalog)
+      examples_str = self.load_examples(final_catalog, validate=validate_examples)
       if examples_str:
         parts.append(f"### Examples:\n{examples_str}")
 

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

@@ -0,0 +1,179 @@
+# 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 copy
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
+
+from jsonschema import Draft202012Validator
+
+if TYPE_CHECKING:
+  from .catalog import A2uiCatalog
+
+from .constants import (
+    BASE_SCHEMA_URL,
+    CATALOG_COMPONENTS_KEY,
+    CATALOG_ID_KEY,
+    CATALOG_STYLES_KEY,
+)
+
+
+def _inject_additional_properties(
+    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 and sets additionalProperties=False.
+
+  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 this node has additionalProperties=True, we inject the source properties
+        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
+
+
+# LLM is instructed to generate a list of messages, so we wrap the bundled schema in an array.
+def _wrap_main_schema(schema: Dict[str, Any]) -> Dict[str, Any]:
+  return {"type": "array", "items": schema}
+
+
+class A2uiValidator:
+  """Validator for A2UI messages."""
+
+  def __init__(self, catalog: "A2uiCatalog"):
+    self._catalog = catalog
+    self._validator = self._build_validator()
+
+  def _build_validator(self) -> Draft202012Validator:
+    """Builds a validator for the A2UI schema."""
+
+    if self._catalog.version == "0.8":
+      return self._build_0_8_validator()
+    return self._build_0_9_validator()
+
+  def _bundle_0_8_schemas(self) -> Dict[str, Any]:
+    if not self._catalog.s2c_schema:
+      return {}
+
+    bundled = copy.deepcopy(self._catalog.s2c_schema)
+
+    # Prepare catalog components and styles for injection
+    source_properties = {}
+    catalog_schema = self._catalog.catalog_schema
+    if catalog_schema:
+      if CATALOG_COMPONENTS_KEY in catalog_schema:
+        # Special mapping for v0.8: "components" -> "component"
+        source_properties["component"] = catalog_schema[CATALOG_COMPONENTS_KEY]
+      if CATALOG_STYLES_KEY in catalog_schema:
+        source_properties[CATALOG_STYLES_KEY] = catalog_schema[CATALOG_STYLES_KEY]
+
+    bundled, _ = _inject_additional_properties(bundled, source_properties)
+    return bundled
+
+  def _build_0_8_validator(self) -> Draft202012Validator:
+    """Builds a validator for the A2UI schema version 0.8."""
+    bundled_schema = self._bundle_0_8_schemas()
+    full_schema = _wrap_main_schema(bundled_schema)
+    return Draft202012Validator(full_schema)
+
+  def _build_0_9_validator(self) -> Draft202012Validator:
+    """Builds a validator for the A2UI schema version 0.9+."""
+    full_schema = _wrap_main_schema(self._catalog.s2c_schema)
+
+    from referencing import Registry, Resource
+
+    # v0.9 schemas (e.g. server_to_client.json) use relative references like
+    # 'catalog.json#/$defs/anyComponent'. Since server_to_client.json has
+    # $id: https://a2ui.org/specification/v0_9/server_to_client.json,
+    # these resolve to https://a2ui.org/specification/v0_9/catalog.json.
+    # We must register them using these absolute URIs.
+    base_uri = self._catalog.s2c_schema.get("$id", BASE_SCHEMA_URL)
+    import os
+
+    def get_sibling_uri(uri, filename):
+      return os.path.join(os.path.dirname(uri), filename)
+
+    catalog_uri = get_sibling_uri(base_uri, "catalog.json")
+    common_types_uri = get_sibling_uri(base_uri, "common_types.json")
+
+    resources = [
+        (
+            common_types_uri,
+            Resource.from_contents(self._catalog.common_types_schema),
+        ),
+        (
+            catalog_uri,
+            Resource.from_contents(self._catalog.catalog_schema),
+        ),
+        # Fallbacks for robustness
+        ("catalog.json", Resource.from_contents(self._catalog.catalog_schema)),
+        (
+            "common_types.json",
+            Resource.from_contents(self._catalog.common_types_schema),
+        ),
+    ]
+    # Also register the catalog ID if it's different from the catalog URI
+    if self._catalog.catalog_id and self._catalog.catalog_id != catalog_uri:
+      resources.append((
+          self._catalog.catalog_id,
+          Resource.from_contents(self._catalog.catalog_schema),
+      ))
+
+    registry = Registry().with_resources(resources)
+    validator_schema = copy.deepcopy(full_schema)
+    validator_schema["$schema"] = "https://json-schema.org/draft/2020-12/schema"
+
+    return Draft202012Validator(validator_schema, registry=registry)
+
+  def validate(self, message: Dict[str, Any]) -> None:
+    """Validates an A2UI message against the schema."""
+    error = next(self._validator.iter_errors(message), None)
+    if error is not None:
+      msg = f"Validation failed: {error.message}"
+      if error.context:
+        msg += "\nContext failures:"
+        for sub_error in error.context:
+          msg += f"\n  - {sub_error.message}"
+      raise ValueError(msg)

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

@@ -27,6 +27,7 @@ def generate_system_prompt(
       allowed_components: List[str] = [],
       include_schema: bool = False,
       include_examples: bool = False,
+      validate_examples: bool = False,
   ) -> str:
     # TODO: Implementation logic for Template Manager
     raise NotImplementedError("This method is not yet implemented.")

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

@@ -49,8 +49,12 @@ def test_catalog_id_missing_raises_error():
 def test_load_examples(tmp_path):
   example_dir = tmp_path / "examples"
   example_dir.mkdir()
-  (example_dir / "example1.json").write_text('{"foo": "bar"}')
-  (example_dir / "example2.json").write_text('{"baz": "qux"}')
+  (example_dir / "example1.json").write_text(
+      '[{"beginRendering": {"surfaceId": "id"}}]'
+  )
+  (example_dir / "example2.json").write_text(
+      '[{"beginRendering": {"surfaceId": "id"}}]'
+  )
   (example_dir / "ignored.txt").write_text("should not be loaded")
 
   catalog = A2uiCatalog(
@@ -63,9 +67,9 @@ def test_load_examples(tmp_path):
 
   examples_str = catalog.load_examples(str(example_dir))
   assert "---BEGIN example1---" in examples_str
-  assert '{"foo": "bar"}' in examples_str
+  assert '[{"beginRendering": {"surfaceId": "id"}}]' in examples_str
   assert "---BEGIN example2---" in examples_str
-  assert '{"baz": "qux"}' in examples_str
+  assert '[{"beginRendering": {"surfaceId": "id"}}]' in examples_str
   assert "ignored" not in examples_str
 
 

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

@@ -218,7 +218,7 @@ def joinpath_side_effect(path):
         return_value="---BEGIN example1---\n{}\n---END example1---",
     ):
       prompt = manager.generate_system_prompt("Role description", include_examples=True)
-      assert "### Examples:" in prompt
+      assert "### Examples" in prompt
       assert "example1" in prompt
 
   # Test without examples