aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/a2ui/middleware.py

@@ -0,0 +1,754 @@
+"""A2UI middleware for protocol-driven prompt enrichment and validation.
+
+Authors:
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, cast
+
+import jsonschema
+import requests
+from a2ui.basic_catalog.provider import BasicCatalog
+from a2ui.schema.catalog import CatalogConfig
+from a2ui.schema.common_modifiers import remove_strict_validation
+from a2ui.schema.constants import A2UI_CLOSE_TAG, A2UI_OPEN_TAG, VERSION_0_8, VERSION_0_9
+from a2ui.schema.manager import A2uiSchemaManager
+
+from aip_agents.a2ui.catalog_provider import (
+    UrlCatalogProvider,
+    materialize_example_urls,
+    release_materialized_example_dirs,
+)
+from aip_agents.a2ui.models import (
+    STANDARD_SCHEMA_VERSION,
+    A2UIActionContext,
+    A2UIActionStatus,
+    A2UIActionValidationResult,
+    A2UIContentMetadata,
+    A2UIContentMetadataStatus,
+    A2UIContext,
+    _normalize_iso_timestamp,
+)
+from aip_agents.middleware.base import AgentMiddleware, ModelRequest
+from aip_agents.types import JsonValue
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+_REQUIRED_ACTION_FIELDS = ("name", "surfaceId", "sourceComponentId", "timestamp", "context")
+_A2UI_REPAIR_ATTEMPT_LIMIT = 1
+
+# Canonical per-version message action keys in protocol stream order. v0.8 and v0.9 use different
+# message structures, so keep the version scope explicit at each constant name.
+_A2UI_V0_8_MESSAGE_ACTION_KEYS = ("beginRendering", "surfaceUpdate", "dataModelUpdate", "deleteSurface")
+_A2UI_V0_9_MESSAGE_ACTION_KEYS = ("createSurface", "updateComponents", "updateDataModel", "deleteSurface")
+
+# v0.9 messages carry `version` as a required sibling field next to the single action key (e.g.
+# ``{"version": "v0.9", "deleteSurface": {...}}``); v0.8 messages have no such top-level field.
+_A2UI_V0_9_SIBLING_FIELDS = ("version",)
+
+_A2UI_MESSAGE_ACTION_KEYS_BY_VERSION: dict[str, tuple[str, ...]] = {
+    VERSION_0_8: _A2UI_V0_8_MESSAGE_ACTION_KEYS,
+    VERSION_0_9: _A2UI_V0_9_MESSAGE_ACTION_KEYS,
+}
+_A2UI_SIBLING_FIELDS_BY_VERSION: dict[str, tuple[str, ...]] = {
+    VERSION_0_9: _A2UI_V0_9_SIBLING_FIELDS,
+}
+
+_A2UI_BUNDLED_EXAMPLE_VERSIONS = frozenset({VERSION_0_8, VERSION_0_9})
+
+
+def _default_examples_path(schema_version: str) -> str | None:
+    """Return the repo-owned default examples directory for the given schema version."""
+    if schema_version not in _A2UI_BUNDLED_EXAMPLE_VERSIONS:
+        return None
+    examples_dir = Path(__file__).resolve().parent / "examples" / schema_version
+    if not examples_dir.is_dir():
+        return None
+    return str(examples_dir)
+
+
+def _build_catalog_configs(
+    context: A2UIContext, schema_version: str
+) -> tuple[list[CatalogConfig], list[str], requests.Session | None]:
+    """Resolve bundled or URL-backed catalogs for the schema manager."""
+    materialized_example_dirs: list[str] = []
+    if not context.catalog_configs:
+        return (
+            [
+                BasicCatalog.get_config(
+                    version=schema_version,
+                    examples_path=_default_examples_path(schema_version),
+                )
+            ],
+            materialized_example_dirs,
+            None,
+        )
+
+    session = requests.Session()
+    catalog_configs: list[CatalogConfig] = []
+    try:
+        for index, entry in enumerate(context.catalog_configs):
+            examples_path = None
+            if entry.example_urls:
+                examples_path = materialize_example_urls(entry.example_urls, session=session)
+                materialized_example_dirs.append(examples_path)
+            catalog_configs.append(
+                CatalogConfig(
+                    name=f"custom-{index}",
+                    provider=UrlCatalogProvider(entry.catalog_url, session=session),
+                    examples_path=examples_path,
+                )
+            )
+        return catalog_configs, materialized_example_dirs, session
+    except Exception:
+        release_materialized_example_dirs(materialized_example_dirs)
+        session.close()
+        raise
+
+
+class A2UIMiddleware(AgentMiddleware):
+    """Middleware for A2UI protocol responsibilities.
+
+    The current implementation stores minimal state required by callers and
+    provides protocol-defined hooks with default no-op behavior.
+    """
+
+    tools: list[Any] = []
+    system_prompt_additions: str | None = None
+
+    def __init__(self, context: A2UIContext | None = None) -> None:
+        """Create a middleware instance.
+
+        Args:
+            context: Optional A2UI context resolved at agent creation.
+        """
+        self.context = context or A2UIContext()
+        self.schema_version = self.context.schema_version or STANDARD_SCHEMA_VERSION
+
+        self._materialized_example_dirs: list[str] = []
+        self._catalog_session: requests.Session | None = None
+        catalog_configs, self._materialized_example_dirs, self._catalog_session = _build_catalog_configs(
+            self.context, self.schema_version
+        )
+
+        try:
+            self._include_examples = any(catalog_config.examples_path is not None for catalog_config in catalog_configs)
+            self.schema_manager = A2uiSchemaManager(
+                version=self.schema_version,
+                catalogs=catalog_configs,
+                # remove_strict_validation has no return type annotation; cast pins the callable
+                # signature so schema_modifiers receives a correctly-typed function.
+                schema_modifiers=[cast(Callable[[dict[str, Any]], dict[str, Any]], remove_strict_validation)],
+            )
+        except Exception:
+            self.close()
+            raise
+
+    def __del__(self) -> None:
+        """Cleanup resources on middleware instance destruction."""
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    def close(self) -> None:
+        """Release downloaded example fixture directories and catalog HTTP session."""
+        if self._materialized_example_dirs:
+            release_materialized_example_dirs(self._materialized_example_dirs)
+            self._materialized_example_dirs = []
+        if self._catalog_session is not None:
+            self._catalog_session.close()
+            self._catalog_session = None
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """No-op hook prior to model call in first implementation pass."""
+        return {}
+
+    async def abefore_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronous mirror of before_model."""
+        return self.before_model(state)
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Append A2UI prompt instructions when enabled."""
+        if not self.context.enabled:
+            return request
+
+        prompt = request.get("system_prompt") or ""
+        prompt = prompt.rstrip("\n")
+        include_examples = self._include_examples
+        prompt = self.schema_manager.generate_system_prompt(
+            role_description=prompt,
+            include_schema=True,
+            include_examples=include_examples,
+            validate_examples=include_examples,
+        )
+        request["system_prompt"] = prompt
+        return request
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """No-op post-model hook."""
+        return {}
+
+    async def aafter_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronous mirror of after_model."""
+        return self.after_model(state)
+
+    def on_final_response(self, content: str, context: dict[str, Any]) -> None:
+        """Capture final response and attach validated A2UI metadata."""
+        if not self.context.enabled:
+            return
+
+        a2ui_content, sanitized_content = self._extract_a2ui_content(content)
+        self._apply_a2ui_content(context, a2ui_content, sanitized_content)
+
+    async def aon_final_response(self, content: str, context: dict[str, Any]) -> None:
+        """Capture final response and repair invalid A2UI content once when configured."""
+        if not self.context.enabled:
+            return
+
+        a2ui_content, sanitized_content = self._extract_a2ui_content(content)
+        if self._should_attempt_repair(content, context, a2ui_content):
+            repaired = await self._repair_a2ui_content(content, context, a2ui_content)
+            if repaired is not None:
+                repaired_a2ui_content, repaired_sanitized_content = self._extract_a2ui_content(repaired)
+                if repaired_a2ui_content.is_valid:
+                    a2ui_content = repaired_a2ui_content
+                    sanitized_content = repaired_sanitized_content
+
+        self._apply_a2ui_content(context, a2ui_content, sanitized_content)
+
+    def _apply_a2ui_content(
+        self, context: dict[str, Any], a2ui_content: A2UIContentMetadata, sanitized_content: str
+    ) -> None:
+        """Store A2UI metadata and sanitized visible content in hook context."""
+        context["a2ui_content"] = a2ui_content.model_dump(exclude_none=True)
+        final_state = context.get("final_state")
+        if isinstance(final_state, dict):
+            final_state["a2ui_content"] = a2ui_content.model_dump(exclude_none=True)
+
+        context["final_content"] = sanitized_content
+        event = context.get("event")
+        if isinstance(event, dict) and "content" in event:
+            event["content"] = sanitized_content
+
+    def _should_attempt_repair(self, content: str, context: dict[str, Any], a2ui_content: A2UIContentMetadata) -> bool:
+        """Return True when invalid marked A2UI content is eligible for one repair attempt."""
+        has_validation_errors = bool(
+            isinstance(a2ui_content.validation, dict) and a2ui_content.validation.get("errors")
+        )
+        if a2ui_content.is_valid and not has_validation_errors:
+            return False
+        if not isinstance(content, str) or not content:
+            return False
+        if A2UI_OPEN_TAG not in content and A2UI_CLOSE_TAG not in content:
+            return False
+        return self._get_repair_handler(context) is not None
+
+    async def _repair_a2ui_content(
+        self, content: str, context: dict[str, Any], a2ui_content: A2UIContentMetadata
+    ) -> str | None:
+        """Invoke the configured A2UI repair handler once and return repaired content."""
+        repair_handler = self._get_repair_handler(context)
+        if repair_handler is None:
+            return None
+
+        try:
+            logger.info("Attempting A2UI final-response repair attempt 1/%d", _A2UI_REPAIR_ATTEMPT_LIMIT)
+            result = repair_handler(
+                content=content,
+                context=context,
+                validation=a2ui_content.validation,
+                schema_version=self.schema_version,
+            )
+            if inspect.isawaitable(result):
+                result = await result
+        except Exception as exc:
+            logger.warning("A2UI final-response repair failed: %s", exc, exc_info=True)
+            return None
+
+        if isinstance(result, str) and result:
+            return result
+        logger.info("A2UI final-response repair returned no content")
+        return None
+
+    def _get_repair_handler(self, context: dict[str, Any]) -> Any:
+        """Resolve a repair handler from hook context or A2UI context extras."""
+        repair_handler = context.get("a2ui_repair_handler")
+        if repair_handler is None:
+            repair_handler = getattr(self.context, "repair_handler", None)
+        if callable(repair_handler):
+            return repair_handler
+        return None
+
+    def _extract_a2ui_content(self, content: str) -> tuple[A2UIContentMetadata, str]:
+        """Extract validated A2UI metadata and sanitized final text."""
+        if not content or not isinstance(content, str):
+            return (
+                A2UIContentMetadata(
+                    schema_version=self.schema_version,
+                    messages=[],
+                    validation=self._build_validation_block(
+                        status=A2UIContentMetadataStatus.FALLBACK,
+                        fallback_reason="missing_content",
+                        errors=["missing_content"],
+                    ),
+                ),
+                "",
+            )
+
+        if not self._contains_a2ui_payload(content):
+            return (
+                A2UIContentMetadata(
+                    schema_version=self.schema_version,
+                    messages=[],
+                    validation=self._build_validation_block(
+                        status=A2UIContentMetadataStatus.FALLBACK,
+                        fallback_reason="missing_content",
+                        errors=["missing_content"],
+                    ),
+                ),
+                content,
+            )
+
+        messages: list[dict[str, Any]] = []
+        selected_catalog = self.schema_manager.get_selected_catalog()
+        errors: list[str] = []
+        repairs: list[dict[str, Any]] = []
+
+        try:
+            response_text_parts, payload_chunks = self._split_a2ui_response(content)
+        except ValueError as e:
+            return (
+                A2UIContentMetadata(
+                    schema_version=self.schema_version,
+                    messages=[],
+                    validation=self._build_validation_block(
+                        status=A2UIContentMetadataStatus.FALLBACK,
+                        fallback_reason="invalid_content",
+                        errors=[str(e)],
+                    ),
+                ),
+                content,
+            )
+
+        for payload_chunk in payload_chunks:
+            try:
+                raw_messages = self._parse_payload_chunk(payload_chunk)
+                normalized_messages, chunk_repairs = self._split_merged_messages(raw_messages)
+                selected_catalog.validator.validate(normalized_messages)
+                messages.extend(cast(list[dict[str, Any]], normalized_messages))
+                repairs.extend(chunk_repairs)
+            except (ValueError, jsonschema.exceptions.ValidationError) as e:
+                errors.append(str(e))
+                logger.warning("Skipping invalid A2UI block: %s", e)
+
+        if not messages:
+            return (
+                A2UIContentMetadata(
+                    schema_version=self.schema_version,
+                    messages=[],
+                    validation=self._build_validation_block(
+                        status=A2UIContentMetadataStatus.FALLBACK,
+                        fallback_reason="invalid_content",
+                        errors=errors or ["invalid_content"],
+                    ),
+                ),
+                "".join(response_text_parts),
+            )
+
+        semantic_errors = self._validate_semantic_renderability(messages)
+        if semantic_errors:
+            return (
+                A2UIContentMetadata(
+                    schema_version=self.schema_version,
+                    messages=messages,
+                    validation=self._build_validation_block(
+                        status=A2UIContentMetadataStatus.FALLBACK,
+                        fallback_reason="invalid_content",
+                        errors=errors + semantic_errors,
+                    ),
+                ),
+                "".join(response_text_parts),
+            )
+
+        return (
+            A2UIContentMetadata(
+                schema_version=self.schema_version,
+                messages=messages,
+                validation=self._build_validation_block(
+                    status=A2UIContentMetadataStatus.VALID,
+                    errors=errors or None,
+                    repairs=repairs or None,
+                ),
+            ),
+            "".join(response_text_parts),
+        )
+
+    def _extract_surface_message(self, message: dict[str, Any]) -> tuple[str | None, dict[str, Any] | None]:
+        """Return the message type and dict payload for semantic renderability checks."""
+        if not isinstance(message, dict) or not message:
+            return None, None
+
+        message_type, payload = next(iter(message.items()))
+        if not isinstance(payload, dict):
+            return None, None
+
+        return message_type, payload
+
+    def _collect_component_ids(self, surface_update: dict[str, Any]) -> set[str]:
+        """Collect component ids from one surface update payload."""
+        component_ids: set[str] = set()
+        components = surface_update.get("components")
+        if not isinstance(components, list):
+            return component_ids
+
+        for component in components:
+            if not isinstance(component, dict):
+                continue
+            component_id = component.get("id")
+            if isinstance(component_id, str) and component_id:
+                component_ids.add(component_id)
+
+        return component_ids
+
+    def _collect_surface_renderability_state(
+        self, messages: list[dict[str, Any]]
+    ) -> tuple[dict[str, list[dict[str, Any]]], dict[str, list[dict[str, Any]]], dict[str, set[str]]]:
+        """Collect beginRendering, surfaceUpdate, and component-id state by surface."""
+        begin_renderings: dict[str, list[dict[str, Any]]] = {}
+        surface_updates: dict[str, list[dict[str, Any]]] = {}
+        component_ids_by_surface: dict[str, set[str]] = {}
+
+        for message in messages:
+            message_type, payload = self._extract_surface_message(message)
+            if payload is None:
+                continue
+
+            surface_id = payload.get("surfaceId")
+            if not isinstance(surface_id, str) or not surface_id:
+                continue
+
+            if message_type == "beginRendering":
+                # Keep the full list because final-response aggregation can still surface
+                # repeated beginRendering payloads for one surface, even though this pass
+                # only validates the first root until duplicate semantics are defined.
+                begin_renderings.setdefault(surface_id, []).append(payload)
+                continue
+
+            if message_type == "surfaceUpdate":
+                surface_updates.setdefault(surface_id, []).append(payload)
+                component_ids_by_surface.setdefault(surface_id, set()).update(self._collect_component_ids(payload))
+
+        return begin_renderings, surface_updates, component_ids_by_surface
+
+    def _iter_layout_child_reference_ids(self, component_payload: dict[str, Any]) -> list[str]:
+        """Collect v0.8 layout child ids from explicitList and template.componentId only."""
+        children = component_payload.get("children")
+        if not isinstance(children, dict):
+            return []
+
+        child_ids = [
+            child_id for child_id in children.get("explicitList", []) if isinstance(child_id, str) and child_id
+        ]
+
+        template = children.get("template")
+        if isinstance(template, dict):
+            template_component_id = template.get("componentId")
+            if isinstance(template_component_id, str) and template_component_id:
+                child_ids.append(template_component_id)
+
+        return child_ids
+
+    def _iter_child_reference_ids(self, node: Any) -> list[str]:
+        """Collect v0.8 child component ids referenced by a rendered component payload."""
+        if not isinstance(node, dict) or not node:
+            return []
+
+        component_type, payload = next(iter(node.items()))
+        if not isinstance(payload, dict):
+            return []
+
+        # This helper intentionally tracks only the v0.8 standard-catalog component
+        # shapes validated by BasicCatalog. Unknown/third-party components are rejected
+        # by the selected schema before this semantic pass runs today.
+
+        if component_type in {"Row", "Column", "List"}:
+            return self._iter_layout_child_reference_ids(payload)
+
+        if component_type in {"Card", "Button"}:
+            child_id = payload.get("child")
+            return [child_id] if isinstance(child_id, str) and child_id else []
+
+        if component_type == "Tabs":
+            tab_items = payload.get("tabItems")
+            if not isinstance(tab_items, list):
+                return []
+
+            return [
+                child_id
+                for tab_item in tab_items
+                if isinstance(tab_item, dict)
+                for child_id in [tab_item.get("child")]
+                if isinstance(child_id, str) and child_id
+            ]
+
+        if component_type == "Modal":
+            return [
+                child_id
+                for child_id in [payload.get("entryPointChild"), payload.get("contentChild")]
+                if isinstance(child_id, str) and child_id
+            ]
+
+        return []
+
+    def _validate_explicit_child_references(
+        self, surface_id: str, surface_update: dict[str, Any], component_ids: set[str]
+    ) -> list[str]:
+        """Check child references against component ids in one update."""
+        components = surface_update.get("components")
+        if not isinstance(components, list):
+            return []
+
+        errors: list[str] = []
+        for component in components:
+            if not isinstance(component, dict):
+                continue
+
+            component_name = component.get("id")
+            rendered_component = component.get("component")
+            if not isinstance(rendered_component, dict):
+                continue
+
+            missing_child_ids = sorted(
+                {
+                    child_id
+                    for child_id in self._iter_child_reference_ids(rendered_component)
+                    if child_id not in component_ids
+                }
+            )
+            if missing_child_ids:
+                component_label = component_name if isinstance(component_name, str) and component_name else "<unknown>"
+                errors.append(
+                    "surfaceUpdate.surfaceId "
+                    f"'{surface_id}' component '{component_label}' references missing child ids: "
+                    f"{', '.join(missing_child_ids)}"
+                )
+
+        return errors
+
+    def _validate_surface_renderability(
+        self,
+        *,
+        surface_id: str,
+        updates: list[dict[str, Any]],
+        begin_renderings: list[dict[str, Any]],
+        component_ids: set[str],
+    ) -> list[str]:
+        """Validate one surface's begin/update/root/child relationships."""
+        if not begin_renderings:
+            return [f"surfaceUpdate.surfaceId '{surface_id}' requires a matching beginRendering.surfaceId"]
+
+        errors: list[str] = []
+        root_id = begin_renderings[0].get("root")
+        if isinstance(root_id, str) and root_id and root_id not in component_ids:
+            errors.append(
+                f"beginRendering.surfaceId '{surface_id}' root '{root_id}' is missing from surfaceUpdate.components ids"
+            )
+
+        for update in updates:
+            errors.extend(self._validate_explicit_child_references(surface_id, update, component_ids))
+
+        return errors
+
+    def _validate_semantic_renderability(self, messages: list[dict[str, Any]]) -> list[str]:
+        """Return deterministic errors for schema-valid but non-renderable v0.8 message sequences."""
+        if self.schema_version != VERSION_0_8:
+            return []
+
+        begin_renderings, surface_updates, component_ids_by_surface = self._collect_surface_renderability_state(
+            messages
+        )
+        errors: list[str] = []
+
+        for surface_id, updates in surface_updates.items():
+            errors.extend(
+                self._validate_surface_renderability(
+                    surface_id=surface_id,
+                    updates=updates,
+                    begin_renderings=begin_renderings.get(surface_id, []),
+                    component_ids=component_ids_by_surface.get(surface_id, set()),
+                )
+            )
+
+        return list(dict.fromkeys(errors))
+
+    def _split_a2ui_response(self, content: str) -> tuple[list[str], list[str]]:
+        """Split tagged A2UI payloads from surrounding visible text."""
+        text_parts: list[str] = []
+        payload_chunks: list[str] = []
+        cursor = 0
+
+        while True:
+            open_index = content.find(A2UI_OPEN_TAG, cursor)
+            if open_index == -1:
+                text_parts.append(content[cursor:])
+                break
+
+            text_parts.append(content[cursor:open_index])
+            payload_start = open_index + len(A2UI_OPEN_TAG)
+            close_index = content.find(A2UI_CLOSE_TAG, payload_start)
+            if close_index == -1:
+                raise ValueError("Missing closing A2UI tag")
+
+            payload_chunks.append(content[payload_start:close_index])
+            cursor = close_index + len(A2UI_CLOSE_TAG)
+
+        return text_parts, payload_chunks
+
+    @staticmethod
+    def _parse_payload_chunk(payload_chunk: str) -> list[JsonValue]:
+        """Parse one raw A2UI JSON chunk into a flat list of messages."""
+        parsed = json.loads(payload_chunk)
+        if isinstance(parsed, dict):
+            return [parsed]
+        if not isinstance(parsed, list):
+            raise ValueError("A2UI payload must decode to a message object or list of messages")
+        return cast(list[JsonValue], parsed)
+
+    def _contains_a2ui_payload(self, content: str) -> bool:
+        """Return True when the response contains an A2UI payload marker."""
+        return A2UI_OPEN_TAG in content and A2UI_CLOSE_TAG in content
+
+    def _split_merged_messages(self, raw_messages: list[JsonValue]) -> tuple[list[JsonValue], list[dict[str, Any]]]:
+        """Canonicalize merged A2UI messages into one message per action key.
+
+        The A2UI protocol requires each streamed message to carry exactly one action key, but the
+        schema does not enforce it, so a model may emit one object that fuses several action keys
+        (e.g. v0.8 ``beginRendering`` + ``surfaceUpdate`` + ``dataModelUpdate``). Such an object is
+        rejected by the renderer even though it passes schema validation. Each action key has its
+        own disjoint body (no shared top-level fields other than the v0.9 sibling fields below), so
+        fanning a merged object out into one message per action key is lossless by construction:
+        every key's value is re-emitted unchanged, in canonical stream order, before validation.
+        v0.9 sibling fields (currently just ``version``) are not action keys and are re-attached to
+        every split message so each one stays independently valid.
+
+        Unknown top-level keys are rejected rather than dropped. This middleware validates with
+        ``remove_strict_validation`` applied (see ``__init__``), which strips
+        ``additionalProperties: false`` from the schema, so the validator will not reject an
+        unrecognized key on its own -- this helper is the only gate for it. Dropping the key would
+        let a likely-malformed payload pass as ``valid`` with no downstream backstop (the H-1
+        regression), so we raise and route it to ``validation.errors`` instead. An unknown key is
+        also a useful signal -- a hallucinated field or a mismatched schema version -- worth
+        surfacing rather than silently discarding. The allowed action-key and sibling-field sets are
+        selected from ``self.schema_version``; a message built for one version (e.g. a v0.9 message
+        carrying ``version``) is rejected under another version's middleware instance rather than
+        silently accepted.
+
+        ``raw_messages`` is the direct `json.loads(...)` output before schema validation, so list
+        elements may still be arbitrary JSON values rather than message dicts. Returns the
+        normalized raw payload plus a list of non-fatal repair notes describing any split performed.
+        """
+        action_key_set = _A2UI_MESSAGE_ACTION_KEYS_BY_VERSION.get(self.schema_version)
+        if action_key_set is None:
+            raise ValueError(f"Unsupported A2UI schema version for message normalization: {self.schema_version}")
+        sibling_field_set = _A2UI_SIBLING_FIELDS_BY_VERSION.get(self.schema_version, ())
+
+        normalized: list[JsonValue] = []
+        repairs: list[dict[str, Any]] = []
+
+        for index, raw_message in enumerate(raw_messages):
+            if not isinstance(raw_message, dict):
+                normalized.append(raw_message)
+                continue
+
+            action_keys = [key for key in action_key_set if key in raw_message]
+            sibling_fields = {field: raw_message[field] for field in sibling_field_set if field in raw_message}
+            unknown_keys = [key for key in raw_message if key not in action_key_set and key not in sibling_field_set]
+
+            if "version" in sibling_fields and sibling_fields["version"] != f"v{self.schema_version}":
+                raise ValueError(
+                    f"A2UI message version {sibling_fields['version']!r} does not match schema version v{self.schema_version}"
+                )
+
+            if unknown_keys:
+                # Strict validation is disabled for this schema (remove_strict_validation strips
+                # additionalProperties: false), so the validator will not reject these keys -- this
+                # raise is the only gate. Dropping them would let a malformed payload pass as
+                # `valid`; reject so it routes to validation.errors instead. See the docstring.
+                unknown_keys_text = ", ".join(sorted(unknown_keys))
+                raise ValueError(f"A2UI message contains unknown top-level keys: {unknown_keys_text}")
+
+            if not action_keys:
+                raise ValueError("A2UI message must contain exactly one action key")
+            if len(action_keys) == 1:
+                normalized.append(raw_message)
+                continue
+
+            normalized.extend({**sibling_fields, key: raw_message[key]} for key in action_keys)
+            repairs.append({"type": "split_merged_message", "from_index": index, "produced": len(action_keys)})
+
+        return normalized, repairs
+
+    def _build_validation_block(
+        self,
+        status: str,
+        fallback_reason: str | None = None,
+        errors: list[str] | None = None,
+        repairs: list[dict[str, Any]] | None = None,
+    ) -> dict[str, Any]:
+        """Build normalized validation result object.
+
+        ``repairs`` records non-fatal deterministic normalizations (e.g. merged-message splits) as
+        sidecar telemetry; it never changes ``status`` and is kept separate from ``errors`` so a
+        valid-but-repaired payload is not mistaken for a failed one.
+        """
+        block: dict[str, Any] = {"status": status}
+        if errors is not None:
+            block["errors"] = errors
+        if repairs:
+            block["repairs"] = repairs
+        if fallback_reason is not None and status == A2UIContentMetadataStatus.FALLBACK:
+            block["fallback_reason"] = fallback_reason
+        return block
+
+    def normalize_user_action(self, action: dict[str, Any]) -> A2UIActionValidationResult:
+        """Normalize raw renderer action payloads to canonical structure."""
+        if not self.context.enabled or not self.context.allow_user_actions:
+            return A2UIActionValidationResult(status=A2UIActionStatus.INVALID, errors=["user_actions_disabled"])
+
+        if not isinstance(action, dict):
+            return A2UIActionValidationResult(status=A2UIActionStatus.INVALID, errors=["action_payload_must_be_object"])
+
+        wrapped_action = action.get("userAction")
+        if wrapped_action is not None and not isinstance(wrapped_action, dict):
+            return A2UIActionValidationResult(status=A2UIActionStatus.INVALID, errors=["missing_userAction_payload"])
+        raw_action = wrapped_action if isinstance(wrapped_action, dict) else action
+
+        missing_fields = []
+        for field in _REQUIRED_ACTION_FIELDS:
+            if field not in raw_action:
+                missing_fields.append(f"missing_{field}")
+        if missing_fields:
+            return A2UIActionValidationResult(status=A2UIActionStatus.INVALID, errors=missing_fields)
+
+        if not isinstance(raw_action["context"], dict):
+            return A2UIActionValidationResult(status=A2UIActionStatus.INVALID, errors=["context_must_be_object"])
+
+        try:
+            normalized_action = A2UIActionContext.from_normalized_payload(
+                name=str(raw_action["name"]),
+                surface_id=str(raw_action["surfaceId"]),
+                source_component_id=str(raw_action["sourceComponentId"]),
+                timestamp=_normalize_iso_timestamp(str(raw_action["timestamp"])),
+                context=raw_action["context"],
+            )
+            return A2UIActionValidationResult(status=A2UIActionStatus.VALID, normalized_action=normalized_action)
+        except ValueError as exc:
+            return A2UIActionValidationResult(status=A2UIActionStatus.INVALID, errors=[str(exc)])