feat(composer): A2UI Dojo — interactive JSONL playback and streaming viewer

agent_sdks/python/src/a2ui/a2a.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 
 import logging
-from typing import Any, Optional, List
+from typing import Any, Optional, List, AsyncIterable, TYPE_CHECKING
 
+if TYPE_CHECKING:
+  from a2ui.core.parser.streaming import A2uiStreamParser
 from a2a.server.agent_execution import RequestContext
 from a2a.types import (
     AgentExtension,
@@ -247,3 +249,39 @@ def try_activate_a2ui_extension(
     return selected_uri.replace(f"{A2UI_EXTENSION_BASE_URI}/v", "")
 
   return None
+
+
+async def stream_response_to_parts(
+    parser: "A2uiStreamParser",
+    token_stream: AsyncIterable[str],
+) -> AsyncIterable[Part]:
+  """Helper to parse a stream of LLM tokens into A2A Parts incrementally.
+
+  Args:
+      parser: A2uiStreamParser instance to process the stream.
+      token_stream: An async iterable of strings (tokens).
+
+  Yields:
+      A2A Part objects as they are discovered in the stream.
+  """
+  async for token in token_stream:
+    logger.info("-----------------------------")
+    logger.info(f"--- AGENT: Received token:\n{token}")
+    response_parts = parser.process_chunk(token)
+    logger.info(
+        f"--- AGENT: Response parts:\n{[part.a2ui_json for part in response_parts]}\n"
+    )
+    logger.info("-----------------------------")
+
+    for part in response_parts:
+      if part.text:
+        yield Part(root=TextPart(text=part.text))
+
+      if part.a2ui_json:
+        json_data = part.a2ui_json
+
+        if isinstance(json_data, list):
+          for message in json_data:
+            yield create_a2ui_part(message)
+        else:
+          yield create_a2ui_part(json_data)

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

@@ -88,9 +88,6 @@ def __init__(self, catalog: "A2uiCatalog" = None):
     # 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._dirty_components: Set[str] = (
-        set()
-    )  # Set of component IDs that need re-yielding
     self._components_with_placeholders: Set[str] = (
         set()
     )  # cid set that were yielded as placeholders
@@ -695,9 +692,8 @@ def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None:
                 delta_msg = {MSG_TYPE_DATA_MODEL_UPDATE: delta_msg_payload}
                 self._yield_messages([delta_msg], messages, strict_integrity=False)
                 self._yielded_data_model.update(contents_dict)
-                # Update internal model to trigger re-yielding of resolved placeholders
+                # Update internal model for path resolution
                 self.update_data_model(dm_obj, messages)
-                self.yield_reachable(messages, check_root=False)
 
         # v0.9: updateDataModel
         elif MSG_TYPE_UPDATE_DATA_MODEL in obj:
@@ -889,9 +885,6 @@ def update_data_model(
 
     self._data_model.update(contents)
     updated_keys = set(contents.keys())
-    for cid, paths in self._comp_paths.items():
-      if paths & updated_keys:
-        self._dirty_components.add(cid)
 
   def _handle_complete_object(
       self,
@@ -1078,9 +1071,7 @@ def yield_reachable(
         return
 
       should_yield = False
-      if (available_reachable - yielded_for_surface) or (
-          available_reachable & self._dirty_components
-      ):
+      if available_reachable - yielded_for_surface:
         should_yield = True
       else:
         # Check if any yielded component's content has changed for this surface
@@ -1154,8 +1145,6 @@ def yield_reachable(
           cid = comp["id"]
           self._yielded_contents[(surface_id, cid)] = json.dumps(comp, sort_keys=True)
           self._yielded_placeholders[cid] = self._get_placeholders(comp)
-          if cid in self._dirty_components:
-            self._dirty_components.remove(cid)
 
     except ValueError as e:
       if "Circular reference detected" in str(e):

agent_sdks/python/src/a2ui/core/schema/constants.py

@@ -52,8 +52,12 @@
 
 DEFAULT_WORKFLOW_RULES = f"""
 The generated response MUST follow these rules:
-1.  The response can contain one or more A2UI JSON blocks.
-2.  Each A2UI JSON block MUST be wrapped in `{A2UI_OPEN_TAG}` and `{A2UI_CLOSE_TAG}` tags.
-3.  Between or around these blocks, you can provide conversational text.
-4.  The JSON part MUST be a single, raw JSON object (usually a list of A2UI messages) and MUST validate against the provided A2UI JSON SCHEMA.
+- The response can contain one or more A2UI JSON blocks.
+- Each A2UI JSON block MUST be wrapped in `{A2UI_OPEN_TAG}` and `{A2UI_CLOSE_TAG}` tags.
+- Between or around these blocks, you can provide conversational text.
+- The JSON part MUST be a single, raw JSON object (usually a list of A2UI messages) and MUST validate against the provided A2UI JSON SCHEMA.
+- Top-Down Component Ordering: Within the `components` list of a message:
+    - The 'root' component MUST be the FIRST element.
+    - Parent components MUST appear before their child components.
+    This specific ordering allows the streaming parser to yield and render the UI incrementally as it arrives.
 """

agent_sdks/python/tests/core/parser/test_streaming_v08.py

@@ -1091,19 +1091,6 @@ def test_data_model_after_components(mock_catalog):
               }],
           }
       },
-      {
-          "surfaceUpdate": {
-              "surfaceId": "s1",
-              "components": [{
-                  "id": "root",
-                  "component": {
-                      "Text": {
-                          "text": {"path": "/name"},
-                      }
-                  },
-              }],
-          }
-      },
   ]
   assertResponseContainsMessages(response_parts, expected)
 
@@ -1518,37 +1505,6 @@ def test_multiple_re_yielding_scenarios(mock_catalog):
               "contents": [{"key": "p1", "valueString": "v1"}],
           }
       },
-      {
-          "surfaceUpdate": {
-              "surfaceId": "s1",
-              "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)
 
@@ -1564,37 +1520,6 @@ def test_multiple_re_yielding_scenarios(mock_catalog):
               "contents": [{"key": "p2", "valueString": "v2"}],
           }
       },
-      {
-          "surfaceUpdate": {
-              "surfaceId": "s1",
-              "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)
 
@@ -1650,38 +1575,13 @@ def test_incremental_data_model_streaming(mock_catalog):
       ' "valueMap": ['
   )
   # The parser yields the data model early once it sniffs the start of it
-  # AND it triggers a re-yield of reachable components that depend on 'items'
   expected = [
       {
           "dataModelUpdate": {
               "contents": [{"key": "items", "valueMap": []}],
               "surfaceId": "default",
           }
       },
-      {
-          "surfaceUpdate": {
-              "surfaceId": "default",
-              "components": [
-                  {
-                      "id": "item-list",
-                      "component": {
-                          "List": {
-                              "children": {
-                                  "template": {
-                                      "componentId": "template-name",
-                                      "dataBinding": "/items",
-                                  }
-                              }
-                          }
-                      },
-                  },
-                  {
-                      "id": "template-name",
-                      "component": {"Text": {"text": {"path": "/name"}}},
-                  },
-              ],
-          }
-      },
   ]
   assertResponseContainsMessages(response_parts, expected)
 
@@ -1697,30 +1597,6 @@ def test_incremental_data_model_streaming(mock_catalog):
               }],
           }
       },
-      {
-          "surfaceUpdate": {
-              "surfaceId": "default",
-              "components": [
-                  {
-                      "id": "item-list",
-                      "component": {
-                          "List": {
-                              "children": {
-                                  "template": {
-                                      "componentId": "template-name",
-                                      "dataBinding": "/items",
-                                  }
-                              }
-                          }
-                      },
-                  },
-                  {
-                      "id": "template-name",
-                      "component": {"Text": {"text": {"path": "/name"}}},
-                  },
-              ],
-          }
-      },
   ]
   assertResponseContainsMessages(response_parts, expected)
 
@@ -1739,30 +1615,6 @@ def test_incremental_data_model_streaming(mock_catalog):
               }],
           }
       },
-      {
-          "surfaceUpdate": {
-              "surfaceId": "default",
-              "components": [
-                  {
-                      "id": "item-list",
-                      "component": {
-                          "List": {
-                              "children": {
-                                  "template": {
-                                      "componentId": "template-name",
-                                      "dataBinding": "/items",
-                                  }
-                              }
-                          }
-                      },
-                  },
-                  {
-                      "id": "template-name",
-                      "component": {"Text": {"text": {"path": "/name"}}},
-                  },
-              ],
-          }
-      },
   ]
   assertResponseContainsMessages(response_parts, expected)