From 663932b8247d49ca6078e693b3104a673a4b4021 Mon Sep 17 00:00:00 2001 From: Nan Yu Date: Thu, 26 Mar 2026 16:52:48 +0000 Subject: [PATCH] feat: implement v0.9 streaming parser --- .../python/src/a2ui/core/parser/streaming.py | 563 ++--- .../src/a2ui/core/parser/streaming_v08.py | 232 ++ .../src/a2ui/core/parser/streaming_v09.py | 262 +++ .../src/a2ui/core/parser/version_handlers.py | 266 --- .../tests/core/parser/test_streaming_v08.py | 32 +- .../tests/core/parser/test_streaming_v09.py | 2034 +++++++++++++++++ .../core/parser/test_version_handlers.py | 133 -- 7 files changed, 2708 insertions(+), 814 deletions(-) create mode 100644 agent_sdks/python/src/a2ui/core/parser/streaming_v08.py create mode 100644 agent_sdks/python/src/a2ui/core/parser/streaming_v09.py delete mode 100644 agent_sdks/python/src/a2ui/core/parser/version_handlers.py create mode 100644 agent_sdks/python/tests/core/parser/test_streaming_v09.py delete mode 100644 agent_sdks/python/tests/core/parser/test_version_handlers.py diff --git a/agent_sdks/python/src/a2ui/core/parser/streaming.py b/agent_sdks/python/src/a2ui/core/parser/streaming.py index 61e2cf8fc..5a23ffe49 100644 --- a/agent_sdks/python/src/a2ui/core/parser/streaming.py +++ b/agent_sdks/python/src/a2ui/core/parser/streaming.py @@ -19,13 +19,19 @@ from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING from .constants import * -from ..schema.constants import VERSION_0_9, VERSION_0_8, A2UI_OPEN_TAG, A2UI_CLOSE_TAG, SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY +from ..schema.constants import ( + VERSION_0_9, + VERSION_0_8, + A2UI_OPEN_TAG, + A2UI_CLOSE_TAG, + SURFACE_ID_KEY, + CATALOG_COMPONENTS_KEY, +) from ..schema.validator import ( analyze_topology, extract_component_ref_fields, extract_component_required_fields, ) -from .version_handlers import A2uiV08Handler, A2uiV09Handler, A2uiVersionHandler from .response_part import ResponsePart @@ -48,16 +54,26 @@ "text", } -# A safe placeholder used when referencing a child component that hasn't yet streamed in. -PLACEHOLDER_COMPONENT = { - "Row": { - "children": {"explicitList": []}, - } -} - class A2uiStreamParser: - """Parses a stream of text for A2UI JSON messages with fine-grained component yielding.""" + """Parses a stream of text for A2UI JSON messages with fine-grained component yielding. + + This class acts as a factory that returns a version-specific parser instance + (V08 or V09) depending on the catalog version. + """ + + def __new__(cls, catalog: "A2uiCatalog" = None, *args, **kwargs): + if cls is A2uiStreamParser: + version = getattr(catalog, "version", None) if catalog else None + if version == VERSION_0_9: + from .streaming_v09 import A2uiStreamParserV09 + + return A2uiStreamParserV09(catalog=catalog, *args, **kwargs) + else: + from .streaming_v08 import A2uiStreamParserV08 + + return A2uiStreamParserV08(catalog=catalog, *args, **kwargs) + return super().__new__(cls) def __init__(self, catalog: "A2uiCatalog" = None): self._ref_fields_map = extract_component_ref_fields(catalog) if catalog else {} @@ -80,55 +96,45 @@ def __init__(self, catalog: "A2uiCatalog" = None): self._seen_components: Dict[str, Dict[str, Any]] = {} # Track data model for path resolution - self._data_model: Dict[str, Any] = {} self._yielded_data_model: Dict[str, Any] = {} - self._yielded_placeholders: Dict[str, Set[str]] = {} # comp_id -> set of paths self._deleted_surfaces: Set[str] = set() - self._comp_paths: Dict[str, Set[str]] = {} # cid -> set of data model paths - # Set of unique component IDs yielded per surface - self._yielded_ids: Dict[str, Set[str]] = {} # surfaceId -> set of cids - self._yielded_contents: Dict[Any, str] = {} # (surfaceId, cid) -> hash of content - self._components_with_placeholders: Set[str] = ( - set() - ) # cid set that were yielded as placeholders - - self._handler: Optional[A2uiVersionHandler] = None - if catalog: - version = getattr(catalog, "version", None) - if version == VERSION_0_9: - self._handler = A2uiV09Handler() - elif version == VERSION_0_8: - self._handler = A2uiV08Handler() - - self._root_id: Optional[str] = ( - DEFAULT_ROOT_ID - if catalog and getattr(catalog, "version", None) == VERSION_0_9 - else None - ) - self._surface_id: Optional[str] = None - self._msg_types: List[str] = [] - self._yielded_begin_rendering_surfaces: Set[str] = set() - self._yielded_real_begin_rendering_surfaces: Set[str] = set() + + # Set of unique component IDs yielded per surface to prevent duplicate yielding + # surfaceId -> set of cids + self._yielded_ids: Dict[str, Set[str]] = {} + # (surfaceId, cid) -> hash of content for change detection + self._yielded_contents: Dict[Any, str] = {} + + self._root_id: Optional[str] = None # The root component ID for the layout tree + self._surface_id: Optional[str] = None # The active surface ID tracking the context + self._msg_types: List[str] = [] # Running list of message types seen in the block + + # A set of surface ids for which we have already yielded a start message + # Tracks if beginRendering or createSurface was emitted + self._yielded_start_messages: Set[str] = set() + + # The current active message type for component grouping self._active_msg_type: Optional[str] = None - # New state for buffering and progress + # State for buffering updates until surface is ready self._pending_messages: Dict[str, List[Dict[str, Any]]] = ( {} - ) # surfaceId -> list of msgs - self._buffered_begin_rendering: Optional[Dict[str, Any]] = None - self._topology_dirty = False - - @property - def seen_components(self) -> Dict[str, Dict[str, Any]]: - return self._seen_components + ) # surfaceId -> list of msgs delayed until start message arrives + self._buffered_start_message: Optional[Dict[str, Any]] = ( + None # The start message to yield before any components + ) + self._topology_dirty = False # Set to true if components are added out of order + self._in_top_level_list = False @property - def buffered_begin_rendering(self) -> Optional[Dict[str, Any]]: - return self._buffered_begin_rendering + def _placeholder_component(self) -> Dict[str, Any]: + """Returns the version-specific placeholder component. - @buffered_begin_rendering.setter - def buffered_begin_rendering(self, value: Optional[Dict[str, Any]]): - self._buffered_begin_rendering = value + This is used when a component references a child component that hasn't yet + streamed in. The placeholder component is added to the components list and + the reference is updated to point to the placeholder component. + """ + raise NotImplementedError("Subclasses must implement _placeholder_component") @property def surface_id(self) -> Optional[str]: @@ -160,20 +166,29 @@ def add_msg_type(self, msg_type: str): ): self._active_msg_type = msg_type + @property + def _yielded_surfaces_set(self) -> Set[str]: + """Provides access to version-specific yielded surfaces set.""" + raise NotImplementedError("Subclasses must implement _yielded_surfaces_set") + + def is_protocol_msg(self, obj: Dict[str, Any]) -> bool: + """Checks if the object is a recognized A2UI message for this version.""" + raise NotImplementedError("Subclasses must implement is_protocol_msg") + + @property + def _data_model_msg_type(self) -> str: + """Returns the message type identifier for data model updates.""" + raise NotImplementedError("Subclasses must implement _data_model_msg_type") + def _get_active_msg_type_for_components(self) -> Optional[str]: """Determines which msg_type to use when wrapping component updates.""" - if self._active_msg_type: - return self._active_msg_type - for mt in self._msg_types: - if mt in ( - MSG_TYPE_SURFACE_UPDATE, - MSG_TYPE_UPDATE_COMPONENTS, - MSG_TYPE_CREATE_SURFACE, - ): - self._active_msg_type = mt - return mt - # Fallback to the first found message type or default - return self._msg_types[0] if self._msg_types else None + raise NotImplementedError( + "Subclasses must implement _get_active_msg_type_for_components" + ) + + def _deduplicate_data_model(self, m: Dict[str, Any], strict_integrity: bool) -> bool: + """Returns True if message should be yielded, False if skipped.""" + return True def _yield_messages( self, @@ -183,56 +198,8 @@ def _yield_messages( ): """Validates and appends messages to the final output list.""" for m in messages_to_yield: - # Deduplicate dataModelUpdate (v0.8) - if MSG_TYPE_DATA_MODEL_UPDATE in m: - dm = m[MSG_TYPE_DATA_MODEL_UPDATE] - raw_contents = dm.get("contents", {}) - contents_dict = {} - if isinstance(raw_contents, list): - for entry in raw_contents: - if isinstance(entry, dict) and "key" in entry: - key = entry["key"] - val = ( - entry.get("valueString") - or entry.get("valueNumber") - or entry.get("valueBoolean") - or entry.get("valueMap") - ) - if key and val is not None: - contents_dict[key] = val - elif isinstance(raw_contents, dict): - contents_dict = raw_contents - - if contents_dict: - is_new = False - for k, v in contents_dict.items(): - if self._yielded_data_model.get(k) != v: - is_new = True - break - if not is_new and strict_integrity: - # Only skip if strict (complete message) and truly old - continue - # Note: we don't update self._yielded_data_model here if it was already updated by sniff - self._yielded_data_model.update(contents_dict) - - # Deduplicate updateDataModel (v0.9) - if MSG_TYPE_UPDATE_DATA_MODEL in m: - udm = m[MSG_TYPE_UPDATE_DATA_MODEL] - if isinstance(udm, dict): - is_new = False - for k, v in udm.items(): - if ( - k not in (SURFACE_ID_KEY, "root") - and self._yielded_data_model.get(k) != v - ): - is_new = True - break - if not is_new and strict_integrity: - continue - # Update yielded model - for k, v in udm.items(): - if k not in (SURFACE_ID_KEY, "root"): - self._yielded_data_model[k] = v + if not self._deduplicate_data_model(m, strict_integrity): + continue # Each surface update message must specify a surfaceId and satisfy catalog validation. if self._validator: @@ -242,11 +209,8 @@ def _yield_messages( ) except ValueError as e: if strict_integrity: - # Fatal error: the agent produced invalid A2UI. raise e else: - # Sniffing error: the fragment is probably not yet complete enough for A2UI. - # We swallow this and wait for more data. logger.debug(f"Validation failed for partial/sniffed message: {e}") continue @@ -261,18 +225,14 @@ def _yield_messages( def _delete_surface(self, sid: str) -> None: """Clears all state related to a specific surface.""" self._pending_messages.pop(sid, None) - # Clear component paths and placeholders for components yielded on this surface - yielded_cids = self._yielded_ids.pop(sid, set()) - for cid in yielded_cids: - self._comp_paths.pop(cid, None) - self._yielded_placeholders.pop(cid, None) + self._yielded_ids.pop(sid, None) # Clear contents for this surface self._yielded_contents = { k: v for k, v in self._yielded_contents.items() if k[0] != sid } - self._yielded_begin_rendering_surfaces.discard(sid) - self._yielded_real_begin_rendering_surfaces.discard(sid) + self._yielded_surfaces_set.discard(sid) + self._yielded_start_messages.discard(sid) self._deleted_surfaces.add(sid) @@ -459,22 +419,6 @@ def _fix_json(self, fragment: str) -> str: return fixed - def _process_json_chunk(self, chunk: str, messages: List[Dict[str, Any]]): - """Processes raw JSON characters and manages the brace stack. - - This method implements a lightweight stream-parsing state machine using a - brace stack. It identifies potential JSON objects within the top-level - A2UI message array (`[...]`). - - When an object is fully closed (brace count returns to zero), it triggers - `_handle_complete_object`. While an object is nested, it continuously - triggers `_sniff_metadata` to discover identifiers as early as possible. - - Args: - chunk: The raw string of bytes suspected to be part of the JSON array. - messages: The list to which newly parsed A2UI messages will be added. - """ - def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]): for char in chunk: char_handled = False @@ -543,15 +487,12 @@ def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]): try: obj = json.loads(obj_buffer) if isinstance(obj, dict): - is_v08_msg = self._in_top_level_list and any( - k in obj - for k in ( - MSG_TYPE_BEGIN_RENDERING, - MSG_TYPE_SURFACE_UPDATE, - MSG_TYPE_DATA_MODEL_UPDATE, - MSG_TYPE_DELETE_SURFACE, - ) + logger.debug( + f"[Parsed Dict] Keys: {list(obj.keys())}, protocol check" + " follows..." ) + + is_protocol = self._in_top_level_list and self.is_protocol_msg(obj) is_comp = obj.get("id") and obj.get("component") # Process objects at top-level OR items in top-level list # When in a list, we are top-level if the ONLY thing on the stack is the list opener @@ -562,7 +503,7 @@ def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]): ) if is_comp: self._handle_partial_component(obj, messages) - elif is_top_level or is_v08_msg: + elif is_top_level or is_protocol: if not self._handle_complete_object( obj, self.surface_id, messages ): @@ -570,21 +511,26 @@ def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]): self._yield_messages([obj], messages) if self._brace_count == 0 or ( - self._in_top_level_list and self._brace_count == 1 + self._in_top_level_list and len(self._brace_stack) == 1 ): # Aggressively clear processed objects from the buffer to prevent slowdown. - # We slice up to the current position. - self._json_buffer = self._json_buffer[len(obj_buffer) :] - if self._brace_stack: - # Adjust stack indices after shift - shift = len(obj_buffer) - self._brace_stack = [ - (b_t, i - shift) for b_t, i in self._brace_stack - ] + if len(self._brace_stack) == 1 and self._brace_stack[0][0] == "[": + # Keep '[' and remove the object after it + self._json_buffer = ( + self._json_buffer[:start_idx] + + self._json_buffer[start_idx + len(obj_buffer) :] + ) else: - self._brace_stack = [] + self._json_buffer = self._json_buffer[len(obj_buffer) :] + if self._brace_stack: + shift = len(obj_buffer) + self._brace_stack = [ + (b_t, i - shift) for b_t, i in self._brace_stack + ] + except json.JSONDecodeError as e: logger.debug(f"Object recognition failed: {e}") + elif char == "[": self._brace_stack.append(("[", len(self._json_buffer))) self._json_buffer += "[" @@ -614,12 +560,15 @@ def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]): self.yield_reachable(messages, check_root=False, raise_on_orphans=False) self._topology_dirty = False + def _construct_sniffed_data_model_message( + self, active_msg_type: str, delta_msg_payload: Dict[str, Any] + ) -> Dict[str, Any]: + """Returns the message to yield for a partial data model update.""" + return {active_msg_type: delta_msg_payload} + def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None: - """Sniffs for partial data model updates in the buffer.""" - if ( - f'"{MSG_TYPE_DATA_MODEL_UPDATE}"' not in self._json_buffer - and f'"{MSG_TYPE_UPDATE_DATA_MODEL}"' not in self._json_buffer - ): + msg_type = self._data_model_msg_type + if f'"{msg_type}"' not in self._json_buffer: return # Look through the brace stack for objects that might contain data model updates for b_type, start_idx in reversed(self._brace_stack): @@ -649,10 +598,15 @@ def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None: continue if obj and isinstance(obj, dict): - # v0.8: dataModelUpdate - if MSG_TYPE_DATA_MODEL_UPDATE in obj: - dm_obj = obj[MSG_TYPE_DATA_MODEL_UPDATE] + active_msg_type = None + msg_type = self._data_model_msg_type + if msg_type in obj: + active_msg_type = msg_type + + if active_msg_type: + dm_obj = obj[active_msg_type] if isinstance(dm_obj, dict) and "contents" in dm_obj: + raw_contents = dm_obj["contents"] contents_dict = self._parse_contents_to_dict(raw_contents) @@ -689,19 +643,15 @@ def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None: if "path" in dm_obj: delta_msg_payload["path"] = dm_obj["path"] - delta_msg = {MSG_TYPE_DATA_MODEL_UPDATE: delta_msg_payload} + delta_msg = self._construct_sniffed_data_model_message( + active_msg_type, delta_msg_payload + ) self._yield_messages([delta_msg], messages, strict_integrity=False) + self._yielded_data_model.update(contents_dict) # Update internal model for path resolution self.update_data_model(dm_obj, messages) - # v0.9: updateDataModel - elif MSG_TYPE_UPDATE_DATA_MODEL in obj: - dm_obj = obj[MSG_TYPE_UPDATE_DATA_MODEL] - if isinstance(dm_obj, dict): - self.update_data_model(dm_obj, messages) - return - def _sniff_partial_component(self, messages: List[ResponsePart]): """Attempts to parse a partial component from the current buffer.""" # We only care about components if we are inside a "components" array @@ -719,49 +669,19 @@ def _sniff_partial_component(self, messages: List[ResponsePart]): try: obj = json.loads(fixed_fragment) if isinstance(obj, dict) and obj.get("id") and obj.get("component"): - # Ignore components that are effectively empty (no type keys) - if isinstance(obj["component"], dict) and len(obj["component"]) > 0: + if isinstance(obj["component"], str): + # Flat style (v0.9+): component type is a string + self._handle_partial_component(obj, messages) + elif isinstance(obj["component"], dict) and len(obj["component"]) > 0: + # Structured style (v0.8): Ignore components that are effectively empty (no type keys) self._handle_partial_component(obj, messages) except Exception: continue - def _sniff_metadata(self): - """Sniffs for surfaceId, root, and msg_types in the current json_buffer. - - This method is called frequently during the parsing of an object fragment. - It attempts to detect the A2UI version (if not already known) and delegates - to the version handler to perform regex-based "sniffing" of keys like - `surfaceId` or `root`. - - Why we need it: - In a streaming context, we want to know *which* UI surface we are updating - long before the closing brace of the message arrives. This allows us to - yield components that have been fully received even if the rest of the - message (e.g., more components or metadata) is still on the wire. - """ - if not self._handler: - version = A2uiVersionHandler.detect_version(self._json_buffer) - if version == VERSION_0_8: - self._handler = A2uiV08Handler() - else: - # Default to 0.9. - # Create new handler if the version is not compatible with 0.9. - self._handler = A2uiV09Handler() - - if self._handler: - self._handler.sniff_metadata(self._json_buffer, self) - else: - # Fallback to generic sniffing if version not yet known - if not self.surface_id: - match = re.search(rf'"{SURFACE_ID_KEY}":\s*"([^"]+)"', self._json_buffer) - if match: - self.surface_id = match.group(1) - - if not self.root_id: - match = re.search(r'"root":\s*"([^"]+)"', self._json_buffer) - if match: - self.root_id = match.group(1) + def _sniff_metadata(self) -> None: + """Sniffs for surfaceId, root, and msg_types in the current json_buffer.""" + raise NotImplementedError("Subclasses must implement _sniff_metadata") def _prune_incomplete_datamodel_entries(self, entries: Any) -> Any: """Recursively removes data model entries that only contain 'key' and no valid values.""" @@ -826,7 +746,12 @@ def _has_empty_dict(obj: Any) -> bool: return False component_def = comp.get("component") - if _has_empty_dict(component_def): + if isinstance(component_def, str): + # v0.9 flat style: check the whole component object for empty dicts + if _has_empty_dict(comp): + return + elif _has_empty_dict(component_def): + # v0.8 nested style: check properties inside component return if isinstance(component_def, dict) and hasattr(self, "_required_fields_map"): @@ -883,113 +808,14 @@ def update_data_model( if k not in (SURFACE_ID_KEY, "root", "contents") } - self._data_model.update(contents) - updated_keys = set(contents.keys()) - def _handle_complete_object( self, obj: Dict[str, Any], sid: Optional[str], messages: List[ResponsePart], ) -> bool: - """Handles a fully-closed top-level A2UI message object. - - This is called when a top-level object (e.g., `beginRendering`, - `surfaceUpdate`, `deleteSurface`) is fully parsed. It finalizes the parser - state for that message and allows version handlers to perform final - processing (like yielding any remaining buffered components). - - Args: - obj: The fully parsed message object. - messages: The list of final A2UI messages to be delivered to the client. - """ - if not isinstance(obj, dict): - return - - # Validate against the schema - if self._validator: - self._validator.validate(obj, root_id=sid, strict_integrity=False) - - # Update state based on the message content - surface_id = obj.get(SURFACE_ID_KEY, self.surface_id) - if MSG_TYPE_SURFACE_UPDATE in obj: - val = obj[MSG_TYPE_SURFACE_UPDATE] - if isinstance(val, dict): - surface_id = val.get(SURFACE_ID_KEY) or surface_id - elif MSG_TYPE_BEGIN_RENDERING in obj: - val = obj[MSG_TYPE_BEGIN_RENDERING] - if isinstance(val, dict): - surface_id = val.get(SURFACE_ID_KEY) or surface_id - elif MSG_TYPE_CREATE_SURFACE in obj: - val = obj[MSG_TYPE_CREATE_SURFACE] - if isinstance(val, dict): - surface_id = val.get(SURFACE_ID_KEY) or surface_id - elif MSG_TYPE_DELETE_SURFACE in obj: - val = obj[MSG_TYPE_DELETE_SURFACE] - if isinstance(val, str): - surface_id = val - elif isinstance(val, dict): - surface_id = val.get(SURFACE_ID_KEY) or surface_id - - self.surface_id = surface_id - sid = self.surface_id or "unknown" - - if MSG_TYPE_DELETE_SURFACE in obj: - # Only process actual deletion if we've already established the surface. - # Otherwise, we buffer it (handled below) so it can be processed after beginRendering. - if ( - sid in self._yielded_begin_rendering_surfaces - or self._buffered_begin_rendering - ): - self._delete_surface(sid) - - # Handle buffering if beginRendering is missing - if sid in self._deleted_surfaces: - return True - - if ( - (MSG_TYPE_SURFACE_UPDATE in obj or MSG_TYPE_DELETE_SURFACE in obj) - and sid not in self._yielded_begin_rendering_surfaces - and not self._buffered_begin_rendering - ): - if sid not in self._pending_messages: - self._pending_messages[sid] = [] - self._pending_messages[sid].append(obj) - - # deleteSurface should NOT trigger progress - return True - - # Let the handler try to deal with it first - if self._handler and self._handler.handle_complete_object(obj, self, messages): - # If beginRendering just completed, flush pending messages - if MSG_TYPE_BEGIN_RENDERING in obj or MSG_TYPE_CREATE_SURFACE in obj: - # ... - # Yield beginRendering immediately when it completes - if sid not in self._yielded_real_begin_rendering_surfaces: - self._yield_messages([obj], messages) - self._yielded_real_begin_rendering_surfaces.add(sid) - self._yielded_begin_rendering_surfaces.add(sid) - - self._buffered_begin_rendering = None - - if sid in self._pending_messages: - pending_list = self._pending_messages.pop(sid) - for pending_msg in pending_list: - self._handle_complete_object(pending_msg, sid, messages) - - # Also try to yield reachable components if we just got a root - self.yield_reachable(messages) - return True - - # Fallback or shared logic (like deleteSurface) - if MSG_TYPE_DELETE_SURFACE in obj: - self._yield_messages([obj], messages) - return True - - # If message type is still unknown, we might have missed it - # We yield anything else as is - self._yield_messages([obj], messages) - return True + """Handles an object that has been fully parsed. To be implemented by subclasses.""" + raise NotImplementedError("Subclasses must implement _handle_complete_object") def yield_reachable( self, @@ -1017,10 +843,7 @@ def yield_reachable( return sid = self.surface_id - if ( - sid not in self._yielded_begin_rendering_surfaces - and not self._buffered_begin_rendering - ): + if sid not in self._yielded_surfaces_set and not self._buffered_start_message: return try: @@ -1083,68 +906,31 @@ def yield_reachable( should_yield = True break - # OR check if any path used by this component was updated in the most recent dataModelUpdate - # (Wait! yield_reachable doesn't know what was just updated. - # But we can check if any of its paths are in our current data model?) - # No, that's not enough. We need to know it was UPDATED. - # Actually, the test updates p1. - # I'll just assume if it has a path that is IN data_model, and it wasn't yielded yet? - # No, _yielded_contents already tracks the hash. - # If hash didn't change (because path is still /p1), then we need another signal. - # I'll add a 'dirty' set to the parser. - if should_yield: current_sid = self.surface_id or "unknown" if ( - self._buffered_begin_rendering - and current_sid not in self._yielded_real_begin_rendering_surfaces + self._buffered_start_message + and current_sid not in self._yielded_start_messages ): self._yield_messages( - [self._buffered_begin_rendering], messages, strict_integrity=True + [self._buffered_start_message], messages, strict_integrity=True ) - self._yielded_real_begin_rendering_surfaces.add(current_sid) - self._yielded_begin_rendering_surfaces.add(current_sid) + self._yielded_start_messages.add(current_sid) + self._yielded_surfaces_set.add(current_sid) # Construct a partial message of the correct type - if self._handler and self._handler.get_version() == VERSION_0_8: - # v0.8: Always yield components via surfaceUpdate - payload = { - SURFACE_ID_KEY: self._surface_id, - CATALOG_COMPONENTS_KEY: processed_components, - } - partial_msg = {MSG_TYPE_SURFACE_UPDATE: payload} - elif active_msg_type == MSG_TYPE_SURFACE_UPDATE: - payload = { - SURFACE_ID_KEY: self._surface_id, - CATALOG_COMPONENTS_KEY: processed_components, - } - partial_msg = {MSG_TYPE_SURFACE_UPDATE: payload} - elif active_msg_type in (MSG_TYPE_UPDATE_COMPONENTS, MSG_TYPE_CREATE_SURFACE): - # v0.9 logic - payload = { - CATALOG_COMPONENTS_KEY: processed_components, - "root": self._root_id or DEFAULT_ROOT_ID, - } - partial_msg = {MSG_TYPE_UPDATE_COMPONENTS: payload} - if self._surface_id: - partial_msg[SURFACE_ID_KEY] = self._surface_id - else: - # Fallback - partial_msg = { - MSG_TYPE_SURFACE_UPDATE: { - SURFACE_ID_KEY: self._surface_id, - CATALOG_COMPONENTS_KEY: processed_components, - } - } + partial_msg = self._construct_partial_message( + processed_components, active_msg_type + ) # Use strict_integrity=False for partial fragments yielded during streaming self._yield_messages([partial_msg], messages, strict_integrity=False) self._yielded_ids.setdefault(surface_id, set()).update(available_reachable) + # Update content/placeholder tracking for comp in processed_components: cid = comp["id"] self._yielded_contents[(surface_id, cid)] = json.dumps(comp, sort_keys=True) - self._yielded_placeholders[cid] = self._get_placeholders(comp) except ValueError as e: if "Circular reference detected" in str(e): @@ -1162,6 +948,10 @@ def yield_reachable( logger.debug(f"yield_reachable error (strict={check_root}): {msg}") raise e + def _get_placeholder_id(self, child_id: str) -> str: + """Returns the ID to use for a missing child placeholder.""" + return f"loading_{child_id}" + def _process_component_topology( self, comp: Dict[str, Any], @@ -1170,6 +960,7 @@ def _process_component_topology( ): """Recursively processes path placeholders and child pruning in one pass.""" comp_id = comp.get("id", "unknown") + # Deduce the component type for better placeholder typing comp_type = ( next(iter(comp.get("component", {}).keys())) @@ -1187,17 +978,9 @@ def traverse(obj, parent_key=None): ): path = obj["path"] key = path.lstrip("/") - # Always track dependencies for re-yielding when data arrives - self._comp_paths.setdefault(comp_id, set()).add(key) - - # Always keep path as-is, client will resolve it or leave undefined - # Always ensure path has a leading slash for v0.8 validation - obj.clear() + if "componentId" not in obj: + obj.clear() obj.update({"path": "/" + key}) - elif "dataBinding" in obj and isinstance(obj["dataBinding"], str): - path = obj["dataBinding"] - if path.startswith("/"): - self._comp_paths.setdefault(comp_id, set()).add(path.lstrip("/")) else: # If not in data model, still ensure path has leading slash if it's a bindable object current_path = obj.get("path") @@ -1222,11 +1005,11 @@ def traverse(obj, parent_key=None): valid_children.append(child_id) else: # Individual placeholder for missing child - placeholder_id = f"loading_{child_id}" + placeholder_id = self._get_placeholder_id(child_id) valid_children.append(placeholder_id) placeholder_comp = { "id": placeholder_id, - "component": PLACEHOLDER_COMPONENT, + **self._placeholder_component, } # Avoid duplicates in extra_components if not any(ec["id"] == placeholder_id for ec in extra_components): @@ -1244,7 +1027,7 @@ def traverse(obj, parent_key=None): valid_children.append(placeholder_id) placeholder_comp = { "id": placeholder_id, - "component": PLACEHOLDER_COMPONENT, + **self._placeholder_component, } if not any(ec["id"] == placeholder_id for ec in extra_components): extra_components.append(placeholder_comp) @@ -1252,11 +1035,11 @@ def traverse(obj, parent_key=None): elif isinstance(obj[field], str): child_id = obj[field] if child_id not in self._seen_components: - placeholder_id = f"loading_{child_id}" + placeholder_id = self._get_placeholder_id(child_id) obj[field] = placeholder_id placeholder_comp = { "id": placeholder_id, - "component": PLACEHOLDER_COMPONENT, + **self._placeholder_component, } if not any(ec["id"] == placeholder_id for ec in extra_components): extra_components.append(placeholder_comp) @@ -1269,22 +1052,8 @@ def traverse(obj, parent_key=None): traverse(item, parent_key) # Start recursion from the component content - traverse(comp.get("component", {})) - - def _get_placeholders(self, comp: Dict[str, Any]) -> Set[str]: - """Returns the set of placeholder paths currently in the component.""" - placeholders = set() - - def traverse(obj): - if isinstance(obj, dict): - if "literalString" in obj and isinstance(obj["literalString"], str): - if obj["literalString"].startswith("Loading from dataModel at path "): - placeholders.add(obj["literalString"]) - for v in obj.values(): - traverse(v) - elif isinstance(obj, list): - for item in obj: - traverse(item) - - traverse(comp.get("component", {})) - return placeholders + if isinstance(comp.get("component"), dict): + traverse(comp.get("component", {})) + else: + # Flat style properties are siblings to 'component' type key + traverse(comp) diff --git a/agent_sdks/python/src/a2ui/core/parser/streaming_v08.py b/agent_sdks/python/src/a2ui/core/parser/streaming_v08.py new file mode 100644 index 000000000..6678b77dc --- /dev/null +++ b/agent_sdks/python/src/a2ui/core/parser/streaming_v08.py @@ -0,0 +1,232 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import json +from typing import Any, List, Dict, Optional, Set + +from .streaming import A2uiStreamParser +from .response_part import ResponsePart +from .constants import * +from ..schema.constants import VERSION_0_8, SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY + + +class A2uiStreamParserV08(A2uiStreamParser): + """Streaming parser implementation for A2UI v0.8 specification.""" + + def __init__(self, catalog=None): + super().__init__(catalog=catalog) + self._root_id = None # v0.8 root is determined by beginRendering + self._yielded_begin_rendering_surfaces: Set[str] = set() + + @property + def _placeholder_component(self) -> Dict[str, Any]: + """Returns the placeholder component.""" + return { + 'component': { + 'Row': { + 'children': {'explicitList': []}, + } + } + } + + @property + def _yielded_surfaces_set(self) -> Set[str]: + """Provides access to version-specific yielded surfaces set.""" + return self._yielded_begin_rendering_surfaces + + def is_protocol_msg(self, obj: Dict[str, Any]) -> bool: + """Checks if the object is a recognized v0.8 message.""" + return any( + k in obj + for k in ( + MSG_TYPE_BEGIN_RENDERING, + MSG_TYPE_SURFACE_UPDATE, + MSG_TYPE_DATA_MODEL_UPDATE, + MSG_TYPE_DELETE_SURFACE, + ) + ) + + @property + 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): + """Sniffs for v0.8 metadata in the json_buffer.""" + if not self.surface_id: + match = re.search(r'"surfaceId"\s*:\s*"([^"]+)"', self._json_buffer) + if match: + self.surface_id = match.group(1) + + if not self.root_id: + match = re.search(r'"root"\s*:\s*"([^"]+)"', self._json_buffer) + if match: + self.root_id = match.group(1) + + if f'"{MSG_TYPE_BEGIN_RENDERING}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_BEGIN_RENDERING) + if f'"{MSG_TYPE_SURFACE_UPDATE}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_SURFACE_UPDATE) + if f'"{MSG_TYPE_DATA_MODEL_UPDATE}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_DATA_MODEL_UPDATE) + if f'"{MSG_TYPE_DELETE_SURFACE}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_DELETE_SURFACE) + + def _handle_complete_object( + self, + obj: Dict[str, Any], + sid: Optional[str], + messages: List[ResponsePart], + ) -> bool: + """Handles v0.8 specific complete objects.""" + if not isinstance(obj, dict): + return False + + if self._validator: + self._validator.validate(obj, root_id=sid, strict_integrity=False) + + # Update state based on the message content + surface_id = obj.get(SURFACE_ID_KEY, self.surface_id) + if MSG_TYPE_SURFACE_UPDATE in obj: + val = obj[MSG_TYPE_SURFACE_UPDATE] + if isinstance(val, dict): + surface_id = val.get(SURFACE_ID_KEY) or surface_id + elif MSG_TYPE_BEGIN_RENDERING in obj: + val = obj[MSG_TYPE_BEGIN_RENDERING] + if isinstance(val, dict): + surface_id = val.get(SURFACE_ID_KEY) or surface_id + elif MSG_TYPE_DELETE_SURFACE in obj: + val = obj[MSG_TYPE_DELETE_SURFACE] + if isinstance(val, str): + surface_id = val + elif isinstance(val, dict): + surface_id = val.get(SURFACE_ID_KEY) or surface_id + + self.surface_id = surface_id + sid = self.surface_id or 'unknown' + + if MSG_TYPE_DELETE_SURFACE in obj: + if sid in self._yielded_surfaces_set or self._buffered_start_message: + self._delete_surface(sid) + + if sid in self._deleted_surfaces: + return True + + if ( + (MSG_TYPE_SURFACE_UPDATE in obj or MSG_TYPE_DELETE_SURFACE in obj) + and sid not in self._yielded_surfaces_set + and not self._buffered_start_message + ): + if sid not in self._pending_messages: + self._pending_messages[sid] = [] + self._pending_messages[sid].append(obj) + return True + + if MSG_TYPE_BEGIN_RENDERING in obj: + br_val = obj[MSG_TYPE_BEGIN_RENDERING] + if isinstance(br_val, dict): + self.surface_id = br_val.get(SURFACE_ID_KEY, self.surface_id) + self.root_id = br_val.get('root', self.root_id or DEFAULT_ROOT_ID) + self._buffered_start_message = obj + + # Yield beginRendering immediately when it completes + if sid not in self._yielded_start_messages: + self._yield_messages([obj], messages) + self._yielded_start_messages.add(sid) + self._yielded_surfaces_set.add(sid) + self._buffered_start_message = None + + if sid in self._pending_messages: + pending_list = self._pending_messages.pop(sid) + for pending_msg in pending_list: + self._handle_complete_object(pending_msg, sid, messages) + + self.yield_reachable(messages) + return True + + if MSG_TYPE_SURFACE_UPDATE in obj: + self.add_msg_type(MSG_TYPE_SURFACE_UPDATE) + components = obj[MSG_TYPE_SURFACE_UPDATE].get('components', []) + for comp in components: + if isinstance(comp, dict) and 'id' in comp: + self._seen_components[comp['id']] = comp + self.yield_reachable(messages, check_root=True, raise_on_orphans=False) + return True + + if MSG_TYPE_DATA_MODEL_UPDATE in obj: + self.add_msg_type(MSG_TYPE_DATA_MODEL_UPDATE) + self.update_data_model(obj[MSG_TYPE_DATA_MODEL_UPDATE], messages) + self._yield_messages([obj], messages) + self.yield_reachable(messages, check_root=False, raise_on_orphans=False) + return True + + if MSG_TYPE_DELETE_SURFACE in obj: + self._yield_messages([obj], messages) + return True + + # If unknown, let base class yield it or yield it here + self._yield_messages([obj], messages) + return True + + def _construct_partial_message( + self, processed_components: List[Dict[str, Any]], active_msg_type: str + ) -> Dict[str, Any]: + """Constructs a partial message for v0.8 (always surfaceUpdate).""" + payload = { + SURFACE_ID_KEY: self.surface_id, + CATALOG_COMPONENTS_KEY: processed_components, + } + return {MSG_TYPE_SURFACE_UPDATE: payload} + + def _get_active_msg_type_for_components(self) -> Optional[str]: + """Determines which msg_type to use when wrapping component updates.""" + if self._active_msg_type: + return self._active_msg_type + for mt in self._msg_types: + if mt in (MSG_TYPE_SURFACE_UPDATE, MSG_TYPE_BEGIN_RENDERING): + self._active_msg_type = mt + return mt + return self._msg_types[0] if self._msg_types else None + + def _deduplicate_data_model(self, m: Dict[str, Any], strict_integrity: bool) -> bool: + if MSG_TYPE_DATA_MODEL_UPDATE in m: + dm = m[MSG_TYPE_DATA_MODEL_UPDATE] + raw_contents = dm.get('contents', {}) + contents_dict = {} + if isinstance(raw_contents, list): + for entry in raw_contents: + if isinstance(entry, dict) and 'key' in entry: + key = entry['key'] + val = ( + entry.get('valueString') + or entry.get('valueNumber') + or entry.get('valueBoolean') + or entry.get('valueMap') + ) + if key and val is not None: + contents_dict[key] = val + elif isinstance(raw_contents, dict): + contents_dict = raw_contents + + if contents_dict: + is_new = False + for k, v in contents_dict.items(): + if self._yielded_data_model.get(k) != v: + is_new = True + break + if not is_new and strict_integrity: + return False + self._yielded_data_model.update(contents_dict) + return True diff --git a/agent_sdks/python/src/a2ui/core/parser/streaming_v09.py b/agent_sdks/python/src/a2ui/core/parser/streaming_v09.py new file mode 100644 index 000000000..5284c0136 --- /dev/null +++ b/agent_sdks/python/src/a2ui/core/parser/streaming_v09.py @@ -0,0 +1,262 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import json +from typing import Any, List, Dict, Optional, Set + +from .streaming import A2uiStreamParser +from .response_part import ResponsePart +from .constants import * +from ..schema.constants import VERSION_0_9, SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY + + +class A2uiStreamParserV09(A2uiStreamParser): + """Streaming parser implementation for A2UI v0.9 specification.""" + + def __init__(self, catalog=None): + super().__init__(catalog=catalog) + # v0.9 default root is "root" + self.root_id = DEFAULT_ROOT_ID + + @property + def _placeholder_component(self) -> Dict[str, Any]: + """Returns a v0.9 flat style placeholder component specification.""" + return { + 'component': 'Row', + 'children': [], + } + + @property + def _data_model_msg_type(self) -> str: + """Returns the message type identifier for data model updates.""" + return MSG_TYPE_UPDATE_DATA_MODEL + + def is_protocol_msg(self, obj: Dict[str, Any]) -> bool: + """Checks if the object is a recognized v0.9 message.""" + return any( + k in obj + for k in ( + MSG_TYPE_CREATE_SURFACE, + MSG_TYPE_UPDATE_COMPONENTS, + MSG_TYPE_UPDATE_DATA_MODEL, + ) + ) + + def _sniff_metadata(self): + """Sniffs for v0.9 metadata in the json_buffer.""" + if not self.surface_id: + match = re.search(r'"surfaceId"\s*:\s*"([^"]+)"', self._json_buffer) + if match: + self.surface_id = match.group(1) + + if not self.root_id or self.root_id == DEFAULT_ROOT_ID: + match = re.search(r'"root"\s*:\s*"([^"]+)"', self._json_buffer) + if match: + self.root_id = match.group(1) + + if f'"{MSG_TYPE_CREATE_SURFACE}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_CREATE_SURFACE) + if f'"{MSG_TYPE_UPDATE_COMPONENTS}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) + if f'"{MSG_TYPE_UPDATE_DATA_MODEL}":' in self._json_buffer: + self.add_msg_type(MSG_TYPE_UPDATE_DATA_MODEL) + + def _handle_complete_object( + self, + obj: Dict[str, Any], + sid: Optional[str], + messages: List[ResponsePart], + ) -> bool: + """Handles v0.9 specific complete objects.""" + if not isinstance(obj, dict): + return False + + if self._validator: + self._validator.validate(obj, root_id=sid, strict_integrity=False) + + # Update state based on the message content + surface_id = obj.get(SURFACE_ID_KEY, self.surface_id) + if MSG_TYPE_UPDATE_COMPONENTS in obj: + val = obj[MSG_TYPE_UPDATE_COMPONENTS] + if isinstance(val, dict): + surface_id = val.get(SURFACE_ID_KEY) or surface_id + elif MSG_TYPE_CREATE_SURFACE in obj: + val = obj[MSG_TYPE_CREATE_SURFACE] + if isinstance(val, dict): + surface_id = val.get(SURFACE_ID_KEY) or surface_id + + self.surface_id = surface_id + sid = self.surface_id or 'unknown' + + # v0.9 Specific Handling + if MSG_TYPE_CREATE_SURFACE in obj: + val = obj[MSG_TYPE_CREATE_SURFACE] + if isinstance(val, dict): + self.root_id = val.get('root', self.root_id or DEFAULT_ROOT_ID) + self._buffered_start_message = obj + + # Yield createSurface immediately when it completes + if sid not in self._yielded_start_messages: + self._yield_messages([obj], messages) + self._yielded_start_messages.add(sid) + self._yielded_surfaces_set.add(sid) + self._buffered_start_message = None + + if sid in self._pending_messages: + # Clear pending messages when createSurface arrives, we want a fresh start! + self._pending_messages.pop(sid) + + self.yield_reachable(messages) + return True + + if MSG_TYPE_UPDATE_COMPONENTS in obj: + self.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) + self.root_id = obj[MSG_TYPE_UPDATE_COMPONENTS].get( + 'root', self.root_id or DEFAULT_ROOT_ID + ) + components = obj[MSG_TYPE_UPDATE_COMPONENTS].get('components', []) + for comp in components: + if isinstance(comp, dict) and 'id' in comp: + self._seen_components[comp['id']] = comp + self.yield_reachable(messages, check_root=True, raise_on_orphans=False) + return True + + if MSG_TYPE_DELETE_SURFACE in obj: + if sid not in self._yielded_start_messages: + self._pending_messages.setdefault(sid, []).append(obj) + return True + self.add_msg_type(MSG_TYPE_DELETE_SURFACE) + self._yield_messages([obj], messages) + return True + + if MSG_TYPE_UPDATE_DATA_MODEL in obj: + + self.add_msg_type(MSG_TYPE_UPDATE_DATA_MODEL) + self.update_data_model(obj[MSG_TYPE_UPDATE_DATA_MODEL], messages) + self._yield_messages([obj], messages) + return True + + return False + + def _construct_sniffed_data_model_message( + self, active_msg_type: str, delta_msg_payload: Dict[str, Any] + ) -> Dict[str, Any]: + """Returns the message to yield for a partial data model update for v0.9.""" + return {'version': 'v0.9', active_msg_type: delta_msg_payload} + + def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None: + """Sniffs for partial data model updates in v0.9 (value property).""" + msg_type = MSG_TYPE_UPDATE_DATA_MODEL + if f'"{msg_type}"' not in self._json_buffer: + return + + for b_type, start_idx in reversed(self._brace_stack): + if b_type != '{': + continue + raw_fragment = self._json_buffer[start_idx:] + if not raw_fragment: + continue + + fixed_fragment = self._fix_json(raw_fragment) + obj = None + try: + obj = json.loads(fixed_fragment) + except json.JSONDecodeError: + # Fallback: iteratively strip from the last comma + trimmed = raw_fragment + while ',' in trimmed: + trimmed = trimmed.rsplit(',', 1)[0] + try: + fixed_trimmed = self._fix_json(trimmed) + if fixed_trimmed: + obj = json.loads(fixed_trimmed) + break + except json.JSONDecodeError: + continue + + if obj and isinstance(obj, dict) and msg_type in obj: + + dm_obj = obj[msg_type] + if isinstance(dm_obj, dict) and 'value' in dm_obj: + value_map = dm_obj['value'] + if isinstance(value_map, dict): + # Find delta against yielded data model + delta = {} + for k, v in value_map.items(): + if self._yielded_data_model.get(k) != v: + delta[k] = v + + if delta: + sid = dm_obj.get(SURFACE_ID_KEY) or self._surface_id or 'default' + delta_msg_payload = { + SURFACE_ID_KEY: sid, + 'value': delta, + } + delta_msg = self._construct_sniffed_data_model_message( + msg_type, delta_msg_payload + ) + self._yield_messages([delta_msg], messages, strict_integrity=False) + # Do NOT update _yielded_data_model here, let update_data_model do it when complete + # Wait! If we don't update it, will we over-yield it in the next chunk? + # Yes, we might. So we should update it or track it! + # The base class updates it (line 644 approx). So we should update it too! + self._yielded_data_model.update(delta) + + 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 = { + CATALOG_COMPONENTS_KEY: processed_components, + } + if self.surface_id: + payload[SURFACE_ID_KEY] = self.surface_id + if self.root_id: + payload['root'] = self.root_id + return {'version': 'v0.9', MSG_TYPE_UPDATE_COMPONENTS: payload} + + @property + def _yielded_surfaces_set(self) -> Set[str]: + """Provides access to version-specific yielded surfaces set.""" + if not hasattr(self, '_yielded_create_surfaces'): + self._yielded_create_surfaces: Set[str] = set() + return self._yielded_create_surfaces + + def _get_active_msg_type_for_components(self) -> Optional[str]: + """Determines which msg_type to use when wrapping component updates.""" + if self._active_msg_type: + return self._active_msg_type + for mt in self._msg_types: + if mt in (MSG_TYPE_UPDATE_COMPONENTS, MSG_TYPE_CREATE_SURFACE): + self._active_msg_type = mt + return mt + return self._msg_types[0] if self._msg_types else None + + def _deduplicate_data_model(self, m: Dict[str, Any], strict_integrity: bool) -> bool: + if MSG_TYPE_UPDATE_DATA_MODEL in m: + udm = m[MSG_TYPE_UPDATE_DATA_MODEL] + if isinstance(udm, dict): + is_new = False + for k, v in udm.items(): + if k not in (SURFACE_ID_KEY, 'root') and self._yielded_data_model.get(k) != v: + is_new = True + break + if not is_new and strict_integrity: + return False + # Update yielded model + for k, v in udm.items(): + if k not in (SURFACE_ID_KEY, 'root'): + self._yielded_data_model[k] = v + return True diff --git a/agent_sdks/python/src/a2ui/core/parser/version_handlers.py b/agent_sdks/python/src/a2ui/core/parser/version_handlers.py deleted file mode 100644 index cb6939b8d..000000000 --- a/agent_sdks/python/src/a2ui/core/parser/version_handlers.py +++ /dev/null @@ -1,266 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import re -import json -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, TYPE_CHECKING -from ..schema.constants import ( - VERSION_0_9, - VERSION_0_8, - SURFACE_ID_KEY, -) -from .response_part import ResponsePart -from .constants import ( - DEFAULT_ROOT_ID, - MSG_TYPE_BEGIN_RENDERING, - MSG_TYPE_SURFACE_UPDATE, - MSG_TYPE_DATA_MODEL_UPDATE, - MSG_TYPE_DELETE_SURFACE, - MSG_TYPE_CREATE_SURFACE, - MSG_TYPE_UPDATE_COMPONENTS, - MSG_TYPE_UPDATE_DATA_MODEL, -) - -if TYPE_CHECKING: - from .streaming import A2uiStreamParser - - -class A2uiVersionHandler(ABC): - """Base class for version-specific A2UI message handling. - - This class decouples A2UI protocol version differences from the core streaming - parser logic. It encapsulates logic for version-specific message identification, - metadata extraction, and state transitions. - - Subclasses should be created for each supported A2UI version (e.g., v0.8, v0.9). - This allows the `A2uiStreamParser` to remain version-agnostic and easily - extensible for future protocol updates. - """ - - @staticmethod - def detect_version(json_buffer: str) -> Optional[str]: - """Detects the A2UI version from the JSON buffer. - - This method is used by the `A2uiStreamParser` at the start of a stream or - when a new A2UI message block is encountered to determine which version - handler to use. It performs lightweight pattern matching on the buffered - JSON text. - - Args: - json_buffer: The raw JSON string buffered so far. - - Returns: - The version string (e.g., "0.8", "0.9") if detected, None otherwise. - """ - # v0.9 markers: "version": "v0.9" or specific message types - if ( - re.search(rf'"version"\s*:\s*"v{VERSION_0_9}"', json_buffer, re.I) - or f'"{MSG_TYPE_UPDATE_COMPONENTS}"' in json_buffer - or f'"{MSG_TYPE_CREATE_SURFACE}"' in json_buffer - ): - return VERSION_0_9 - # v0.8 markers - if ( - f'"{MSG_TYPE_BEGIN_RENDERING}"' in json_buffer - or f'"{MSG_TYPE_SURFACE_UPDATE}"' in json_buffer - ): - return VERSION_0_8 - return None - - @abstractmethod - def sniff_metadata(self, json_buffer: str, parser: 'A2uiStreamParser'): - """Sniffs for surfaceId, root, and msg_types in the current json_buffer. - - This method allows for incremental metadata discovery during streaming. - By looking at the raw JSON buffer fragment, the handler can populate - the parser's state even before an object is fully closed. This is crucial - for "fine-grained" streaming where we want to know the context (like - surfaceId or the root component) as early as possible. - - Should be called as new characters are added to the parser's top-level - JSON buffer. - - Args: - json_buffer: The raw JSON string currently being sniffed. - parser: The A2uiStreamParser instance whose state should be updated. - """ - pass - - @abstractmethod - def handle_complete_object( - self, - obj: Dict[str, Any], - parser: 'A2uiStreamParser', - messages: List[ResponsePart], - ) -> bool: - """Handles a completed object from the top-level list. - - This is called when the parser identifies a complete JSON object at the - top level of the A2UI message array. The handler should check if the - object is a known message type for its version and update the parser's - state or yield final/partial messages accordingly. - - Args: - obj: The fully parsed JSON object (dictionary). - parser: The A2uiStreamParser instance. - messages: The list of accumulated A2UI messages to append to. - - Returns: - True if the object was fully handled by this version handler, False - if the parser should fall back to shared logic or append it as-is. - """ - pass - - @abstractmethod - def is_v08_msg(self, obj: Dict[str, Any]) -> bool: - """Checks if the object is a recognized v0.8 message.""" - pass - - @abstractmethod - def get_version(self) -> str: - """Returns the version string.""" - pass - - -class A2uiV08Handler(A2uiVersionHandler): - """Handler for A2UI v0.8 messages.""" - - def get_version(self) -> str: - return VERSION_0_8 - - def is_v08_msg(self, obj: Dict[str, Any]) -> bool: - return any( - k in obj - for k in ( - MSG_TYPE_BEGIN_RENDERING, - MSG_TYPE_SURFACE_UPDATE, - MSG_TYPE_DATA_MODEL_UPDATE, - MSG_TYPE_DELETE_SURFACE, - ) - ) - - def sniff_metadata(self, json_buffer: str, parser: 'A2uiStreamParser'): - if not parser.surface_id: - match = re.search(r'"surfaceId"\s*:\s*"([^"]+)"', json_buffer) - if match: - parser.surface_id = match.group(1) - - if not parser.root_id: - match = re.search(r'"root"\s*:\s*"([^"]+)"', json_buffer) - if match: - parser.root_id = match.group(1) - - if f'"{MSG_TYPE_BEGIN_RENDERING}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_BEGIN_RENDERING) - if f'"{MSG_TYPE_SURFACE_UPDATE}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_SURFACE_UPDATE) - if f'"{MSG_TYPE_DATA_MODEL_UPDATE}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_DATA_MODEL_UPDATE) - if f'"{MSG_TYPE_DELETE_SURFACE}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_DELETE_SURFACE) - - def handle_complete_object( - self, - obj: Dict[str, Any], - parser: 'A2uiStreamParser', - messages: List[ResponsePart], - ) -> bool: - if MSG_TYPE_BEGIN_RENDERING in obj: - br_val = obj[MSG_TYPE_BEGIN_RENDERING] - if isinstance(br_val, dict): - parser.surface_id = br_val.get(SURFACE_ID_KEY, parser.surface_id) - parser.root_id = br_val.get('root', parser.root_id or DEFAULT_ROOT_ID) - parser.buffered_begin_rendering = obj - return True - - if MSG_TYPE_SURFACE_UPDATE in obj: - parser.add_msg_type(MSG_TYPE_SURFACE_UPDATE) - components = obj[MSG_TYPE_SURFACE_UPDATE].get('components', []) - for comp in components: - if isinstance(comp, dict) and 'id' in comp: - parser.seen_components[comp['id']] = comp - parser.yield_reachable(messages, check_root=True, raise_on_orphans=False) - return True - - if MSG_TYPE_DATA_MODEL_UPDATE in obj: - parser.add_msg_type(MSG_TYPE_DATA_MODEL_UPDATE) - parser.update_data_model(obj[MSG_TYPE_DATA_MODEL_UPDATE], messages) - parser._yield_messages([obj], messages) - parser.yield_reachable(messages, check_root=False, raise_on_orphans=False) - return True - - return False - - -class A2uiV09Handler(A2uiVersionHandler): - """Handler for A2UI v0.9 messages.""" - - def get_version(self) -> str: - return VERSION_0_9 - - def is_v08_msg(self, obj: Dict[str, Any]) -> bool: - return False - - def sniff_metadata(self, json_buffer: str, parser: 'A2uiStreamParser'): - if not parser.surface_id: - match = re.search(r'"surfaceId"\s*:\s*"([^"]+)"', json_buffer) - if match: - parser.surface_id = match.group(1) - - if not parser.root_id: - # v0.9 default root is "root", but it can be overridden - match = re.search(r'"root"\s*:\s*"([^"]+)"', json_buffer) - if match: - parser.root_id = match.group(1) - - if f'"{MSG_TYPE_CREATE_SURFACE}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_CREATE_SURFACE) - if f'"{MSG_TYPE_UPDATE_COMPONENTS}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) - if f'"{MSG_TYPE_UPDATE_DATA_MODEL}":' in json_buffer: - parser.add_msg_type(MSG_TYPE_UPDATE_DATA_MODEL) - - def handle_complete_object( - self, - obj: Dict[str, Any], - parser: 'A2uiStreamParser', - messages: List[ResponsePart], - ) -> bool: - if MSG_TYPE_CREATE_SURFACE in obj: - # createSurface in v0.9 is similar to beginRendering but maybe doesn't establish root yet - # Actually, v0.9 says root is usually "root". - parser.root_id = parser.root_id or DEFAULT_ROOT_ID - parser.buffered_begin_rendering = obj - return True - - if MSG_TYPE_UPDATE_COMPONENTS in obj: - parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) - parser.root_id = obj[MSG_TYPE_UPDATE_COMPONENTS].get( - 'root', parser.root_id or DEFAULT_ROOT_ID - ) - components = obj[MSG_TYPE_UPDATE_COMPONENTS].get('components', []) - for comp in components: - if isinstance(comp, dict) and 'id' in comp: - parser.seen_components[comp['id']] = comp - parser.yield_reachable(messages, check_root=True, raise_on_orphans=False) - return True - - if MSG_TYPE_UPDATE_DATA_MODEL in obj: - parser.add_msg_type(MSG_TYPE_UPDATE_DATA_MODEL) - parser.update_data_model(obj[MSG_TYPE_UPDATE_DATA_MODEL], messages) - parser._yield_messages([obj], messages) - return True - - return False diff --git a/agent_sdks/python/tests/core/parser/test_streaming_v08.py b/agent_sdks/python/tests/core/parser/test_streaming_v08.py index 3cef6bc47..51cfed5bc 100644 --- a/agent_sdks/python/tests/core/parser/test_streaming_v08.py +++ b/agent_sdks/python/tests/core/parser/test_streaming_v08.py @@ -30,7 +30,7 @@ MSG_TYPE_DATA_MODEL_UPDATE, ) from a2ui.core.schema.catalog import A2uiCatalog -from a2ui.core.parser.streaming import A2uiStreamParser, PLACEHOLDER_COMPONENT +from a2ui.core.parser.streaming import A2uiStreamParser from a2ui.core.parser.response_part import ResponsePart @@ -223,10 +223,6 @@ def _normalize_messages(messages): payload = msg[MSG_TYPE_SURFACE_UPDATE] if CATALOG_COMPONENTS_KEY in payload: payload[CATALOG_COMPONENTS_KEY].sort(key=lambda x: x.get("id", "")) - elif "updateComponents" in msg: - payload = msg["updateComponents"] - if CATALOG_COMPONENTS_KEY in payload: - payload[CATALOG_COMPONENTS_KEY].sort(key=lambda x: x.get("id", "")) return res @@ -335,7 +331,7 @@ def test_incremental_yielding_v08(mock_catalog): }, { "id": "loading_children_root-column", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -356,7 +352,7 @@ def test_incremental_yielding_v08(mock_catalog): }, { "id": "loading_c1", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -379,11 +375,11 @@ def test_incremental_yielding_v08(mock_catalog): }, { "id": "loading_c1", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, { "id": "loading_c2", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -422,7 +418,7 @@ def test_incremental_yielding_v08(mock_catalog): }, { "id": "loading_c2", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -802,7 +798,7 @@ def test_partial_single_child_string(mock_catalog): }, { "id": "loading_c1", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -870,7 +866,7 @@ def test_partial_template_componentId(mock_catalog): }, { "id": "loading_c1", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -938,15 +934,15 @@ def test_partial_children_lists(mock_catalog): }, { "id": "loading_c1", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, { "id": "loading_c2", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, { "id": "loading_c3", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -979,11 +975,11 @@ def test_partial_children_lists(mock_catalog): }, { "id": "loading_c2", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, { "id": "loading_c3", - "component": PLACEHOLDER_COMPONENT, + **parser._placeholder_component, }, ], } @@ -1881,8 +1877,8 @@ def test_sniff_partial_component_discards_empty_children_dict(mock_catalog): "id": "root-column", }, { - "component": PLACEHOLDER_COMPONENT, "id": "loading_item-list", + **parser._placeholder_component, }, ], } diff --git a/agent_sdks/python/tests/core/parser/test_streaming_v09.py b/agent_sdks/python/tests/core/parser/test_streaming_v09.py new file mode 100644 index 000000000..6255467e3 --- /dev/null +++ b/agent_sdks/python/tests/core/parser/test_streaming_v09.py @@ -0,0 +1,2034 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import copy +from unittest.mock import MagicMock +import pytest +from a2ui.core.schema.constants import ( + A2UI_OPEN_TAG, + A2UI_CLOSE_TAG, + VERSION_0_9, + SURFACE_ID_KEY, + CATALOG_COMPONENTS_KEY, +) +from a2ui.core.parser.constants import ( + MSG_TYPE_CREATE_SURFACE, + MSG_TYPE_UPDATE_COMPONENTS, + MSG_TYPE_DELETE_SURFACE, + MSG_TYPE_DATA_MODEL_UPDATE, +) +from a2ui.core.schema.catalog import A2uiCatalog +from a2ui.core.parser.streaming import A2uiStreamParser +from a2ui.core.parser.response_part import ResponsePart + + +@pytest.fixture +def mock_catalog(): + s2c_schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://a2ui.org/specification/v0_9/server_to_client.json", + "title": "A2UI Message Schema", + "type": "object", + "oneOf": [ + {"$ref": "#/$defs/CreateSurfaceMessage"}, + {"$ref": "#/$defs/UpdateComponentsMessage"}, + {"$ref": "#/$defs/UpdateDataModelMessage"}, + {"$ref": "#/$defs/DeleteSurfaceMessage"}, + ], + "$defs": { + "CreateSurfaceMessage": { + "type": "object", + "properties": { + "version": {"const": "v0.9"}, + "createSurface": { + "type": "object", + "properties": { + "surfaceId": { + "type": "string", + }, + "catalogId": { + "type": "string", + }, + "theme": { + "type": "object", + "additionalProperties": True, + }, + }, + "required": ["surfaceId", "catalogId"], + "additionalProperties": False, + }, + }, + "required": ["version", "createSurface"], + "additionalProperties": False, + }, + "UpdateComponentsMessage": { + "type": "object", + "properties": { + "version": {"const": "v0.9"}, + "updateComponents": { + "type": "object", + "properties": { + "surfaceId": { + "type": "string", + }, + "root": { + "type": "string", + }, + "components": { + "type": "array", + "minItems": 1, + "items": {"$ref": "catalog.json#/$defs/anyComponent"}, + }, + }, + "required": ["surfaceId", "components"], + "additionalProperties": False, + }, + }, + "required": ["version", "updateComponents"], + "additionalProperties": False, + }, + "UpdateDataModelMessage": { + "type": "object", + "properties": { + "version": {"const": "v0.9"}, + "updateDataModel": { + "type": "object", + "properties": { + "surfaceId": { + "type": "string", + }, + "value": {"additionalProperties": True}, + }, + "required": ["surfaceId"], + "additionalProperties": False, + }, + }, + "required": ["version", "updateDataModel"], + "additionalProperties": False, + }, + "DeleteSurfaceMessage": { + "type": "object", + "properties": { + "version": {"const": "v0.9"}, + "deleteSurface": { + "type": "object", + "properties": {"surfaceId": {"type": "string"}}, + "required": ["surfaceId"], + }, + }, + "required": ["version", "deleteSurface"], + }, + }, + } + catalog_schema = { + "catalogId": "test_catalog", + "components": { + "Container": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + ], + "properties": { + "component": {"const": "Container"}, + "children": { + "type": "array", + "items": {"type": "string"}, + }, + }, + "required": ["component", "children"], + }, + "Card": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + ], + "properties": { + "component": {"const": "Card"}, + "child": {"$ref": "common_types.json#/$defs/ComponentId"}, + }, + "required": ["component", "child"], + }, + "Text": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + ], + "properties": { + "component": {"const": "Text"}, + "text": {"$ref": "common_types.json#/$defs/DynamicString"}, + }, + "required": ["component", "text"], + }, + "Column": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + ], + "properties": { + "component": {"const": "Column"}, + "children": {"$ref": "common_types.json#/$defs/ChildList"}, + }, + "required": ["component", "children"], + }, + "AudioPlayer": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + { + "type": "object", + "properties": { + "component": {"const": "AudioPlayer"}, + "url": {"$ref": "common_types.json#/$defs/DynamicString"}, + "description": { + "$ref": "common_types.json#/$defs/DynamicString" + }, + }, + "required": ["component", "url"], + }, + ], + }, + "List": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + { + "type": "object", + "properties": { + "component": {"const": "List"}, + "children": {"$ref": "common_types.json#/$defs/ChildList"}, + "direction": { + "type": "string", + "enum": ["vertical", "horizontal"], + }, + }, + "required": ["component", "children"], + }, + ], + }, + "Row": { + "type": "object", + "allOf": [ + {"$ref": "common_types.json#/$defs/ComponentCommon"}, + {"$ref": "#/$defs/CatalogComponentCommon"}, + { + "type": "object", + "properties": { + "component": {"const": "Row"}, + "children": {"$ref": "common_types.json#/$defs/ChildList"}, + }, + "required": ["component", "children"], + }, + ], + }, + }, + "$defs": { + "CatalogComponentCommon": { + "type": "object", + "properties": {"weight": {"type": "number"}}, + }, + "anyComponent": { + "oneOf": [ + {"$ref": "#/components/Container"}, + {"$ref": "#/components/Card"}, + {"$ref": "#/components/Text"}, + {"$ref": "#/components/Column"}, + {"$ref": "#/components/AudioPlayer"}, + {"$ref": "#/components/List"}, + {"$ref": "#/components/Row"}, + ], + "discriminator": {"propertyName": "component"}, + }, + }, + } + common_types_schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://a2ui.org/specification/v0_9/common_types.json", + "title": "A2UI Common Types", + "$defs": { + "ComponentId": { + "type": "string", + }, + "AccessibilityAttributes": { + "type": "object", + "properties": { + "label": { + "$ref": "#/$defs/DynamicString", + } + }, + }, + "Action": {"type": "object", "additionalProperties": True}, + "ComponentCommon": { + "type": "object", + "properties": {"id": {"$ref": "#/$defs/ComponentId"}}, + "required": ["id"], + }, + "DataBinding": {"type": "object"}, + "DynamicString": { + "anyOf": [{"type": "string"}, {"$ref": "#/$defs/DataBinding"}] + }, + "DynamicValue": { + "anyOf": [ + {"type": "object"}, + {"type": "array"}, + {"$ref": "#/$defs/DataBinding"}, + ] + }, + "DynamicNumber": { + "anyOf": [{"type": "number"}, {"$ref": "#/$defs/DataBinding"}] + }, + "ChildList": { + "oneOf": [ + {"type": "array", "items": {"$ref": "#/$defs/ComponentId"}}, + { + "type": "object", + "properties": { + "componentId": {"$ref": "#/$defs/ComponentId"}, + "path": {"type": "string"}, + }, + "required": ["componentId", "path"], + "additionalProperties": False, + }, + ] + }, + }, + } + return A2uiCatalog( + version=VERSION_0_9, + name="test_catalog", + s2c_schema=s2c_schema, + common_types_schema=common_types_schema, + catalog_schema=catalog_schema, + ) + + +def _normalize_messages(messages): + """Sorts components in messages for stable comparison.""" + # Support ResponsePart list by extracting a2ui_json + res = [] + for m in messages: + if isinstance(m, ResponsePart): + if m.a2ui_json: + if isinstance(m.a2ui_json, list): + res.extend(copy.deepcopy(m.a2ui_json)) + else: + res.append(copy.deepcopy(m.a2ui_json)) + else: + res.append(copy.deepcopy(m)) + + for msg in res: + if MSG_TYPE_UPDATE_COMPONENTS in msg: + payload = msg[MSG_TYPE_UPDATE_COMPONENTS] + if CATALOG_COMPONENTS_KEY in payload: + payload[CATALOG_COMPONENTS_KEY].sort(key=lambda x: x.get("id", "")) + return res + + +def assertResponseContainsMessages(response, expected_messages): + """Asserts that the response parts contain the expected messages.""" + assert _normalize_messages(response) == _normalize_messages(expected_messages) + + +def assertResponseContainsNoA2UI(response): + assert len(response) == 0 or response[0].a2ui_json == None + + +def assertResponseContainsText(response, expected_text): + """Asserts that the response parts contain the expected text.""" + assert any( + (p.text if isinstance(p, ResponsePart) else p) == expected_text for p in response + ) + + +def test_incremental_yielding_v09(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # 1. CreateSurface initializes the session + chunk1 = "Here is your" + chunk2 = f" response.{A2UI_OPEN_TAG}[" + chunk3 = '{"version": "v0.9", "createSurface": {"surfaceId": "s1", "catalogId":' + chunk4 = ' "test_catalog' + chunk5 = '"' + chunk6 = "}}," + + response_parts = parser.process_chunk(chunk1) + assertResponseContainsText(response_parts, "Here is your") + assertResponseContainsNoA2UI(response_parts) + + response_parts = parser.process_chunk(chunk2) + assertResponseContainsText(response_parts, " response.") + assertResponseContainsNoA2UI(response_parts) + + response_parts = parser.process_chunk(chunk3) + assertResponseContainsNoA2UI(response_parts) + + response_parts = parser.process_chunk(chunk4) + assertResponseContainsNoA2UI(response_parts) + + # CreateSurface should not be yielded yet because createSurface is not closing. + response_parts = parser.process_chunk(chunk5) + assertResponseContainsNoA2UI(response_parts) + + response_parts = parser.process_chunk(chunk6) + expected = [{ + "version": "v0.9", + "createSurface": {"surfaceId": "s1", "catalogId": "test_catalog"}, + }] + assertResponseContainsMessages(response_parts, expected) + + # 2. Add root component via updateComponents + surface_chunk_1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [' + ) + surface_chunk_2 = '{"id": "root", "component": "Column", "children": [' + surface_chunk_3 = '"c1",' + surface_chunk_4 = '"c2"]}]}}' + + response_parts = parser.process_chunk(surface_chunk_1) + assertResponseContainsNoA2UI(response_parts) + # Root is seen but incomplete + response_parts = parser.process_chunk(surface_chunk_2) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Column", + "children": ["loading_children_root"], + }, + { + "id": "loading_children_root", + **parser._placeholder_component, + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + # Child c1 loading + response_parts = parser.process_chunk(surface_chunk_3) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Column", + "children": ["loading_c1"], + }, + { + "id": "loading_c1", + **parser._placeholder_component, + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + # Child c2 loading + response_parts = parser.process_chunk(surface_chunk_4) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Column", + "children": ["loading_c1", "loading_c2"], + }, + { + "id": "loading_c1", + **parser._placeholder_component, + }, + { + "id": "loading_c2", + **parser._placeholder_component, + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + # 3. Add child components + c1_chunk = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "components": [' + '{"id": "c1", "component": "Text", "text": "hello"}]}}' + ) + c2_chunk = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "components": [' + '{"id": "c2", "component": "Text", "text": "world"}]}}' + ) + + response_parts = parser.process_chunk(c1_chunk) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Column", + "children": ["c1", "loading_c2"], + }, + { + "id": "c1", + "component": "Text", + "text": "hello", + }, + { + "id": "loading_c2", + **parser._placeholder_component, + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + response_parts = parser.process_chunk(c2_chunk) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Column", + "children": ["c1", "c2"], + }, + { + "id": "c1", + "component": "Text", + "text": "hello", + }, + { + "id": "c2", + "component": "Text", + "text": "world", + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_waiting_for_root_component(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # Establish root + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # Send a non-root component but keep the updateComponents message open + chunk = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "c1", "component": "Text", "text": "hello"}' + ) + response_parts = parser.process_chunk(chunk) + # Should not yield anything because no reachable components from root yet + assertResponseContainsNoA2UI(response_parts) + + # Now send the root and close the updateComponents message + chunk_root = ( + f', {{"id": "root", "component": "Card", "child": "c1"}}]}}}} {A2UI_CLOSE_TAG}' + ) + response_parts = parser.process_chunk(chunk_root) + # Should yield createSurface and the updateComponents with both components + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "c1", + "component": "Text", + "text": "hello", + }, + { + "id": "root", + "component": "Card", + "child": "c1", + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_complete_surface_ignore_orphan_component(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + chunk = ( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}}, ' + ) + parser.process_chunk(chunk) + + chunk += ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": "root"}, {"id":' + ' "orphan", "component": "Text", "text": "orphan"}]}}]' + ) + chunk += A2UI_CLOSE_TAG + + response_parts = parser.process_chunk(chunk) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Text", + "text": "root", + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_circular_reference_detection(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # c1 -> c2 + parser.process_chunk( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "c1",' + ' "components": [{"id": "c1", "component": "Card", "child": "c2"}]}},' + ) + + # c2 -> c1 (Cycle!) + with pytest.raises(ValueError, match="Circular reference detected"): + parser.process_chunk( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "components":' + ' [{"id": "c2", "component": "Card", "child": "c1"}]}}' + ) + + +def test_self_reference_detection(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # c1 -> c1 (Self-reference!) + with pytest.raises(ValueError, match="Self-reference detected"): + parser.process_chunk( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "c1",' + ' "components": [{"id": "c1", "component": "Card", "child": "c1"}]}}' + ) + + +def test_interleaved_conversational_text(): + # No catalog needed for basic text extraction and unvalidated JSON parsing + parser = A2uiStreamParser(catalog=None) + + # Chunk 1: purely conversational + messages = parser.process_chunk("Hello! ") + assert messages == [ResponsePart(text="Hello! ", a2ui_json=None)] + + # Chunk 2: start of A2UI block + messages = parser.process_chunk(f"Here is your UI: {A2UI_OPEN_TAG}") + assert messages == [ResponsePart(text="Here is your UI: ", a2ui_json=None)] + + # Chunk 3: A2UI content + messages = parser.process_chunk( + '[{"createSurface": {"root": "root", "surfaceId": "s1"}}]' + ) + assert any( + m.a2ui_json + and ( + any("createSurface" in msg for msg in m.a2ui_json) + if isinstance(m.a2ui_json, list) + else "createSurface" in m.a2ui_json + ) + for m in messages + ) + + # Chunk 4: Closing A2UI and more text + messages = parser.process_chunk(f"{A2UI_CLOSE_TAG} That's all!") + assert messages == [ResponsePart(text=" That's all!", a2ui_json=None)] + + +def test_split_tag_handling_for_text(): + parser = A2uiStreamParser(catalog=None) + + # Split A2UI_OPEN_TAG: "" + tag_part1 = A2UI_OPEN_TAG[:4] + tag_part2 = A2UI_OPEN_TAG[4:] + + # Chunk 1: text followed by split tag + messages = parser.process_chunk(f"Talking {tag_part1}") + # "Talking " should be yielded because it's definitively not part of the tag prefix + assert messages == [ResponsePart(text="Talking ", a2ui_json=None)] + + # Chunk 2: completes the tag + messages = parser.process_chunk(tag_part2) + # Should transition to A2UI mode, no new text yielded + assert len(messages) == 0 + + # Chunk 3: A2UI content + close tag + text + messages = parser.process_chunk( + f'[{{"createSurface": {{"root": "r", "surfaceId": "s"}}}}] {A2UI_CLOSE_TAG} End.' + ) + texts = [m.text for m in messages if m.text] + assert " End." in texts + + +def test_create_surface_missing_catalog_id(mock_catalog): + """Verifies that v0.9 createSurface fails when catalogId is missing.""" + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(A2UI_OPEN_TAG) + + # Missing catalogId + chunk = '[{"version": "v0.9", "createSurface": {"surfaceId": "s1"}}]' + with pytest.raises(ValueError, match="required property"): + parser.process_chunk(chunk) + + +def test_only_create_surface(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + response_parts = parser.process_chunk(A2UI_OPEN_TAG) + assert response_parts == [] + + response_parts = parser.process_chunk("[") + assert response_parts == [] + + chunk = ( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1", "theme": {"primaryColor": "#FF0000",' + ) + response_parts = parser.process_chunk(chunk) + # createSurface is not yielded until the entire message is received + assert response_parts == [] + + response_parts = parser.process_chunk('"font": "Roboto"}') # closing styles + assert response_parts == [] + + response_parts = parser.process_chunk("}") # closing createSurface + assert response_parts == [] + + response_parts = parser.process_chunk("}") # closing the item in the list + expected = [ + { + "version": "v0.9", + "createSurface": { + "catalogId": "test_catalog", + "surfaceId": "s1", + "theme": {"primaryColor": "#FF0000", "font": "Roboto"}, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_add_msg_type_deduplication(): + parser = A2uiStreamParser() + parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) + parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) + assert parser.msg_types == [MSG_TYPE_UPDATE_COMPONENTS] + + parser.add_msg_type(MSG_TYPE_CREATE_SURFACE) + assert parser.msg_types == [MSG_TYPE_UPDATE_COMPONENTS, MSG_TYPE_CREATE_SURFACE] + parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) + assert parser.msg_types == [MSG_TYPE_UPDATE_COMPONENTS, MSG_TYPE_CREATE_SURFACE] + + +def test_streaming_msg_type_deduplication(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + # 1. Send partial chunk that triggers sniffing + chunk1 = ( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": "Hello"}' + ) + parser.process_chunk(chunk1) + + assert MSG_TYPE_UPDATE_COMPONENTS in parser.msg_types + assert parser.msg_types.count(MSG_TYPE_UPDATE_COMPONENTS) == 1 + + # 2. Send the rest, which triggers handle_complete_object + chunk2 = f', {{"id": "c1", "component": "Text", "text": "hi"}}]}}}} {A2UI_CLOSE_TAG}' + parser.process_chunk(chunk2) + + # After completion, msg_types is reset + assert not parser.msg_types + + +def test_multiple_a2ui_blocks(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # Block 1: createSurface + chunk1 = ( + f'Some text here {A2UI_OPEN_TAG}[{{"version": "v0.9", "createSurface":' + f' {{"catalogId": "test_catalog", "surfaceId": "s1"}}}}] {A2UI_CLOSE_TAG} mid' + " text" + ) + response_parts = parser.process_chunk(chunk1) + assert len(response_parts) == 2 + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "s1"}, + }, + ] + assertResponseContainsMessages(response_parts, expected) + assertResponseContainsText(response_parts, "Some text here ") + assertResponseContainsText(response_parts, " mid text") + + # Block 2: updateComponents + chunk2 = ( + f' more text {A2UI_OPEN_TAG}[{{"version": "v0.9", "updateComponents":' + ' {"surfaceId": "s1", "root": "root", "components": [{"id": "root",' + f' "component": "Text", "text": "block2"}}]}}}}]}}] {A2UI_CLOSE_TAG} trailing' + " text" + ) + response_parts = parser.process_chunk(chunk2) + assert len(response_parts) == 2 + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": "block2", + }], + }, + }] + assertResponseContainsMessages(response_parts, expected) + assertResponseContainsText(response_parts, " more text ") + assertResponseContainsText(response_parts, " trailing text") + + +def test_partial_json_and_incremental_yielding(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # 1. Open A2UI block + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # 2. Send a partial component in a updateComponents + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": "hello' + ) + # The JSON fixer should close the quotes, braces, and brackets + response_parts = parser.process_chunk(chunk1) + + # Should yield the partial root component + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": "hello", + }], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_partial_single_child_string(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # Establish root + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # Send a container with a single string child, but we haven't seen the child yet (e.g. Card with child) + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Card", "child": "c1"}' + ) + response_parts = parser.process_chunk(chunk1) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "loading_c1", + **parser._placeholder_component, + }, + { + "id": "root", + "component": "Card", + "child": "loading_c1", + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + # Send the child component + chunk2 = ', {"id": "c1", "component": "Text", "text": "child 1"}]}}' + response_parts = parser.process_chunk(chunk2) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Card", + "child": "c1", + }, + { + "id": "c1", + "component": "Text", + "text": "child 1", + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_partial_template_componentId(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # Establish root + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # Send a container with a template, but we haven't seen the template component yet + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "List", "children": {"componentId":' + ' "c1"' + ) + response_parts = parser.process_chunk(chunk1) + # No complete object is yielded due to missing path + expected = [] + assertResponseContainsMessages(response_parts, expected) + + # Send the path to complete the template, and the template component + chunk2 = ', "path": "/items"' + response_parts = parser.process_chunk(chunk2) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "List", + "children": {"componentId": "loading_c1", "path": "/items"}, + }, + { + "id": "loading_c1", + **parser._placeholder_component, + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + chunk3 = ( + '}}, {"id": "c1", "component": "Text", "text": "child 1"}]}} ' + A2UI_CLOSE_TAG + ) + response_parts = parser.process_chunk(chunk3) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "List", + "children": {"componentId": "c1", "path": "/items"}, + }, + { + "id": "c1", + "component": "Text", + "text": "child 1", + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_partial_children_lists(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # Establish root + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # Send a container with 3 children, but we've only "seen" the first one + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Container", "children": ["c1",' + ' "c2", "c3"]}' + ) + response_parts = parser.process_chunk(chunk1) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "loading_c1", + **parser._placeholder_component, + }, + { + "id": "loading_c2", + **parser._placeholder_component, + }, + { + "id": "loading_c3", + **parser._placeholder_component, + }, + { + "id": "root", + "component": "Container", + "children": ["loading_c1", "loading_c2", "loading_c3"], + }, + ], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + chunk2 = ', {"id": "c1", "component": "Text", "text": "child 1"}]}}' + response_parts = parser.process_chunk(chunk2) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "loading_c2", + **parser._placeholder_component, + }, + { + "id": "loading_c3", + **parser._placeholder_component, + }, + { + "id": "root", + "component": "Container", + "children": ["c1", "loading_c2", "loading_c3"], + }, + { + "id": "c1", + "component": "Text", + "text": "child 1", + }, + ], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_data_model_before_components(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + + # 1. Data model comes BEFORE components + chunk1 = ( + '{"version": "v0.9", "updateDataModel": {"surfaceId": "s1", "value": {"name":' + ' "Alice"}}}, ' + ) + response_parts = parser.process_chunk(chunk1) + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "s1", + "value": {"name": "Alice"}, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # 2. createSurface + response_parts = parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}, ' + ) + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "s1"}, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # 3. Component with path + chunk2 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": {"path":' + ' "/name"}}]}}]' + + A2UI_CLOSE_TAG + ) + response_parts = parser.process_chunk(chunk2) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": {"path": "/name"}, + }], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_data_model_after_components(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}, ' + ) + + # 1. Component arrives first + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": {"path":' + ' "/name"}}]}}, ' + ) + response_parts = parser.process_chunk(chunk1) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": {"path": "/name"}, + }], + }, + }, + ] + + assertResponseContainsMessages(response_parts, expected) + + # 2. Send data model update + chunk2 = ( + '{"version": "v0.9", "updateDataModel": {"surfaceId": "s1", "value": {"name":' + ' "Alice"}}}]' + + A2UI_CLOSE_TAG + ) + response_parts = parser.process_chunk(chunk2) + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "s1", + "value": {"name": "Alice"}, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_partial_paths(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}, ' + ) + + # Partial path arrives + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": {"path": "/loca' + ) + response_parts = parser.process_chunk(chunk1) + assert len(response_parts) == 0 + + # Complete the path + chunk2 = 'tion"}]}}]' + A2UI_CLOSE_TAG + response_parts = parser.process_chunk(chunk2) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": {"path": "/location"}, + }], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_cut_atomic_id(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + + # Atomic key component `surfaceId` is cut. Complete it in the next chunk. + response_parts = parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "contact' + ) + assert len(response_parts) == 0 + + # Should have createSurface only + response_parts = parser.process_chunk( + '-card"}}, {"version": "v0.9", "updateComponents": {"surfaceId": "contact-card",' + ' "root": "root", "components": [{"id": "button' + ) + assert len(response_parts) == 1 + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "contact-card"}, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + response_parts = parser.process_chunk('-text"') + # Waiting for the component definition to complete + assert len(response_parts) == 0 + + # 3. Complete the component AND make it reachable from root + response_parts = parser.process_chunk( + ', "component": "Text", "text": "hi"}, {"id": "root", "component":' + ' "Card", "child": "button-text"}]}}]' + + A2UI_CLOSE_TAG + ) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "contact-card", + "root": "root", + "components": [ + {"id": "button-text", "component": "Text", "text": "hi"}, + {"id": "root", "component": "Card", "child": "button-text"}, + ], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_cut_cuttable_text(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}, ' + ) + + # Cuttable key 'text' is cut + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": "Em' + ) + response_parts = parser.process_chunk(chunk1) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": "Em", + }], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # Complete the text + chunk2 = 'ail"}]}}]' + A2UI_CLOSE_TAG + response_parts = parser.process_chunk(chunk2) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": "Email", + }], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_message_ordering_buffering(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + + # 1. updateComponents arrives BEFORE createSurface + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": "hi"}]}}, ' + ) + response_parts = parser.process_chunk(chunk1) + + # Should yield nothing yet as createSurface is missing + assert len(response_parts) == 0 + + # 2. createSurface arrives + chunk2 = ( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}]' + + A2UI_CLOSE_TAG + ) + response_parts = parser.process_chunk(chunk2) + + # Should now yield createSurface AND the buffered updateComponents + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "s1"}, + }, + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{"id": "root", "component": "Text", "text": "hi"}], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_delete_surface_buffering(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + + # deleteSurface before createSurface -> should be ignored + response_parts = parser.process_chunk( + '{"version": "v0.9", "deleteSurface": {"surfaceId": "s1"}}, ' + ) + assert len(response_parts) == 0 + + # createSurface for s1 -> creating a new surface + response_parts = parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}]' + ) + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "s1"}, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_cut_path(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}, ' + ) + + # "path" is non-cuttable, so this should not yield until closing quote + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Text", "text": {"path": "/user' + ) + response_parts = parser.process_chunk(chunk1) + # Awaiting for the closing quote of "path" + assert len(response_parts) == 0 + + # Now close it + chunk2 = '/profile"}}}}]}]' + A2UI_CLOSE_TAG + response_parts = parser.process_chunk(chunk2) + + # Should now have the placeholder with full path + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "root", + "component": "Text", + "text": {"path": "/user/profile"}, + }, + ], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_strict_begin_rendering_validation(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(A2UI_OPEN_TAG + "[") + + # 1. Totally invalid message (v0.8) + chunk = '{"unknownMessage": "invalid"}]' + with pytest.raises(ValueError, match="Validation failed"): + parser.process_chunk(chunk) + + +def test_yield_validation_failure(mock_catalog): + # Setup a more strict schema for Text component + mock_catalog.catalog_schema[CATALOG_COMPONENTS_KEY]["Text"]["required"] = ["text"] + parser = A2uiStreamParser(catalog=mock_catalog) + + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}}, ' + ) + + # Send an invalid message (missing required field inside envelope) + chunk = '{"updateComponents": {"components": []}}' + + with pytest.raises(ValueError, match="Validation failed"): + parser.process_chunk(chunk) + + +def test_delta_streaming_correctness(mock_catalog): + """Verifies that the parser correctly assembles components from small deltas.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + # 1. Start with open tag + parser.process_chunk(A2UI_OPEN_TAG) + response_parts = parser.process_chunk("[") + assert response_parts == [] + + # 2. Stream createSurface char by char + v09_json = ( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}' + ) + for char in v09_json[:-1]: + response_parts = parser.process_chunk(char) + assert response_parts == [] + + response_parts = parser.process_chunk(v09_json[-1]) + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "s1"}, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # 3. Stream updateComponents with component splitting across deltas + response_parts = parser.process_chunk( + ', {"updateComponents": {"surfaceId": "s1", "components": [' + ) + assert response_parts == [] + response_parts = parser.process_chunk('{"id": "root", "compon') + assert response_parts == [] + + # The parser is eager and yields immediately because 'text' is in CUTTABLE_KEYS + response_parts = parser.process_chunk('ent": "Text", "text": "hello') + + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": "hello", + }], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # Stream the rest of the text + response_parts = parser.process_chunk(" world") + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [{ + "id": "root", + "component": "Text", + "text": "hello world", + }], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_multiple_re_yielding_scenarios(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "s1"}}, ' + ) + + # 1. Component with 2 paths + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "components":' + ' [{"id": "root", "component": "Container", "children": ["c1", "c2"]}, {"id":' + ' "c1", "component": "Text", "text": {"path": "/p1"}}, {"id": "c2", "component":' + ' "Text", "text": {"path": "/p2"}}]}}, ' + ) + response_parts = parser.process_chunk(chunk1) + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "c1", + "component": "Text", + "text": {"path": "/p1"}, + }, + { + "id": "c2", + "component": "Text", + "text": {"path": "/p2"}, + }, + { + "id": "root", + "component": "Container", + "children": ["c1", "c2"], + }, + ], + }, + }, + ] + + assertResponseContainsMessages(response_parts, expected) + + # 2. Add p1 + response_parts = parser.process_chunk( + '{"version": "v0.9", "updateDataModel": {"surfaceId": "s1", "value": {"p1":' + ' "v1"}}}, ' + ) + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "s1", + "value": {"p1": "v1"}, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # 3. Add p2 + response_parts = parser.process_chunk( + '{"version": "v0.9", "updateDataModel": {"surfaceId": "s1", "value": {"p2":' + ' "v2"}}}' + ) + + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "s1", + "value": {"p2": "v2"}, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_incremental_data_model_streaming(mock_catalog): + """Verifies that the parser yields surface updates as items arrive in a updateDataModel stream.""" + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk(f"{A2UI_OPEN_TAG}[") + + # 1. Establish surface + parser.process_chunk( + '{"version": "v0.9", "createSurface": {"catalogId": "test_catalog", "surfaceId":' + ' "default"}}, ' + ) + + # 2. Establish surface components with a data binding + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "default", "root":' + ' "item-list", "components": [{"id": "item-list", "component": "List",' + ' "children": {"componentId": "template-name", "path": "/items"}}, {"id":' + ' "template-name", "component": "Text", "text": {"path": "/name"}}]}}, ' + ) + response_parts = parser.process_chunk(chunk1) + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "default", + "root": "item-list", + "components": [ + { + "id": "item-list", + "component": "List", + "children": { + "componentId": "template-name", + "path": "/items", + }, + }, + { + "id": "template-name", + "component": "Text", + "text": {"path": "/name"}, + }, + ], + }, + }] + assertResponseContainsMessages(response_parts, expected) + + # 3. Start streaming updateDataModel + response_parts = parser.process_chunk( + '{"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value":' + ' {"items": {' + ) + # The parser yields the data model early once it sniffs the start of it + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": {"items": {}}, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # Add Item 1 + response_parts = parser.process_chunk('"item1": {"name": "Item 1"}, ') + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "items": {"item1": {"name": "Item 1"}}, + }, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + # Add Item 2 + response_parts = parser.process_chunk( + '"item2": {"name": "Item 2"}}}}}] ' + A2UI_CLOSE_TAG + ) + expected = [ + { + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "items": { + "item1": {"name": "Item 1"}, + "item2": {"name": "Item 2"}, + }, + }, + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_sniff_partial_invalid_datamodel_fails_gracefully(mock_catalog): + """Verifies that sniffing a partial DM update that would fail validation doesn't crash.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + # 1. Provide a partial chunk where the LAST item in value is incomplete + chunk = ( + '[ {"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value": ' + '{"name": "John", "incomplete": ' + ) + + response_parts = parser.process_chunk(A2UI_OPEN_TAG + chunk) + # Should yield only the FIRST valid entry (name: John) + expected = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": {"name": "John"}, + }, + }] + assertResponseContainsMessages(response_parts, expected) + + # 2. Provide a chunk that is COMPLETELY invalid (missing value on only entry) + chunk2 = ( + A2UI_OPEN_TAG + + '{"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value":' + ' {"unnamed":' + ) + response_parts = parser.process_chunk(chunk2) + assert response_parts == [] + + +def test_sniff_partial_datamodel_with_cut_key(mock_catalog): + """Verifies that an unclosed string like {"key or {"key": "val still yields previous valid entries.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + # Chunk ending with an open key name + chunk1 = ( + A2UI_OPEN_TAG + + '[ {"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value": ' + '{"infoLink": "[More Info](x)", "rat' + ) + response_parts = parser.process_chunk(chunk1) + + expected1 = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": {"infoLink": "[More Info](x)"}, + }, + }] + assertResponseContainsMessages(response_parts, expected1) + + # Chunk completing the key and value + chunk2 = 'ing": "5 star"' + response_parts = parser.process_chunk(chunk2) + expected2 = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": {"rating": "5 star"}, + }, + }] + assertResponseContainsMessages(response_parts, expected2) + + # Chunk ending inside a string value that IS NOT cuttable due to URL heuristics + chunk3_url = ', "imageUrl": "http://localhost:10002/static/map' + response_parts = parser.process_chunk(chunk3_url) + # Should yield nothing because imageUrl is incomplete + assert response_parts == [] + + # Finish the chunk + chunk4 = 's.png"}}] ' + A2UI_CLOSE_TAG + response_parts = parser.process_chunk(chunk4) + expected4 = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "imageUrl": "http://localhost:10002/static/maps.png", + }, + }, + }] + assertResponseContainsMessages(response_parts, expected4) + + +def test_sniff_partial_datamodel_cumulative_unmodified_keys(mock_catalog): + """Verifies that unchanged keys are retained in partial updateDataModels.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + # Chunk 1: 'title' is complete, 'items' is unclosed + chunk1 = ( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value":' + ' {"title": "Top Restaurants", "items": {' + ) + + response_parts = parser.process_chunk(chunk1) + + expected1 = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "title": "Top Restaurants", + "items": {}, + }, + }, + }] + assertResponseContainsMessages(response_parts, expected1) + + # Chunk 2: 'items' gets populated with nested object + chunk2 = '"item1": {"name": "Restaurant A"}' + response_parts = parser.process_chunk(chunk2) + + expected2 = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "items": {"item1": {"name": "Restaurant A"}}, + }, + }, + }] + assertResponseContainsMessages(response_parts, expected2) + + +def test_sniff_partial_datamodel_prunes_empty_keys(mock_catalog): + """Verifies that entries with only a key and no value are pruned from partial updateDataModels.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + chunk = ( + A2UI_OPEN_TAG + + '[ {"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value": {' + + '"title": "Top Restaurants",' + + '"items": {"item1": {"name": "Food", "detail": "Spicy", "imageUrl":' + ) + response_parts = parser.process_chunk(chunk) + + # Should yield title, name, detail, but NOT imageUrl because it is incomplete + expected = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "title": "Top Restaurants", + "items": { + "item1": { + "name": "Food", + "detail": "Spicy", + } + }, + }, + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_sniff_partial_datamodel_prunes_empty_trailing_dict(mock_catalog): + """Verifies that an incomplete trailing empty dictionary '{}' is dropped.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + chunk = ( + A2UI_OPEN_TAG + + '[ {"version": "v0.9", "updateDataModel": {"surfaceId": "default", "value": {' + + '"title": "Top Restaurants",' + + '"items": {"item2": {"name": "Han Dynasty", "imageUrl":' + ' "http://localhost:10002/static/mapotofu.jpeg", "broken":' + ) + response_parts = parser.process_chunk(chunk) + + expected = [{ + "version": "v0.9", + "updateDataModel": { + "surfaceId": "default", + "value": { + "title": "Top Restaurants", + "items": { + "item2": { + "name": "Han Dynasty", + "imageUrl": "http://localhost:10002/static/mapotofu.jpeg", + } + }, + }, + }, + }] + assertResponseContainsMessages(response_parts, expected) + + +def test_sniff_partial_component_discards_empty_children_dict(mock_catalog): + """Verifies that an incomplete component with an empty children dictionary is discarded until populated.""" + parser = A2uiStreamParser(catalog=mock_catalog) + + chunk = ( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "default"}},' + + '{"version": "v0.9", "updateComponents": {"surfaceId": "default", "root":' + ' "root-column", "components": [' + + '{"id": "root-column", "component": "Column", "children": ["item-list"]},' + + '{"id": "item-list", "component": "List", "direction": "vertical",' + ' "children": {' + ) + + response_parts = parser.process_chunk(chunk) + + # item-list has {"children": {}}. It should be completely discarded from _seen_components. + # Its parent, root-column, will then replace the missing item-list with a loading placeholder. + expected = [ + { + "version": "v0.9", + "createSurface": {"catalogId": "test_catalog", "surfaceId": "default"}, + }, + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "default", + "root": "root-column", + "components": [ + { + "id": "loading_item-list", + "component": "Row", + "children": [], + }, + { + "id": "root-column", + "component": "Column", + "children": ["loading_item-list"], + }, + ], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_partial_empty_dict_discarded(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + + # Establish root + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + # Send a component with an empty dictionary that will be closed by the fixer + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "root",' + ' "components": [{"id": "root", "component": "Column", "children": {' + ) + response_parts = parser.process_chunk(chunk1) + assertResponseContainsNoA2UI(response_parts) + + chunk2 = ( + '"componentId": "c1", "path": "/items"}}, {"id": "c1", "component": "Text",' + ' "text": "Child 1"}]}} ' + + A2UI_CLOSE_TAG + ) + response_parts = parser.process_chunk(chunk2) + + expected = [ + { + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "root", + "components": [ + { + "id": "c1", + "component": "Text", + "text": "Child 1", + }, + { + "id": "root", + "component": "Column", + "children": {"componentId": "c1", "path": "/items"}, + }, + ], + }, + }, + ] + assertResponseContainsMessages(response_parts, expected) + + +def test_sniff_partial_component_enforces_required_fields(mock_catalog): + parser = A2uiStreamParser(catalog=mock_catalog) + parser.process_chunk( + A2UI_OPEN_TAG + + '[{"version": "v0.9", "createSurface": {"catalogId": "test_catalog",' + ' "surfaceId": "s1"}},' + ) + + chunk1 = ( + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "root": "c1",' + ' "components": [{"id": "c1", "component": "AudioPlayer", "description": "almost' + ' ready"' + ) + response_parts = parser.process_chunk(chunk1) + assertResponseContainsNoA2UI(response_parts) + + chunk2 = ', "url": "http://audio.mp3"}]}} ' + A2UI_CLOSE_TAG + response_parts = parser.process_chunk(chunk2) + + expected = [{ + "version": "v0.9", + "updateComponents": { + "surfaceId": "s1", + "root": "c1", + "components": [{ + "id": "c1", + "component": "AudioPlayer", + "description": "almost ready", + "url": "http://audio.mp3", + }], + }, + }] + assertResponseContainsMessages(response_parts, expected) diff --git a/agent_sdks/python/tests/core/parser/test_version_handlers.py b/agent_sdks/python/tests/core/parser/test_version_handlers.py deleted file mode 100644 index 2b3291f4b..000000000 --- a/agent_sdks/python/tests/core/parser/test_version_handlers.py +++ /dev/null @@ -1,133 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import pytest -from unittest.mock import MagicMock -from a2ui.core.parser.version_handlers import A2uiVersionHandler, A2uiV08Handler, A2uiV09Handler -from a2ui.core.schema.constants import VERSION_0_8, VERSION_0_9 -from a2ui.core.parser.constants import ( - MSG_TYPE_BEGIN_RENDERING, - MSG_TYPE_SURFACE_UPDATE, - MSG_TYPE_CREATE_SURFACE, - MSG_TYPE_UPDATE_COMPONENTS, -) - - -def test_detect_version(): - # v0.9 detection - assert A2uiVersionHandler.detect_version('{"version": "v0.9"}') == VERSION_0_9 - assert A2uiVersionHandler.detect_version('{"updateComponents": {}}') == VERSION_0_9 - assert A2uiVersionHandler.detect_version('{"createSurface": {}}') == VERSION_0_9 - - # v0.8 detection - assert A2uiVersionHandler.detect_version('{"beginRendering": {}}') == VERSION_0_8 - assert A2uiVersionHandler.detect_version('{"surfaceUpdate": {}}') == VERSION_0_8 - - # Unknown - assert A2uiVersionHandler.detect_version('{"foo": "bar"}') is None - - -def test_v08_handler_sniff_metadata(): - handler = A2uiV08Handler() - parser = MagicMock() - parser.surface_id = None - parser.root_id = None - parser.msg_types = [] - - json_buffer = '{"surfaceId": "s1", "beginRendering": {"root": "r1"}}' - handler.sniff_metadata(json_buffer, parser) - - assert parser.surface_id == "s1" - assert parser.root_id == "r1" - parser.add_msg_type.assert_called_with(MSG_TYPE_BEGIN_RENDERING) - - json_buffer = json.dumps([ - {"beginRendering": {"surfaceId": "s1", "root": "r1"}}, - { - "surfaceUpdate": { - "surfaceId": "s1", - "components": [{ - "id": "c1", - "component": {"Text": {"text": {"literalString": "hello"}}}, - }], - } - }, - ]) - handler.sniff_metadata(json_buffer, parser) - assert parser.surface_id == "s1" - assert parser.root_id == "r1" - parser.add_msg_type.assert_any_call(MSG_TYPE_BEGIN_RENDERING) - parser.add_msg_type.assert_any_call(MSG_TYPE_SURFACE_UPDATE) - - -def test_v08_handler_handle_complete_object(): - handler = A2uiV08Handler() - parser = MagicMock() - parser.root_id = None - parser.seen_components = {} - messages = [] - - # beginRendering - obj = {"beginRendering": {"root": "r1"}} - assert handler.handle_complete_object(obj, parser, messages) is True - assert parser.root_id == "r1" - assert parser.buffered_begin_rendering == obj - - # surfaceUpdate - obj = {"surfaceUpdate": {"components": [{"id": "c1", "component": {}}]}} - assert handler.handle_complete_object(obj, parser, messages) is True - assert "c1" in parser.seen_components - parser.yield_reachable.assert_called_once() - - -def test_v09_handler_sniff_metadata(): - handler = A2uiV09Handler() - parser = MagicMock() - parser.surface_id = None - parser.root_id = None - parser.msg_types = [] - - json_buffer = '{"surfaceId": "s1", "updateComponents": {"root": "r2"}}' - handler.sniff_metadata(json_buffer, parser) - - assert parser.surface_id == "s1" - assert parser.root_id == "r2" - parser.add_msg_type.assert_called_with(MSG_TYPE_UPDATE_COMPONENTS) - - -def test_v09_handler_handle_complete_object(): - handler = A2uiV09Handler() - parser = MagicMock() - parser.root_id = None - parser.seen_components = {} - messages = [] - - # createSurface - obj = {"createSurface": {"surfaceId": "s1"}} - assert handler.handle_complete_object(obj, parser, messages) is True - assert parser.root_id == "root" # Default for v0.9 - assert parser.buffered_begin_rendering == obj - - # updateComponents - obj = { - "updateComponents": { - "root": "custom", - "components": [{"id": "custom", "component": {}}], - } - } - assert handler.handle_complete_object(obj, parser, messages) is True - assert parser.root_id == "custom" - assert "custom" in parser.seen_components - parser.yield_reachable.assert_called_once()