diff --git a/agent_sdks/python/src/a2ui/parser/streaming.py b/agent_sdks/python/src/a2ui/parser/streaming.py index 2c7ea3e0e4..384a8ded55 100644 --- a/agent_sdks/python/src/a2ui/parser/streaming.py +++ b/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 diff --git a/agent_sdks/python/tests/parser/test_streaming_v09.py b/agent_sdks/python/tests/parser/test_streaming_v09.py index 2f78ba598e..01b8fabb80 100644 --- a/agent_sdks/python/tests/parser/test_streaming_v09.py +++ b/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"