Migrate a2ui_agent to strict type checking and enable mypy in CI (#1903)

This change enables mypy type checking across the `a2ui_agent` package, rebased onto latest `main`.

Key changes:
- Enabled type checking for `a2ui_agent` in `pyproject.toml` (`[tool.mypy]`).
- Fixed PEP 484 compliance for re-exports in `a2ui_core` (`__init__.py`).
- Added missing type annotations (return types, variable types) across `a2ui_agent` schema, parser, and ADK integration modules.
- Added defensive validation checks (`isinstance(comps, list)`, `isinstance(comp_schema, dict)`) in `validator.py`.
- Integrated `uv run --all-packages mypy agent_sdks/python/a2ui_agent/src` into the consolidated `python-ci` job in `.github/workflows/python_ci.yml`.
- Aligned indentation and formatting (`pyink`) with upstream workspace rules.

TAG=agy
CONV=f0f6116c-47f0-466b-af69-e8d4ecabedd6

Author: jiahaog

.github/workflows/python_ci.yml

@@ -59,8 +59,9 @@ jobs:
           uv run pyink --check .
       - name: Type Check
         run: |
-          uv run mypy -p a2ui.core
-          uv run mypy agent_sdks/python/a2ui_core/scripts
+          uv run --all-packages mypy agent_sdks/python/a2ui_agent/src
+          uv run --all-packages mypy -p a2ui.core
+          uv run --all-packages mypy agent_sdks/python/a2ui_core/scripts
       - name: Run Unit Tests
         run: uv run pytest
       - name: Verify Load Real

agent_sdks/python/a2ui_agent/pyproject.toml

@@ -63,4 +63,5 @@ default = true
 dev = [
     "antlr4-tools>=0.2.1",
     "hatchling>=1.30.1",
+    "types-jsonschema",
 ]

agent_sdks/python/a2ui_agent/src/a2ui/a2a/extension.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import logging
-from typing import Optional, List
+from typing import Optional, List, Any, Dict
 
+from packaging.version import Version, parse as parse_version
 from a2a.server.agent_execution import RequestContext
 from a2a.types import AgentCard, AgentExtension
 
@@ -40,7 +41,7 @@ def get_a2ui_agent_extension(
     Returns:
         The configured A2UI AgentExtension.
     """
-    params = {}
+    params: Dict[str, Any] = {}
     if accepts_inline_catalogs:
         params[AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY] = (
             True  # Only set if not default of False
@@ -107,10 +108,8 @@ def _select_newest_a2ui_extension(
     if not matched_extensions:
         return None
 
-    def _version_key(uri: str) -> tuple:
+    def _version_key(uri: str) -> Version:
         version_str = uri.replace(f"{A2UI_EXTENSION_BASE_URI}/v", "")
-        from packaging.version import parse as parse_version
-
         return parse_version(version_str)
 
     return max(matched_extensions, key=_version_key)

agent_sdks/python/a2ui_agent/src/a2ui/a2a/parts.py

@@ -63,7 +63,7 @@ def is_a2ui_part(part: Part) -> bool:
     Returns:
         True if the part contains A2UI data, False otherwise.
     """
-    return (
+    return bool(
         isinstance(part.root, DataPart)
         and part.root.metadata
         and part.root.metadata.get(MIME_TYPE_KEY)
@@ -80,7 +80,7 @@ def get_a2ui_datapart(part: Part) -> Optional[DataPart]:
     Returns:
         The DataPart containing A2UI data if present, None otherwise.
     """
-    if is_a2ui_part(part):
+    if is_a2ui_part(part) and isinstance(part.root, DataPart):
         return part.root
     return None
 

agent_sdks/python/a2ui_agent/src/a2ui/adk/a2a/event_converter.py

@@ -33,6 +33,7 @@
   ```
 """
 
+from __future__ import annotations
 from typing import TYPE_CHECKING, Optional
 
 from a2ui.adk.a2a.part_converter import A2uiPartConverter
@@ -78,6 +79,7 @@ def __call__(
         )
 
         catalog = invocation_context.session.state.get(self._catalog_key)
+        effective_converter: GenAIPartToA2APartConverter
         if catalog:
             # Use the catalog-aware part converter
             effective_converter = A2uiPartConverter(

agent_sdks/python/a2ui_agent/src/a2ui/adk/a2a/part_converter.py

@@ -115,10 +115,9 @@ def convert(self, part: genai_types.Part) -> list[a2a_types.Part]:
                     logger.info("No result in A2UI tool response")
                     return []
 
-            # Handle generic/other tool responses that returned a string containing A2UI tags.
-            if function_response.response and function_response.response.get("result"):
+            if function_response.response:
                 result = function_response.response.get("result")
-                if has_a2ui_parts(result):
+                if isinstance(result, str) and has_a2ui_parts(result):
                     return parse_response_to_parts(
                         result,
                         validator=self._catalog.validator,

agent_sdks/python/a2ui_agent/src/a2ui/adk/send_a2ui_to_client_toolset.py

@@ -142,7 +142,8 @@ def __init__(
     ):
         super().__init__()
         self._a2ui_enabled = a2ui_enabled
-        self._ui_tools = [self._SendA2uiJsonToClientTool(a2ui_catalog, a2ui_examples)]
+        self._ui_tool = self._SendA2uiJsonToClientTool(a2ui_catalog, a2ui_examples)
+        self._ui_tools: list[base_tool.BaseTool] = [self._ui_tool]
 
     async def _resolve_a2ui_enabled(
         self, ctx: readonly_context.ReadonlyContext
@@ -196,7 +197,7 @@ async def get_part_converter(
         Returns:
             A configured A2uiPartConverter.
         """
-        catalog = await self._ui_tools[0]._resolve_a2ui_catalog(ctx)
+        catalog = await self._ui_tool._resolve_a2ui_catalog(ctx)
         return A2uiPartConverter(catalog)
 
     class _SendA2uiJsonToClientTool(base_tool.BaseTool):

agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming.py

@@ -18,7 +18,7 @@
 import json
 import logging
 import re
-from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING
+from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING, Union
 
 from .constants import *
 from ..schema.constants import (
@@ -30,11 +30,11 @@
     CATALOG_COMPONENTS_KEY,
 )
 from ..schema.validator import (
-    analyze_topology,
     extract_component_ref_fields,
     extract_component_required_fields,
 )
 from ..schema.validator import A2uiValidator
+from a2ui.core.validating import analyze_topology
 from .response_part import ResponsePart
 from a2ui.core.validating.validator import RELAXED_VALIDATION, STRICT_VALIDATION, ValidationConfig
 from a2ui.core import A2uiParseError, A2uiIntegrityError
@@ -53,7 +53,7 @@ class A2uiStreamParser:
     (V08 or V09) depending on the catalog version.
     """
 
-    def __new__(cls, catalog: A2uiCatalog):
+    def __new__(cls, catalog: A2uiCatalog) -> A2uiStreamParser:
         if cls is A2uiStreamParser:
             version = catalog.version
             # Lazy import inside __new__ to prevent circular import errors, as the
@@ -78,7 +78,7 @@ def __init__(self, catalog: A2uiCatalog):
         self._found_delimiter = False
         self._buffer = ""
         self._json_buffer = ""
-        self._brace_stack: List[Tuple[str, int]] = []
+        self._brace_stack: list[tuple[str, int]] = []
         self._brace_count = 0
         self._in_top_level_list = False
         self._in_string = False
@@ -143,7 +143,7 @@ def surface_id(self) -> Optional[str]:
         return self._surface_id
 
     @surface_id.setter
-    def surface_id(self, value: Optional[str]):
+    def surface_id(self, value: Optional[str]) -> None:
         self._surface_id = value
         if value is not None and self._unbound_root_id is not None:
             self._root_ids[value] = self._unbound_root_id
@@ -161,7 +161,7 @@ def root_id(self) -> Optional[str]:
         )
 
     @root_id.setter
-    def root_id(self, value: Optional[str]):
+    def root_id(self, value: Optional[str]) -> None:
         if self._surface_id:
             if value is not None:
                 self._root_ids[self._surface_id] = value
@@ -174,7 +174,7 @@ def root_id(self, value: Optional[str]):
     def msg_types(self) -> List[str]:
         return self._msg_types
 
-    def add_msg_type(self, msg_type: str):
+    def add_msg_type(self, msg_type: str) -> None:
         if msg_type not in self._msg_types:
             self._msg_types.append(msg_type)
         if msg_type in (
@@ -204,6 +204,14 @@ def _get_active_msg_type_for_components(self) -> Optional[str]:
             "Subclasses must implement _get_active_msg_type_for_components"
         )
 
+    def _construct_partial_message(
+        self, components: List[Dict[str, Any]], active_msg_type: str
+    ) -> Dict[str, Any]:
+        """Constructs a partial message of the correct type. Subclasses must implement."""
+        raise NotImplementedError(
+            "Subclasses must implement _construct_partial_message"
+        )
+
     def _deduplicate_data_model(self, m: Dict[str, Any]) -> bool:
         """Returns True if message should be yielded, False if skipped."""
         return True
@@ -213,7 +221,7 @@ def _yield_messages(
         messages_to_yield: List[Dict[str, Any]],
         messages: List[ResponsePart],
         config: ValidationConfig = STRICT_VALIDATION,
-    ):
+    ) -> None:
         """Validates and appends messages to the final output list."""
         for m in messages_to_yield:
             if not self._deduplicate_data_model(m):
@@ -358,7 +366,7 @@ def process_chunk(self, chunk: str) -> List[ResponsePart]:
             )
         return messages
 
-    def _reset_json_state(self):
+    def _reset_json_state(self) -> None:
         """Resets the JSON-specific parsing state (e.g., at the end of a block)."""
         self._json_buffer = ""
         self._brace_stack = []
@@ -450,7 +458,7 @@ def _fix_json(self, fragment: str) -> str:
 
         return fixed
 
-    def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]):
+    def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]) -> None:
         for char in chunk:
             char_handled = False
             if self._brace_count == 0:
@@ -670,24 +678,27 @@ def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None:
                                     or "default"
                                 )
                                 # Deduplicate delta_contents by only keeping the LATEST entry for each dirty key
+                                delta_contents: Union[
+                                    List[Dict[str, Any]], Dict[str, Any]
+                                ]
                                 if isinstance(raw_contents, list):
-                                    delta_contents = []
+                                    list_contents: List[Dict[str, Any]] = []
                                     seen_keys = set()
                                     for entry in reversed(raw_contents):
                                         if not isinstance(entry, dict):
                                             continue
-                                        k = entry.get("key")
+                                        entry_key = entry.get("key")
                                         # Only include entries that have a valid parsed key (cumulative)
                                         if (
-                                            k
-                                            and k in contents_dict
-                                            and k not in seen_keys
+                                            entry_key
+                                            and entry_key in contents_dict
+                                            and entry_key not in seen_keys
                                         ):
-                                            delta_contents.insert(0, entry)
-                                            seen_keys.add(k)
+                                            list_contents.insert(0, entry)
+                                            seen_keys.add(entry_key)
                                     delta_contents = (
                                         self._prune_incomplete_datamodel_entries(
-                                            delta_contents
+                                            list_contents
                                         )
                                     )
                                 else:
@@ -711,7 +722,7 @@ def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None:
                                 # Update internal model for path resolution
                                 self.update_data_model(dm_obj, messages)
 
-    def _sniff_partial_component(self, messages: List[ResponsePart]):
+    def _sniff_partial_component(self, messages: List[ResponsePart]) -> None:
         """Attempts to parse a partial component from the current buffer."""
         # We only care about components if we are inside a "components" array
         if f'"{CATALOG_COMPONENTS_KEY}"' not in self._json_buffer:
@@ -779,7 +790,7 @@ def _prune_incomplete_datamodel_entries(self, entries: Any) -> Any:
 
     def _handle_partial_component(
         self, comp: Dict[str, Any], messages: List[ResponsePart]
-    ):
+    ) -> None:
         """Handles a component discovered before its parent message is finished.
 
         When the parser sees a full JSON object that looks like a component
@@ -880,10 +891,10 @@ def _handle_complete_object(
 
     def yield_reachable(
         self,
-        messages: List[Dict[str, Any]],
+        messages: List[ResponsePart],
         check_root: bool = False,
         raise_on_orphans: bool = False,
-    ):
+    ) -> None:
         """Yields a partial message containing all reachable and seen components.
 
         This is the core of the streaming logic. Instead of waiting for a UI message
@@ -934,8 +945,8 @@ def yield_reachable(
                 )
 
             # 1. Process placeholders and partial children
-            processed_components = []
-            extra_components = []
+            processed_components: List[Dict[str, Any]] = []
+            extra_components: List[Dict[str, Any]] = []
             surface_id = self.surface_id or "unknown"
             yielded_for_surface = self._yielded_ids.get(surface_id, set())
 
@@ -1022,7 +1033,7 @@ def _process_component_topology(
         comp: Dict[str, Any],
         extra_components: List[Dict[str, Any]],
         inline_resolved: bool = False,
-    ):
+    ) -> None:
         """Recursively processes path placeholders and child pruning in one pass."""
         comp_id = comp.get("id", "unknown")
 
@@ -1033,7 +1044,7 @@ def _process_component_topology(
             else ""
         )
 
-        def traverse(obj, parent_key=None):
+        def traverse(obj: Any, parent_key: Optional[str] = None) -> None:
             if isinstance(obj, dict):
                 # 1. Handle Path Placeholders (from _apply_placeholders)
                 if (

agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v08.py

@@ -68,7 +68,7 @@ def _data_model_msg_type(self) -> str:
         """Returns the message type identifier for data model updates."""
         return MSG_TYPE_DATA_MODEL_UPDATE
 
-    def _sniff_metadata(self):
+    def _sniff_metadata(self) -> None:
         """Sniffs for v0.8 metadata in the json_buffer."""
 
         def get_latest_value(key: str) -> Optional[str]:

agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v09.py

@@ -60,7 +60,7 @@ def is_protocol_msg(self, obj: Dict[str, Any]) -> bool:
             )
         )
 
-    def _sniff_metadata(self):
+    def _sniff_metadata(self) -> None:
         """Sniffs for v0.9 metadata in the json_buffer."""
 
         def get_latest_value(key: str) -> Optional[str]:
@@ -237,7 +237,7 @@ def _construct_partial_message(
         self, processed_components: List[Dict[str, Any]], active_msg_type: str
     ) -> Dict[str, Any]:
         """Constructs a partial message for v0.9 (updateComponents)."""
-        payload = {
+        payload: Dict[str, Any] = {
             CATALOG_COMPONENTS_KEY: processed_components,
         }
         if self.surface_id: