From 215b64aa8338046bc9e90a7b843e4f6fb3743639 Mon Sep 17 00:00:00 2001 From: Sarthak Agarwal Date: Thu, 12 Mar 2026 10:45:56 -0400 Subject: [PATCH 1/3] Fix multi-surface validation and inline catalog handling - Validator: support multiple surfaces with different root IDs by tracking per-surface roots instead of a single global root - Validator: distinguish initial renders (beginRendering/createSurface) from incremental updates; skip root and orphan checks on updates while still catching cycles, self-refs, and duplicates - Manager: merge inline catalog components onto the base catalog instead of using inline-only; allow both inlineCatalogs and supportedCatalogIds in client capabilities - Agent: remove schema from system prompt (provided per-request via client capabilities); pass client_ui_capabilities through to catalog selection and validation - Executor: extract a2uiClientCapabilities from all DataParts (not just request parts) so UI events use the correct catalog - Client: inline OrgChart action schema (was unresolvable $ref) and add path-reference oneOf for chain (matches agent examples) --- .../python/src/a2ui/core/schema/manager.py | 44 ++-- .../tests/core/schema/test_validator.py | 194 ++++++++++++++++++ .../agent_executor.py | 22 +- .../prompt_builder.py | 2 +- 4 files changed, 232 insertions(+), 30 deletions(-) diff --git a/agent_sdks/python/src/a2ui/core/schema/manager.py b/agent_sdks/python/src/a2ui/core/schema/manager.py index 97e9eefd4..4d78f40e2 100644 --- a/agent_sdks/python/src/a2ui/core/schema/manager.py +++ b/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,25 @@ 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, diff --git a/agent_sdks/python/tests/core/schema/test_validator.py b/agent_sdks/python/tests/core/schema/test_validator.py index 47224b794..60a78ecd9 100644 --- a/agent_sdks/python/tests/core/schema/test_validator.py +++ b/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) diff --git a/samples/agent/adk/contact_multiple_surfaces/agent_executor.py b/samples/agent/adk/contact_multiple_surfaces/agent_executor.py index 72effe18b..82ce090fb 100644 --- a/samples/agent/adk/contact_multiple_surfaces/agent_executor.py +++ b/samples/agent/adk/contact_multiple_surfaces/agent_executor.py @@ -75,23 +75,23 @@ async def execute( ) for i, part in enumerate(context.message.parts): if isinstance(part.root, DataPart): + # Extract client UI capabilities from any DataPart that has them + if ( + agent.schema_manager.accepts_inline_catalogs + and "metadata" in part.root.data + and "a2uiClientCapabilities" in part.root.data["metadata"] + ): + logger.info(f" Part {i}: Found 'a2uiClientCapabilities' in DataPart.") + client_ui_capabilities = part.root.data["metadata"][ + "a2uiClientCapabilities" + ] + if "userAction" in part.root.data: logger.info(f" Part {i}: Found a2ui UI ClientEvent payload.") ui_event_part = part.root.data["userAction"] elif "request" in part.root.data: logger.info(f" Part {i}: Found 'request' in DataPart.") query = part.root.data["request"] - - # Check for inline catalog - if ( - agent.schema_manager.accepts_inline_catalogs - and "metadata" in part.root.data - and "a2uiClientCapabilities" in part.root.data["metadata"] - ): - logger.info(f" Part {i}: Found 'a2uiClientCapabilities' in DataPart.") - client_ui_capabilities = part.root.data["metadata"][ - "a2uiClientCapabilities" - ] else: logger.info(f" Part {i}: DataPart (data: {part.root.data})") elif isinstance(part.root, TextPart): diff --git a/samples/agent/adk/contact_multiple_surfaces/prompt_builder.py b/samples/agent/adk/contact_multiple_surfaces/prompt_builder.py index 8bc8cbeb0..b1492091e 100644 --- a/samples/agent/adk/contact_multiple_surfaces/prompt_builder.py +++ b/samples/agent/adk/contact_multiple_surfaces/prompt_builder.py @@ -88,7 +88,7 @@ def get_text_prompt() -> str: client_ui_capabilities_str = ( '{"inlineCatalogs":[{"catalogId": "inline_catalog",' - ' "components":{"OrgChart":{"type":"object","properties":{"chain":{"type":"array","items":{"type":"object","properties":{"title":{"type":"string"},"name":{"type":"string"}},"required":["title","name"]}},"action":{"$ref":"#/definitions/Action"}},"required":["chain"]},"WebFrame":{"type":"object","properties":{"url":{"type":"string"},"html":{"type":"string"},"height":{"type":"number"},"interactionMode":{"type":"string","enum":["readOnly","interactive"]},"allowedEvents":{"type":"array","items":{"type":"string"}}}}}}]}' + ' "components":{"OrgChart":{"type":"object","properties":{"chain":{"oneOf":[{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]},{"type":"array","items":{"type":"object","properties":{"title":{"type":"string"},"name":{"type":"string"}},"required":["title","name"]}}]},"action":{"type":"object","properties":{"name":{"type":"string"},"context":{"type":"array","items":{"type":"object","properties":{"key":{"type":"string"},"value":{"type":"object","properties":{"path":{"type":"string"},"literalString":{"type":"string"},"literalNumber":{"type":"number"},"literalBoolean":{"type":"boolean"}}}},"required":["key","value"]}}},"required":["name"]}},"required":["chain"]},"WebFrame":{"type":"object","properties":{"url":{"type":"string"},"html":{"type":"string"},"height":{"type":"number"},"interactionMode":{"type":"string","enum":["readOnly","interactive"]},"allowedEvents":{"type":"array","items":{"type":"string"}}}}}}]}' ) client_ui_capabilities = json.loads(client_ui_capabilities_str) inline_catalog = schema_manager.get_selected_catalog( From 87bae8a820ec35317170418cd32b887a1378db75 Mon Sep 17 00:00:00 2001 From: Sarthak Agarwal Date: Thu, 19 Mar 2026 12:58:05 -0400 Subject: [PATCH 2/3] Format manager.py with pyink --- agent_sdks/python/src/a2ui/core/schema/manager.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/agent_sdks/python/src/a2ui/core/schema/manager.py b/agent_sdks/python/src/a2ui/core/schema/manager.py index 4d78f40e2..185df6a2e 100644 --- a/agent_sdks/python/src/a2ui/core/schema/manager.py +++ b/agent_sdks/python/src/a2ui/core/schema/manager.py @@ -144,9 +144,7 @@ def _select_catalog( # 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 - } + 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] From 3ce76301592cdae0dd2f3343fc8c4815da6dd5a1 Mon Sep 17 00:00:00 2001 From: Sarthak Agarwal Date: Thu, 19 Mar 2026 13:08:07 -0400 Subject: [PATCH 3/3] Update test_schema_manager tests for new inline catalog merging behavior --- .../tests/core/schema/test_schema_manager.py | 68 ++++++++++++++----- 1 file changed, 50 insertions(+), 18 deletions(-) diff --git a/agent_sdks/python/tests/core/schema/test_schema_manager.py b/agent_sdks/python/tests/core/schema/test_schema_manager.py index 999b1081f..e4435635d 100644 --- a/agent_sdks/python/tests/core/schema/test_schema_manager.py +++ b/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 (