Fix inline catalog validation and multi-surface support

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

@@ -19,7 +19,7 @@
 import importlib.resources
 from typing import List, Dict, Any, Optional, Callable
 from dataclasses import dataclass, field
-from .utils import load_from_bundled_resource, deep_update
+from .utils import load_from_bundled_resource
 from ..inference_strategy import InferenceStrategy
 from .constants import *
 from .catalog import CatalogConfig, A2uiCatalog
@@ -103,8 +103,11 @@ def _select_catalog(
     """Selects the component catalog for the prompt based on client capabilities.
 
     Selection priority:
-    1. First inline catalog if provided (and accepted by the agent).
-    2. First client-supported catalog ID that is also supported by the agent.
+    1. If inline catalogs are provided (and accepted by the agent), their
+       components are merged on top of a base catalog. The base is determined
+       by supportedCatalogIds (if also provided) or the agent's default catalog.
+    2. If only supportedCatalogIds is provided, pick the first mutually
+       supported catalog.
     3. Fallback to the first agent-supported catalog (usually the bundled catalog).
 
     Args:
@@ -114,8 +117,8 @@ def _select_catalog(
     Returns:
       The resolved A2uiCatalog.
     Raises:
-      ValueError: If capabilities are ambiguous (both inline_catalogs and supported_catalog_ids are provided), if inline
-        catalogs are sent but not accepted, or if no mutually supported catalog is found.
+      ValueError: If inline catalogs are sent but not accepted, or if no
+        mutually supported catalog is found.
     """
     if not self._supported_catalogs:
       raise ValueError("No supported catalogs found.")  # This should not happen.
@@ -136,20 +139,23 @@ def _select_catalog(
           " capabilities. However, the agent does not accept inline catalogs."
       )
 
-    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 inline catalog schema.
-      inline_catalog_schema = inline_catalogs[0]
-      inline_catalog_schema = self._apply_modifiers(inline_catalog_schema)
-
-      # Deep merge the standard catalog properties with the inline catalog
-      merged_schema = copy.deepcopy(self._supported_catalogs[0].catalog_schema)
-      deep_update(merged_schema, inline_catalog_schema)
+      # Determine the base catalog: use supportedCatalogIds if provided,
+      # otherwise fall back to the agent's default catalog.
+      base_catalog = self._supported_catalogs[0]
+      if client_supported_catalog_ids:
+        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:
+            base_catalog = agent_supported_catalogs[cscid]
+            break
+
+      merged_schema = copy.deepcopy(base_catalog.catalog_schema)
+
+      for inline_catalog_schema in inline_catalogs:
+        inline_catalog_schema = self._apply_modifiers(inline_catalog_schema)
+        inline_components = inline_catalog_schema.get(CATALOG_COMPONENTS_KEY, {})
+        merged_schema[CATALOG_COMPONENTS_KEY].update(inline_components)
 
       return A2uiCatalog(
           version=self._version,

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

@@ -412,11 +412,9 @@ def joinpath_side_effect(path):
 
   assert "Role" in prompt
   assert "---BEGIN A2UI JSON SCHEMA---" in prompt
-  assert (
-      '### Catalog Schema:\n{\n  "$schema":'
-      ' "https://json-schema.org/draft/2020-12/schema",\n  "catalogId": "id_inline"'
-      in prompt
-  )
+  # Inline catalog is merged onto the base catalog (catalogId: "basic")
+  assert "### Catalog Schema:" in prompt
+  assert '"catalogId": "basic"' in prompt
   assert '"Button": {}' in prompt
 
 
@@ -469,28 +467,59 @@ def test_select_catalog_logic():
   assert A2uiSchemaManager._select_catalog(manager, {}) == basic
   assert A2uiSchemaManager._select_catalog(manager, None) == basic
 
-  # Rule 2: Exception if both inline and supported IDs are provided
-  with pytest.raises(ValueError, match="Only one is allowed"):
-    A2uiSchemaManager._select_catalog(
-        manager,
-        {
-            INLINE_CATALOGS_KEY: [{"inline": "catalog"}],
-            SUPPORTED_CATALOG_IDS_KEY: ["id_custom1"],
-        },
-    )
-
-  # Rule 3: Inline catalog loading
+  # Rule 2: Both inline and supported IDs are allowed (supportedCatalogIds
+  # selects the base catalog, inlineCatalogs extend it).
+  inline_with_supported = {
+      INLINE_CATALOGS_KEY: [{"components": {"Custom": {}}}],
+      SUPPORTED_CATALOG_IDS_KEY: ["id_custom1"],
+  }
+  # Need components on the base catalog for merging to work
+  custom1_with_components = A2uiCatalog(
+      version=VERSION_0_9,
+      name="custom1",
+      s2c_schema={},
+      common_types_schema={},
+      catalog_schema={
+          "$schema": "https://json-schema.org/draft/2020-12/schema",
+          "catalogId": "id_custom1",
+          "components": {"Base": {}},
+      },
+  )
+  manager._supported_catalogs[1] = custom1_with_components
+  manager._apply_modifiers = MagicMock(side_effect=lambda x: x)
+  catalog_both = A2uiSchemaManager._select_catalog(manager, inline_with_supported)
+  assert catalog_both.name == INLINE_CATALOG_NAME
+  # Base catalog's components are preserved and inline components are merged
+  assert "Base" in catalog_both.catalog_schema["components"]
+  assert "Custom" in catalog_both.catalog_schema["components"]
+  assert catalog_both.catalog_schema["catalogId"] == "id_custom1"
+
+  # Rule 3: Inline catalog loading (merges onto base)
+  basic_with_components = A2uiCatalog(
+      version=VERSION_0_9,
+      name=BASIC_CATALOG_NAME,
+      s2c_schema={},
+      common_types_schema={},
+      catalog_schema={
+          "$schema": "https://json-schema.org/draft/2020-12/schema",
+          "catalogId": "id_basic",
+          "components": {"Text": {}},
+      },
+  )
+  manager._supported_catalogs[0] = basic_with_components
   inline_schema = {
       "$schema": "https://json-schema.org/draft/2020-12/schema",
       "catalogId": "id_inline",
-      "components": {},
+      "components": {"Button": {}},
   }
   manager._apply_modifiers = MagicMock(return_value=inline_schema)
   catalog_inline = A2uiSchemaManager._select_catalog(
       manager, {INLINE_CATALOGS_KEY: [inline_schema]}
   )
   assert catalog_inline.name == INLINE_CATALOG_NAME
-  assert catalog_inline.catalog_schema == inline_schema
+  # Merged: base components + inline components
+  assert "Text" in catalog_inline.catalog_schema["components"]
+  assert "Button" in catalog_inline.catalog_schema["components"]
   assert catalog_inline.s2c_schema == manager._server_to_client_schema
   assert catalog_inline.common_types_schema == manager._common_types_schema
 
@@ -500,6 +529,9 @@ def test_select_catalog_logic():
     A2uiSchemaManager._select_catalog(manager, {INLINE_CATALOGS_KEY: [inline_schema]})
   manager._accepts_inline_catalogs = True
 
+  # Restore original catalogs for remaining tests
+  manager._supported_catalogs = [basic, custom1, custom2]
+
   # Rule 4: Otherwise, find the intersection, return any catalog that matches.
   # The priority is determined by the order in supported_catalog_ids.
   assert (

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

@@ -1037,3 +1037,197 @@ def test_validate_global_recursion_limit_exceeded(self, test_catalog):
     payload = self.make_payload(test_catalog, data_model=deep_data)
     with pytest.raises(ValueError, match="Global recursion limit exceeded"):
       test_catalog.validator.validate(payload)
+
+  # --- Multi-surface tests ---
+
+  def test_validate_multi_surface_v08(self, catalog_0_8):
+    """Tests that multiple surfaces with different root IDs validate correctly."""
+    payload = [
+        {"beginRendering": {"surfaceId": "surface-a", "root": "root-a"}},
+        {"beginRendering": {"surfaceId": "surface-b", "root": "root-b"}},
+        {
+            "surfaceUpdate": {
+                "surfaceId": "surface-a",
+                "components": [
+                    {"id": "root-a", "component": {"Card": {"child": "child-a"}}},
+                    {"id": "child-a", "component": {"Text": {"text": "Hello A"}}},
+                ],
+            }
+        },
+        {
+            "surfaceUpdate": {
+                "surfaceId": "surface-b",
+                "components": [
+                    {"id": "root-b", "component": {"Card": {"child": "child-b"}}},
+                    {"id": "child-b", "component": {"Text": {"text": "Hello B"}}},
+                ],
+            }
+        },
+    ]
+    # Should not raise - each surface has its own root
+    catalog_0_8.validator.validate(payload)
+
+  def test_validate_multi_surface_missing_root_v08(self, catalog_0_8):
+    """Tests that missing root in one surface still fails validation."""
+    payload = [
+        {"beginRendering": {"surfaceId": "surface-a", "root": "root-a"}},
+        {"beginRendering": {"surfaceId": "surface-b", "root": "root-b"}},
+        {
+            "surfaceUpdate": {
+                "surfaceId": "surface-a",
+                "components": [
+                    {"id": "root-a", "component": {"Text": {"text": "Hello A"}}},
+                ],
+            }
+        },
+        {
+            "surfaceUpdate": {
+                "surfaceId": "surface-b",
+                "components": [
+                    # Missing root-b, only has a non-root component
+                    {"id": "not-root-b", "component": {"Text": {"text": "Hello B"}}},
+                ],
+            }
+        },
+    ]
+    with pytest.raises(ValueError, match="Missing root component.*root-b"):
+      catalog_0_8.validator.validate(payload)
+
+  # --- Incremental update tests ---
+
+  def test_incremental_update_no_root_v08(self, catalog_0_8):
+    """Incremental update without root component should pass."""
+    payload = [
+        {
+            "surfaceUpdate": {
+                "surfaceId": "contact-card",
+                "components": [
+                    {"id": "main_card", "component": {"Card": {"child": "col"}}},
+                    {"id": "col", "component": {"Text": {"text": "Updated"}}},
+                ],
+            }
+        },
+    ]
+    # No beginRendering → incremental update → root check skipped
+    catalog_0_8.validator.validate(payload)
+
+  def test_incremental_update_no_root_v09(self, catalog_0_9):
+    """Incremental update without root component should pass (v0.9)."""
+    payload = [
+        {
+            "version": "v0.9",
+            "updateComponents": {
+                "surfaceId": "contact-card",
+                "components": [
+                    {"id": "card1", "component": "Card", "child": "text1"},
+                    {"id": "text1", "component": "Text", "text": "Updated"},
+                ],
+            },
+        },
+    ]
+    # No createSurface → incremental update → root check skipped
+    catalog_0_9.validator.validate(payload)
+
+  def test_incremental_update_orphans_allowed_v08(self, catalog_0_8):
+    """Incremental update with 'orphaned' components should pass."""
+    payload = [
+        {
+            "surfaceUpdate": {
+                "surfaceId": "contact-card",
+                "components": [
+                    {"id": "text1", "component": {"Text": {"text": "Hello"}}},
+                    {"id": "text2", "component": {"Text": {"text": "World"}}},
+                ],
+            }
+        },
+    ]
+    # These are disconnected but it's an incremental update
+    catalog_0_8.validator.validate(payload)
+
+  def test_incremental_update_self_ref_still_fails(self, test_catalog):
+    """Self-references should still be caught in incremental updates."""
+    if test_catalog.version == VERSION_0_8:
+      payload = [
+          {
+              "surfaceUpdate": {
+                  "surfaceId": "s1",
+                  "components": [
+                      {"id": "card1", "component": {"Card": {"child": "card1"}}},
+                  ],
+              }
+          },
+      ]
+    else:
+      payload = [
+          {
+              "version": "v0.9",
+              "updateComponents": {
+                  "surfaceId": "s1",
+                  "components": [
+                      {"id": "card1", "component": "Card", "child": "card1"},
+                  ],
+              },
+          },
+      ]
+    with pytest.raises(ValueError, match="Self-reference detected"):
+      test_catalog.validator.validate(payload)
+
+  def test_incremental_update_cycle_still_fails(self, test_catalog):
+    """Cycles should still be caught in incremental updates."""
+    if test_catalog.version == VERSION_0_8:
+      payload = [
+          {
+              "surfaceUpdate": {
+                  "surfaceId": "s1",
+                  "components": [
+                      {"id": "a", "component": {"Card": {"child": "b"}}},
+                      {"id": "b", "component": {"Card": {"child": "a"}}},
+                  ],
+              }
+          },
+      ]
+    else:
+      payload = [
+          {
+              "version": "v0.9",
+              "updateComponents": {
+                  "surfaceId": "s1",
+                  "components": [
+                      {"id": "a", "component": "Card", "child": "b"},
+                      {"id": "b", "component": "Card", "child": "a"},
+                  ],
+              },
+          },
+      ]
+    with pytest.raises(ValueError, match="Circular reference detected"):
+      test_catalog.validator.validate(payload)
+
+  def test_incremental_update_duplicates_still_fail(self, test_catalog):
+    """Duplicate IDs should still be caught in incremental updates."""
+    if test_catalog.version == VERSION_0_8:
+      payload = [
+          {
+              "surfaceUpdate": {
+                  "surfaceId": "s1",
+                  "components": [
+                      {"id": "text1", "component": {"Text": {"text": "A"}}},
+                      {"id": "text1", "component": {"Text": {"text": "B"}}},
+                  ],
+              }
+          },
+      ]
+    else:
+      payload = [
+          {
+              "version": "v0.9",
+              "updateComponents": {
+                  "surfaceId": "s1",
+                  "components": [
+                      {"id": "text1", "component": "Text", "text": "A"},
+                      {"id": "text1", "component": "Text", "text": "B"},
+                  ],
+              },
+          },
+      ]
+    with pytest.raises(ValueError, match="Duplicate component ID"):
+      test_catalog.validator.validate(payload)