Fix streaming top-level objects without list wrapping

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

@@ -451,14 +451,11 @@ def _fix_json(self, fragment: str) -> str:
   def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]):
     for char in chunk:
       char_handled = False
-      if not self._in_top_level_list:
+      if self._brace_count == 0:
         if char == "[":
-          if self._brace_count == 0:
-            self._in_top_level_list = True
-          self._brace_stack.append(("[", len(self._json_buffer)))
-          self._json_buffer += "["
-          self._brace_count += 1
-          char_handled = True
+          self._in_top_level_list = True
+        elif char == "{":
+          pass
         else:
           continue
 

agent_sdks/python/tests/parser/test_streaming_v09.py

@@ -452,3 +452,42 @@ def test_v09_path_heuristic_absolute_path(mock_catalog):
   assert len(messages) > 0
   comp = messages[0][MSG_TYPE_UPDATE_COMPONENTS]["components"][0]
   assert comp["text"]["path"] == "/absolute/path"
+
+
+def test_v09_single_top_level_object(mock_catalog):
+  """Tests that v0.9 supports a single top-level object without array wrapping."""
+  parser = A2uiStreamParser(catalog=mock_catalog)
+
+  chunk = (
+      A2UI_OPEN_TAG
+      + '{"version": "v0.9", "createSurface": {"surfaceId": "s1", "catalogId": "c1"}}'
+      + A2UI_CLOSE_TAG
+  )
+  messages = []
+  for part in parser.process_chunk(chunk):
+    if part.a2ui_json:
+      messages.extend(part.a2ui_json)
+
+  assert len(messages) == 1
+  assert messages[0]["createSurface"]["surfaceId"] == "s1"
+
+
+def test_v09_multiple_top_level_objects(mock_catalog):
+  """Tests that v0.9 supports multiple consecutive top-level objects without array wrapping."""
+  parser = A2uiStreamParser(catalog=mock_catalog)
+
+  chunk = (
+      A2UI_OPEN_TAG
+      + '{"version": "v0.9", "createSurface": {"surfaceId": "s1", "catalogId": "c1"}}\n'
+      + '{"version": "v0.9", "updateComponents": {"surfaceId": "s1", "components":'
+      ' [{"id": "root", "component": "Text", "text": "Hello"}]}}'
+      + A2UI_CLOSE_TAG
+  )
+  messages = []
+  for part in parser.process_chunk(chunk):
+    if part.a2ui_json:
+      messages.extend(part.a2ui_json)
+
+  assert len(messages) == 2
+  assert messages[0]["createSurface"]["surfaceId"] == "s1"
+  assert messages[1][MSG_TYPE_UPDATE_COMPONENTS]["components"][0]["text"] == "Hello"