A2UISchemaManager

def get_component_schemas(self) -> dict[str, dict[str, Any]]:
    """Get resolved schemas for all components, keyed by component type name.

    Returns:
        A dict mapping component type names (e.g. "Button") to their
        JSON Schema definitions from the catalog ``components`` section.
    """
    schemas: dict[str, dict[str, Any]] = {}
    for catalog in [self._basic_catalog, self._custom_catalog]:
        if catalog is not None:
            components = catalog.get("components", {})
            for name, defn in components.items():
                if isinstance(defn, dict) and name not in schemas:
                    schemas[name] = defn
    return schemas

def build_schema_registry(self) -> Any:
    """Build a jsonschema referencing Registry for cross-file $ref resolution.

    The v0.9 ``server_to_client.json`` schema references ``catalog.json``
    and ``common_types.json`` via relative ``$ref``. This method creates a
    registry that maps those URIs to the loaded schemas so that
    ``jsonschema.validate()`` can resolve them locally.

    When a custom catalog is provided, it is registered as ``catalog.json``
    so that the envelope schema validates against the custom components.

    Returns:
        A ``referencing.Registry`` instance.
    """
    active_catalog = self._get_active_catalog()
    base = _VERSION_CONFIG[self._protocol_version]["schema_base_uri"]
    resources: list[tuple[str, Any]] = [
        (f"{base}catalog.json", Resource.from_contents(active_catalog, default_specification=DRAFT202012)),
        (f"{base}common_types.json", Resource.from_contents(self._common_types, default_specification=DRAFT202012)),
        ("catalog.json", Resource.from_contents(active_catalog, default_specification=DRAFT202012)),
        ("common_types.json", Resource.from_contents(self._common_types, default_specification=DRAFT202012)),
        (
            f"{base}basic_catalog.json",
            Resource.from_contents(self._basic_catalog, default_specification=DRAFT202012),
        ),
        ("basic_catalog.json", Resource.from_contents(self._basic_catalog, default_specification=DRAFT202012)),
    ]
    for schema in [self._server_to_client, self._basic_catalog, self._common_types]:
        schema_id = schema.get("$id")
        if schema_id:
            resources.append((schema_id, Resource.from_contents(schema, default_specification=DRAFT202012)))
    if self._custom_catalog is not None:
        custom_id = self._custom_catalog.get("$id")
        if custom_id:
            resources.append((
                custom_id,
                Resource.from_contents(self._custom_catalog, default_specification=DRAFT202012),
            ))
    return Registry().with_resources(resources)

def generate_prompt_section(
    self,
    include_schema: bool = True,
    include_rules: bool = True,
    response_delimiter: str = A2UI_DEFAULT_DELIMITER,
    actions: list[Any] | None = None,
) -> str:
    """Generate the A2UI portion of the system prompt.

    Args:
        include_schema: Whether to include the JSON schema for validation reference.
        include_rules: Whether to include the plain-text catalog rules.
        response_delimiter: The delimiter the LLM should use to separate text from A2UI JSON.
        actions: List of ``A2UIAction`` objects. If provided, an "Available Actions"
            section is injected into the prompt listing the actions and how to
            create buttons for them.

    Returns:
        A string to append to the agent's system prompt.
    """
    v = self._protocol_version
    sections: list[str] = []

    example_json = self._build_prompt_example()

    sections.append(
        f"## A2UI Response Format ({v})\n\n"
        f"You can generate rich UI responses using the A2UI {v} protocol. "
        "When you want to display a UI, respond with your text message first, "
        f"then the delimiter `{response_delimiter}`, "
        "then a JSON array of A2UI message objects.\n\n"
        "Example response format:\n"
        "```\n"
        "Here is the UI you requested.\n"
        f"{response_delimiter}\n"
        f"{example_json}\n"
        "```"
    )

    message_types = self._prompt_message_types.replace("{catalog_id}", self._catalog_id).replace(
        "{version_string}", v
    )
    sections.append(f"\n\n## A2UI Message Types\n\n{message_types}")

    all_components = self._get_all_components()
    if all_components:
        basic_components = sorted(self._basic_catalog.get("components", {}).keys())
        custom_components = (
            sorted(self._custom_catalog.get("components", {}).keys()) if self._custom_catalog else []
        )

        comp_section = "\n\n## Available Components\n\n"
        if basic_components:
            comp_section += f"**Basic catalog:** {', '.join(basic_components)}\n\n"
        if custom_components:
            comp_section += f"**Custom components:** {', '.join(custom_components)}\n\n"
            for comp_name in custom_components:
                comp_def = self._custom_catalog.get("components", {}).get(comp_name, {})  # type: ignore[union-attr]
                desc = self._extract_component_description(comp_name, comp_def)
                if desc:
                    comp_section += f"- **{comp_name}**: {desc}\n"
            comp_section += "\n"

        comp_section += (
            "Each component uses a flat discriminator format:\n"
            '```json\n{"id": "myId", "component": "Text", "text": "Hello"}\n```\n\n'
            "Components are referenced by `id`. Use `children` (array of ids) for layout components (Column, Row). "
            "Use `child` (single id) for wrapper components (Card, Button)."
        )
        sections.append(comp_section)

    if include_rules:
        rules_parts: list[str] = []
        if self._catalog_rules:
            rules_parts.append(self._catalog_rules)
        if self._custom_catalog_rules:
            rules_parts.append(self._custom_catalog_rules)
        if rules_parts:
            sections.append("\n\n## Component Rules\n\n" + "\n\n".join(rules_parts))

    if actions:
        action_lines = ["\n\n## Available Actions\n"]
        action_lines.append("The following actions can be triggered by buttons in the UI:\n")

        event_actions = [a for a in actions if a.action_type == "event"]
        func_actions = [a for a in actions if a.action_type == "functionCall"]

        if event_actions:
            action_lines.append("### Server Events\n")
            for a in event_actions:
                desc = f": {a.description}" if a.description else ""
                action_lines.append(f"- `{a.name}`{desc}")
                if a.example_context:
                    action_lines.append(f"  Context: {json.dumps(a.example_context)}")
            action_lines.append(
                "\nTo create a button that triggers a server event:\n"
                "```json\n"
                '{"id": "action_btn", "component": "Button", "child": "action_btn_text", '
                '"action": {"event": {"name": "<action_name>", "context": {...}}}},\n'
                '{"id": "action_btn_text", "component": "Text", "text": "Button Label"}\n'
                "```"
            )

        if func_actions:
            action_lines.append("\n### Client Functions\n")
            action_lines.append(
                "Client functions execute on the client without a server round-trip. "
                "Use the **exact** function name, argument structure, and returnType shown below. "
                "Do not add extra properties.\n"
            )
            for a in func_actions:
                desc = f": {a.description}" if a.description else ""
                example_args = json.dumps(a.example_args) if a.example_args else "{}"
                action_lines.append(f"- **`{a.name}`**{desc}")
                action_lines.append(
                    f"  ```json\n"
                    f'  {{"id": "action_btn", "component": "Button", "child": "action_btn_text", '
                    f'"action": {{"functionCall": {{"call": "{a.name}", "args": {example_args}, "returnType": "void"}}}}}}\n'
                    f"  ```"
                )

        sections.append("\n".join(action_lines))

    if include_schema:
        schema_str = json.dumps(self._server_to_client, indent=2)
        sections.append(f"\n\n## A2UI Message Schema ({v})\n\n```json\n{schema_str}\n```")

    return "\n".join(sections)