test(examples-chat): regression coverage for GenUI emit coalescing (PR C) (#299)

* docs(plans): progressive GenUI rendering + bubble coalescing (chat-side wiring)

Three independent PRs per spec's phasing (order A → B → C):
- PR A — Surface store + catalog shape (lib only) [this branch]
- PR B — <a2ui-surface> per-component rendering (lib only)
- PR C — Backend coalescing + envelope reordering (Python only)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test(examples-chat): GenUI emit coalescing + reorder regressions

Asserts emit_generated_surface returns 2 replacements (3-message
thread post-merge, not 4), preserves the upstream tool-call AI's
id, tool_calls, additional_kwargs, response_metadata, and orders
the wrapped envelopes surfaceUpdate -> beginRendering -> dataModelUpdate.

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: blove

examples/chat/python/tests/test_graph_smoke.py

@@ -272,3 +272,105 @@ def test_generate_a2ui_schema_tool_is_removed(self):
         tool_node = _builder.nodes["tools"].runnable
         tool_names = list(tool_node.tools_by_name.keys())
         assert "generate_a2ui_schema" not in tool_names
+
+
+import json
+from uuid import uuid4
+
+
+class TestEmitInPlaceCoalescing:
+    """Regression: emit_generated_surface MUST coalesce the GenUI turn
+    into a single AI message (3-message thread, not 4), preserving the
+    upstream tool-call AI's id, tool_calls, additional_kwargs, and
+    response_metadata. Envelopes inside the wrapped content MUST be
+    ordered surfaceUpdate -> beginRendering -> dataModelUpdate × N."""
+
+    def _run(self, state):
+        from src.graph import emit_generated_surface
+        return asyncio.run(emit_generated_surface(state))
+
+    def test_post_emit_thread_has_three_messages_not_four(self):
+        original_ai_id = str(uuid4())
+        tool_call_id = "call_123"
+        envelopes = [
+            {"dataModelUpdate": {"surfaceId": "s1", "contents": [{"key": "name", "valueString": "Ada"}]}},
+            {"surfaceUpdate": {"surfaceId": "s1", "components": [{"id": "c1", "component": {"TextField": {"value": "{$.name}"}}}]}},
+            {"beginRendering": {"surfaceId": "s1", "root": "c1"}},
+        ]
+        tool_call_ai = AIMessage(
+            id=original_ai_id,
+            content="",
+            tool_calls=[{"id": tool_call_id, "name": "render_a2ui_surface", "args": {}, "type": "tool_call"}],
+        )
+        tool_msg = ToolMessage(
+            id="tool_msg_1",
+            tool_call_id=tool_call_id,
+            content=json.dumps(envelopes),
+        )
+        state = {"messages": [HumanMessage(content="render a card"), tool_call_ai, tool_msg]}
+
+        result = self._run(state)
+
+        # add_messages will REPLACE the tool message (same id) and the
+        # AI message (same id) — net thread length stays 3 after merge.
+        # Here we just assert the returned message list is 2 entries
+        # (replacements only), both targeting the upstream ids.
+        returned = result["messages"]
+        assert len(returned) == 2, f"expected 2 replacements, got {len(returned)}: {returned}"
+        # ToolMessage replacement keeps its id and tool_call_id
+        tool_replacement = next(m for m in returned if isinstance(m, ToolMessage))
+        assert tool_replacement.id == tool_msg.id
+        assert tool_replacement.tool_call_id == tool_call_id
+        # AI replacement keeps the upstream AI id (in-place merge)
+        ai_replacement = next(m for m in returned if isinstance(m, AIMessage))
+        assert ai_replacement.id == original_ai_id, (
+            "AI replacement must reuse upstream tool-call AI id for in-place merge"
+        )
+
+    def test_preserves_tool_calls_additional_kwargs_response_metadata(self):
+        original_ai_id = str(uuid4())
+        tool_call_id = "call_xyz"
+        envelopes = [
+            {"surfaceUpdate": {"surfaceId": "s1", "components": []}},
+            {"beginRendering": {"surfaceId": "s1", "root": "c1"}},
+        ]
+        tool_call_ai = AIMessage(
+            id=original_ai_id,
+            content="",
+            tool_calls=[{"id": tool_call_id, "name": "render_a2ui_surface", "args": {}, "type": "tool_call"}],
+            additional_kwargs={"reasoning": "the user wants a card"},
+            response_metadata={"finish_reason": "tool_calls"},
+        )
+        tool_msg = ToolMessage(id="t1", tool_call_id=tool_call_id, content=json.dumps(envelopes))
+        state = {"messages": [HumanMessage(content="x"), tool_call_ai, tool_msg]}
+
+        result = self._run(state)
+        ai_replacement = next(m for m in result["messages"] if isinstance(m, AIMessage))
+        assert ai_replacement.tool_calls and ai_replacement.tool_calls[0]["id"] == tool_call_id
+        assert ai_replacement.additional_kwargs.get("reasoning") == "the user wants a card"
+        assert ai_replacement.response_metadata.get("finish_reason") == "tool_calls"
+
+    def test_envelopes_reordered_to_surface_begin_data(self):
+        tool_call_id = "call_r"
+        envelopes_unordered = [
+            {"dataModelUpdate": {"surfaceId": "s1", "contents": [{"key": "n", "valueString": "1"}]}},
+            {"dataModelUpdate": {"surfaceId": "s1", "contents": [{"key": "m", "valueString": "2"}]}},
+            {"beginRendering": {"surfaceId": "s1", "root": "c1"}},
+            {"surfaceUpdate": {"surfaceId": "s1", "components": []}},
+        ]
+        tool_call_ai = AIMessage(
+            id="ai-1",
+            content="",
+            tool_calls=[{"id": tool_call_id, "name": "render_a2ui_surface", "args": {}, "type": "tool_call"}],
+        )
+        tool_msg = ToolMessage(id="t-1", tool_call_id=tool_call_id, content=json.dumps(envelopes_unordered))
+        state = {"messages": [HumanMessage(content="x"), tool_call_ai, tool_msg]}
+
+        result = self._run(state)
+        ai = next(m for m in result["messages"] if isinstance(m, AIMessage))
+        # Strip the A2UI_PREFIX wrapper before splitting JSONL.
+        lines = [ln for ln in ai.content.split("\n") if ln.strip() and not ln.startswith("---a2ui_JSON---")]
+        keys = [list(json.loads(ln).keys())[0] for ln in lines]
+        assert keys == ["surfaceUpdate", "beginRendering", "dataModelUpdate", "dataModelUpdate"], (
+            f"expected surfaceUpdate -> beginRendering -> dataModelUpdate × N, got {keys}"
+        )