Address review comments

Author: nan-yu

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

@@ -106,6 +106,7 @@ def __init__(self, catalog: "A2uiCatalog" = None):
     self._yielded_contents: Dict[Any, str] = {}
 
     self._root_ids: Dict[str, str] = {}  # The root component IDs mapped per surface
+    self._default_root_id: Optional[str] = None  # Base default root ID for the protocol
     self._unbound_root_id: Optional[str] = (
         None  # Temporary holding variable for when root arrives before surfaceId
     )
@@ -153,8 +154,13 @@ def surface_id(self, value: Optional[str]):
   @property
   def root_id(self) -> Optional[str]:
     if self._surface_id:
-      return self._root_ids.get(self._surface_id)
-    return self._unbound_root_id
+      return self._root_ids.get(self._surface_id, self._default_root_id)
+    # Return unbound root ID if explicitly sniffed, otherwise use protocol default
+    return (
+        self._unbound_root_id
+        if self._unbound_root_id is not None
+        else self._default_root_id
+    )
 
   @root_id.setter
   def root_id(self, value: Optional[str]):

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

@@ -64,13 +64,22 @@ def _data_model_msg_type(self) -> str:
 
   def _sniff_metadata(self):
     """Sniffs for v0.8 metadata in the json_buffer."""
-    matches = re.findall(r'"surfaceId"\s*:\s*"([^"]+)"', self._json_buffer)
-    if matches:
-      self.surface_id = matches[-1]
 
-    root_matches = re.findall(r'"root"\s*:\s*"([^"]+)"', self._json_buffer)
-    if root_matches:
-      self.root_id = root_matches[-1]
+    def get_latest_value(key: str) -> Optional[str]:
+      idx = len(self._json_buffer)
+      while True:
+        idx = self._json_buffer.rfind(f'"{key}"', 0, idx)
+        if idx == -1:
+          return None
+        match = re.match(rf'"{key}"\s*:\s*"([^"]+)"', self._json_buffer[idx:])
+        if match:
+          return match.group(1)
+
+    self.surface_id = get_latest_value('surfaceId')
+
+    parsed_root = get_latest_value('root')
+    if parsed_root is not None:
+      self.root_id = parsed_root
 
     if f'"{MSG_TYPE_BEGIN_RENDERING}":' in self._json_buffer:
       self.add_msg_type(MSG_TYPE_BEGIN_RENDERING)

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

@@ -28,7 +28,7 @@ class A2uiStreamParserV09(A2uiStreamParser):
   def __init__(self, catalog=None):
     super().__init__(catalog=catalog)
     # v0.9 default root is "root"
-    self.root_id = DEFAULT_ROOT_ID
+    self._default_root_id = DEFAULT_ROOT_ID
 
   @property
   def _placeholder_component(self) -> Dict[str, Any]:
@@ -56,13 +56,22 @@ def is_protocol_msg(self, obj: Dict[str, Any]) -> bool:
 
   def _sniff_metadata(self):
     """Sniffs for v0.9 metadata in the json_buffer."""
-    matches = re.findall(r'"surfaceId"\s*:\s*"([^"]+)"', self._json_buffer)
-    if matches:
-      self.surface_id = matches[-1]
 
-    root_matches = re.findall(r'"root"\s*:\s*"([^"]+)"', self._json_buffer)
-    if root_matches:
-      self.root_id = root_matches[-1]
+    def get_latest_value(key: str) -> Optional[str]:
+      idx = len(self._json_buffer)
+      while True:
+        idx = self._json_buffer.rfind(f'"{key}"', 0, idx)
+        if idx == -1:
+          return None
+        match = re.match(rf'"{key}"\s*:\s*"([^"]+)"', self._json_buffer[idx:])
+        if match:
+          return match.group(1)
+
+    self.surface_id = get_latest_value('surfaceId')
+
+    parsed_root = get_latest_value('root')
+    if parsed_root is not None:
+      self.root_id = parsed_root
 
     if f'"{MSG_TYPE_CREATE_SURFACE}":' in self._json_buffer:
       self.add_msg_type(MSG_TYPE_CREATE_SURFACE)