diff --git a/agent_sdks/python/src/a2ui/core/schema/manager.py b/agent_sdks/python/src/a2ui/core/schema/manager.py index 5a92c8787..5a90a4d16 100644 --- a/agent_sdks/python/src/a2ui/core/schema/manager.py +++ b/agent_sdks/python/src/a2ui/core/schema/manager.py @@ -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,33 @@ 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) + # 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) + + if CATALOG_COMPONENTS_KEY not in merged_schema: + merged_schema[CATALOG_COMPONENTS_KEY] = {} + + 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, name=INLINE_CATALOG_NAME, - catalog_schema=inline_catalog_schema, + catalog_schema=merged_schema, s2c_schema=self._server_to_client_schema, common_types_schema=self._common_types_schema, ) diff --git a/agent_sdks/python/src/a2ui/core/schema/validator.py b/agent_sdks/python/src/a2ui/core/schema/validator.py index 07c1b968a..11ca3fc6a 100644 --- a/agent_sdks/python/src/a2ui/core/schema/validator.py +++ b/agent_sdks/python/src/a2ui/core/schema/validator.py @@ -274,51 +274,78 @@ def validate(self, a2ui_json: Union[Dict[str, Any], List[Any]]) -> None: msg += f"\n - {sub_error.message}" raise ValueError(msg) - root_id = _find_root_id(messages) + root_ids, initial_surfaces = _find_root_ids(messages) + ref_map = _extract_component_ref_fields(self._catalog) for message in messages: if not isinstance(message, dict): continue components = None + surface_id = None if "surfaceUpdate" in message: # v0.8 components = message["surfaceUpdate"].get(COMPONENTS) + surface_id = message["surfaceUpdate"].get("surfaceId") elif "updateComponents" in message and isinstance( message["updateComponents"], dict ): # v0.9 components = message["updateComponents"].get(COMPONENTS) + surface_id = message["updateComponents"].get("surfaceId") if components: - ref_map = _extract_component_ref_fields(self._catalog) - _validate_component_integrity(root_id, components, ref_map) - _validate_topology(root_id, components, ref_map) + root_id = root_ids.get(surface_id, root_ids.get(None, ROOT)) + is_initial = surface_id in initial_surfaces + _validate_component_integrity(root_id, components, ref_map, is_initial) + _validate_topology(root_id, components, ref_map, is_initial) _validate_recursion_and_paths(message) -def _find_root_id(messages: List[Dict[str, Any]]) -> str: +def _find_root_ids( + messages: List[Dict[str, Any]], +) -> Tuple[Dict[Optional[str], str], Set[str]]: """ - Finds the root id from a list of A2UI messages. - - For v0.8, the root id is in the beginRendering message. - - For v0.9+, the root id is 'root'. + Finds root ids and initial surfaces from a list of A2UI messages. + - For v0.8, root ids come from beginRendering messages (surfaceId → root). + - For v0.9+, the root id is always 'root'; initial surfaces come from createSurface. + + Returns: + A tuple of: + - root_ids: dict mapping surfaceId → root_id (None key = default fallback). + - initial_surfaces: set of surfaceIds that have an initial render in this batch. """ + root_ids: Dict[Optional[str], str] = {None: ROOT} + initial_surfaces: Set[str] = set() for message in messages: if not isinstance(message, dict): continue - if "beginRendering" in message: - return message["beginRendering"].get(ROOT, ROOT) - return ROOT + if "beginRendering" in message: # v0.8 + begin = message["beginRendering"] + surface_id = begin.get("surfaceId") + root_id = begin.get(ROOT, ROOT) + root_ids[surface_id] = root_id + if surface_id: + initial_surfaces.add(surface_id) + # Set the first beginRendering as the default fallback + if root_ids[None] == ROOT: + root_ids[None] = root_id + elif "createSurface" in message: # v0.9 + surface_id = message["createSurface"].get("surfaceId") + if surface_id: + initial_surfaces.add(surface_id) + return root_ids, initial_surfaces def _validate_component_integrity( root_id: str, components: List[Dict[str, Any]], ref_fields_map: Dict[str, tuple[Set[str], Set[str]]], + is_initial_render: bool = True, ) -> None: """ Validates that: 1. All component IDs are unique. - 2. A 'root' component exists. + 2. A 'root' component exists (initial renders only). 3. All references point to existing IDs. """ ids: Set[str] = set() @@ -333,8 +360,8 @@ def _validate_component_integrity( raise ValueError(f"Duplicate component ID: {comp_id}") ids.add(comp_id) - # 2. Check for root component - if root_id not in ids: + # 2. Check for root component (only on initial renders) + if is_initial_render and root_id not in ids: raise ValueError(f"Missing root component: No component has id='{root_id}'") # 3. Check for dangling references using helper @@ -351,11 +378,12 @@ def _validate_topology( root_id: str, components: List[Dict[str, Any]], ref_fields_map: Dict[str, tuple[Set[str], Set[str]]], + is_initial_render: bool = True, ) -> None: """ Validates the topology of the component tree: 1. No circular references (including self-references). - 2. No orphaned components (all components must be reachable from 'root'). + 2. No orphaned components (initial renders only). """ adj_list: Dict[str, List[str]] = {} all_ids: Set[str] = set() @@ -401,16 +429,22 @@ def dfs(node_id: str, depth: int): recursion_stack.remove(node_id) - if root_id in all_ids: - dfs(root_id, 0) - - # Check for Orphans - orphans = all_ids - visited - if orphans: - sorted_orphans = sorted(list(orphans)) - raise ValueError( - f"Component '{sorted_orphans[0]}' is not reachable from '{root_id}'" - ) + if is_initial_render: + # Initial render: DFS from root, then check for orphans + if root_id in all_ids: + dfs(root_id, 0) + orphans = all_ids - visited + if orphans: + sorted_orphans = sorted(list(orphans)) + raise ValueError( + f"Component '{sorted_orphans[0]}' is not reachable from '{root_id}'" + ) + else: + # Incremental update: DFS from all nodes to detect cycles, + # but skip orphan check (partial trees are expected) + for node_id in all_ids: + if node_id not in visited: + dfs(node_id, 0) def _extract_component_ref_fields( diff --git a/agent_sdks/python/tests/core/schema/test_validator.py b/agent_sdks/python/tests/core/schema/test_validator.py index d46d093fa..dfc902470 100644 --- a/agent_sdks/python/tests/core/schema/test_validator.py +++ b/agent_sdks/python/tests/core/schema/test_validator.py @@ -760,22 +760,31 @@ def test_validate_duplicate_ids(self, test_catalog): def test_validate_missing_root(self, test_catalog): # This payload has components but none are 'root' - # bypass make_payload as it adds root if missing + # Include beginRendering/createSurface to mark as initial render if test_catalog.version == VERSION_0_8: - payload = { - "surfaceUpdate": { - "surfaceId": "test", - "components": [{"id": "c1", "component": {"Text": {"text": "hi"}}}], - } - } + payload = [ + {"beginRendering": {"surfaceId": "test"}}, + { + "surfaceUpdate": { + "surfaceId": "test", + "components": [{"id": "c1", "component": {"Text": {"text": "hi"}}}], + } + }, + ] else: - payload = [{ - "version": "v0.9", - "updateComponents": { - "surfaceId": "test", - "components": [{"id": "c1", "component": "Text", "text": "hi"}], + payload = [ + { + "version": "v0.9", + "createSurface": {"surfaceId": "test", "catalogId": "standard"}, }, - }] + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "test", + "components": [{"id": "c1", "component": "Text", "text": "hi"}], + }, + }, + ] with pytest.raises(ValueError, match="Missing root component"): test_catalog.validator.validate(payload) @@ -838,6 +847,17 @@ def test_validate_orphaned_component(self, test_catalog): {"id": "orphan", "component": "Text", "text": "Orphan"}, ] payload = self.make_payload(test_catalog, components=components) + # Add initial render marker so orphan check runs + if test_catalog.version == VERSION_0_8: + payload = [ + {"beginRendering": {"surfaceId": "test-surface"}}, + payload, + ] + else: + payload.insert(0, { + "version": "v0.9", + "createSurface": {"surfaceId": "test-surface", "catalogId": "standard"}, + }) with pytest.raises( ValueError, match="Component 'orphan' is not reachable from 'root'" ): @@ -850,6 +870,17 @@ def test_validate_recursion_limit_exceeded(self, test_catalog): components.append({"id": f"c{55}", "component": "Text", "text": "End"}) payload = self.make_payload(test_catalog, components=components) + # Add initial render marker so DFS starts from root + if test_catalog.version == VERSION_0_8: + payload = [ + {"beginRendering": {"surfaceId": "test-surface"}}, + payload, + ] + else: + payload.insert(0, { + "version": "v0.9", + "createSurface": {"surfaceId": "test-surface", "catalogId": "standard"}, + }) with pytest.raises( ValueError, match="Global recursion limit exceeded: logical depth" ): @@ -862,6 +893,17 @@ def test_validate_recursion_limit_valid(self, test_catalog): components.append({"id": f"c{40}", "component": "Text", "text": "End"}) payload = self.make_payload(test_catalog, components=components) + # Add initial render marker so DFS starts from root + if test_catalog.version == VERSION_0_8: + payload = [ + {"beginRendering": {"surfaceId": "test-surface"}}, + payload, + ] + else: + payload.insert(0, { + "version": "v0.9", + "createSurface": {"surfaceId": "test-surface", "catalogId": "standard"}, + }) test_catalog.validator.validate(payload) def test_validate_template_reachability(self, test_catalog): @@ -1007,3 +1049,136 @@ 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) + + 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.""" + components = [ + {"id": "card1", "component": "Card", "child": "card1"}, + ] + payload = self.make_payload(test_catalog, components=components) + 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.""" + components = [ + {"id": "a", "component": "Card", "child": "b"}, + {"id": "b", "component": "Card", "child": "a"}, + ] + payload = self.make_payload(test_catalog, components=components) + 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.""" + components = [ + {"id": "text1", "component": "Text", "text": "A"}, + {"id": "text1", "component": "Text", "text": "B"}, + ] + payload = self.make_payload(test_catalog, components=components) + with pytest.raises(ValueError, match="Duplicate component ID"): + test_catalog.validator.validate(payload) diff --git a/samples/agent/adk/contact_multiple_surfaces/agent.py b/samples/agent/adk/contact_multiple_surfaces/agent.py index 1b1d45811..9f15656f1 100644 --- a/samples/agent/adk/contact_multiple_surfaces/agent.py +++ b/samples/agent/adk/contact_multiple_surfaces/agent.py @@ -130,7 +130,7 @@ def _build_agent(self, use_ui: bool) -> LlmAgent: workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, include_examples=True, - include_schema=True, + include_schema=False, validate_examples=False, # Missing inline_catalogs for OrgChart and WebFrame validation ) if use_ui @@ -145,7 +145,7 @@ def _build_agent(self, use_ui: bool) -> LlmAgent: tools=[get_contact_info], ) - async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: + async def stream(self, query, session_id, client_ui_capabilities=None) -> AsyncIterable[dict[str, Any]]: session_state = {"base_url": self.base_url} session = await self._runner.session_service.get_session( @@ -169,7 +169,9 @@ async def stream(self, query, session_id) -> AsyncIterable[dict[str, Any]]: current_query_text = query # Ensure schema was loaded - selected_catalog = self.schema_manager.get_selected_catalog() + selected_catalog = self.schema_manager.get_selected_catalog( + client_ui_capabilities=client_ui_capabilities + ) if self.use_ui and not selected_catalog.catalog_schema: logger.error( "--- ContactAgent.stream: A2UI_SCHEMA is not loaded. " diff --git a/samples/agent/adk/contact_multiple_surfaces/agent_executor.py b/samples/agent/adk/contact_multiple_surfaces/agent_executor.py index c56d488c4..230172cf8 100644 --- a/samples/agent/adk/contact_multiple_surfaces/agent_executor.py +++ b/samples/agent/adk/contact_multiple_surfaces/agent_executor.py @@ -53,6 +53,7 @@ async def execute( query = "" ui_event_part = None action = None + client_ui_capabilities = None logger.info(f"--- Client requested extensions: {context.requested_extensions} ---") use_ui = try_activate_a2ui_extension(context) @@ -74,6 +75,17 @@ 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"] @@ -81,16 +93,7 @@ async def execute( 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" - ] + if client_ui_capabilities: catalog = agent.schema_manager.get_selected_catalog( client_ui_capabilities=client_ui_capabilities ) @@ -164,7 +167,7 @@ async def execute( await event_queue.enqueue_event(task) updater = TaskUpdater(event_queue, task.id, task.context_id) - async for item in agent.stream(query, task.context_id): + async for item in agent.stream(query, task.context_id, client_ui_capabilities=client_ui_capabilities): is_task_complete = item["is_task_complete"] if not is_task_complete: await updater.update_status( 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( diff --git a/samples/client/lit/contact/ui/custom-components/register-components.ts b/samples/client/lit/contact/ui/custom-components/register-components.ts index 047a0438c..b2470d812 100644 --- a/samples/client/lit/contact/ui/custom-components/register-components.ts +++ b/samples/client/lit/contact/ui/custom-components/register-components.ts @@ -25,17 +25,53 @@ export function registerContactComponents() { type: "object", properties: { chain: { - type: "array", - items: { - type: "object", - properties: { - title: { type: "string" }, - name: { type: "string" }, + oneOf: [ + { + type: "object", + description: "A path reference to an array in the data model.", + properties: { path: { type: "string" } }, + required: ["path"], + }, + { + type: "array", + items: { + type: "object", + properties: { + title: { type: "string" }, + name: { type: "string" }, + }, + required: ["title", "name"], + }, + }, + ], + }, + action: { + type: "object", + description: "The client-side action to be dispatched when a node is clicked.", + 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: ["title", "name"], }, + required: ["name"], }, - action: { $ref: "#/definitions/Action" }, }, required: ["chain"], });