From 92668e733598f4971a0fd2e98a7916d2b6528fe5 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Tue, 21 Jul 2026 23:19:22 +0800
Subject: [PATCH 1/3] Preserve model emission order in AG-UI messages snapshot

---
 .../ag-ui/agent_framework_ag_ui/_agent_run.py | 54 +++++++++++
 .../agent_framework_ag_ui/_run_common.py      | 43 ++++++++
 python/packages/ag-ui/tests/ag_ui/test_run.py | 97 ++++++++++++++++++-
 3 files changed, 191 insertions(+), 3 deletions(-)

diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
index 84820bba408..777b6c8fa7c 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
@@ -67,6 +67,7 @@
     _resume_contract_error,  # type: ignore
     _resolve_ui_payload,  # type: ignore
     _stringify_tool_result,  # type: ignore
+    _track_tool_call_segment,  # type: ignore
 )
 from ._snapshots import (
     _DEFAULT_STATE_INPUT_KEY,
@@ -1529,6 +1530,54 @@ def _merge_resolved_approval_results_into_snapshot(
     snapshot_messages[:] = merged_messages
 
 
+def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict[str, Any]]) -> None:
+    """Append this turn's messages in the order the model emitted them.
+
+    Segments tracked during streaming record whether text came before or after
+    tool calls (issue #7223); tool results still follow the tool-call message
+    they answer. Anything not covered by segment tracking falls back to the
+    legacy grouping so no content is dropped.
+    """
+    text_message_ids = {segment["id"] for segment in flow.snapshot_segments if segment["kind"] == "text"}
+    # A tool-only opening message (TextMessageStart with no text segment) lets
+    # the first tool-call message reuse the streamed message id, matching the
+    # legacy layout; every other tool message gets a fresh id.
+    tool_open_id = flow.message_id if flow.message_id and flow.message_id not in text_message_ids else None
+    emitted_call_ids: set[str] = set()
+
+    for segment in flow.snapshot_segments:
+        kind = segment["kind"]
+        if kind == "text":
+            if segment["text"]:
+                all_messages.append({"id": segment["id"], "role": "assistant", "content": segment["text"]})
+        elif kind == "tool_calls":
+            calls = [
+                flow.tool_calls_by_id[call_id] for call_id in segment["call_ids"] if call_id in flow.tool_calls_by_id
+            ]
+            if not calls:
+                continue
+            message_id = tool_open_id or generate_event_id()
+            tool_open_id = None
+            all_messages.append({"id": message_id, "role": "assistant", "tool_calls": [call.copy() for call in calls]})
+            segment_call_ids = set(segment["call_ids"])
+            emitted_call_ids.update(segment_call_ids)
+            all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in segment_call_ids)
+        elif kind == "reasoning":
+            all_messages.extend(entry for entry in flow.reasoning_messages if entry.get("id") == segment["id"])
+
+    leftover_calls = [tc for tc in flow.pending_tool_calls if tc.get("id") not in emitted_call_ids]
+    if leftover_calls:
+        all_messages.append(
+            {
+                "id": tool_open_id or generate_event_id(),
+                "role": "assistant",
+                "tool_calls": [call.copy() for call in leftover_calls],
+            }
+        )
+        emitted_call_ids.update(call_id for call in leftover_calls if (call_id := call.get("id")) is not None)
+    all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") not in emitted_call_ids)
+
+
 def _build_messages_snapshot(
     flow: FlowState,
     snapshot_messages: list[dict[str, Any]],
@@ -1536,6 +1585,10 @@ def _build_messages_snapshot(
     """Build MessagesSnapshotEvent from current flow state."""
     all_messages = list(snapshot_messages)
 
+    if flow.snapshot_segments:
+        _append_segmented_snapshot_messages(flow, all_messages)
+        return MessagesSnapshotEvent(messages=all_messages)  # type: ignore[arg-type]
+
     # Add assistant message with tool calls only (no content)
     if flow.pending_tool_calls:
         tool_call_message_id = (
@@ -2239,6 +2292,7 @@ async def run_agent_stream(
                         flow.pending_tool_calls.append(confirm_entry)
                         flow.tool_calls_by_id[confirm_id] = confirm_entry
                         flow.tool_calls_ended.add(confirm_id)  # Mark as ended since we emit End event
+                        _track_tool_call_segment(flow, confirm_id)
                         flow.waiting_for_approval = True
                         flow.interrupts.append(
                             _approval_interrupt_for_function_call(
diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
index e59ea0941f4..b36007e528c 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
@@ -467,6 +467,10 @@ class FlowState:
     reasoning_messages: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
     accumulated_reasoning: dict[str, str] = field(default_factory=dict)  # pyright: ignore[reportUnknownVariableType]
     reasoning_message_id: str | None = None
+    # Ordered text/tool_calls/reasoning segment boundaries in the order the
+    # model emitted them, so MESSAGES_SNAPSHOT can preserve that order
+    # (issue #7223) instead of the fixed tool-calls -> results -> text layout.
+    snapshot_segments: list[dict[str, Any]] = field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
 
     def get_tool_name(self, call_id: str | None) -> str | None:
         """Get tool name by call ID."""
@@ -480,6 +484,37 @@ def get_pending_without_end(self) -> list[dict[str, Any]]:
         return [tc for tc in self.pending_tool_calls if tc.get("id") not in self.tool_calls_ended]
 
 
+def _open_text_segment(flow: FlowState, message_id: str) -> dict[str, Any]:
+    """Start a new snapshot text segment for a freshly opened text message."""
+    segment = {"kind": "text", "id": message_id, "text": ""}
+    flow.snapshot_segments.append(segment)
+    return segment
+
+
+def _text_segment_for(flow: FlowState, message_id: str) -> dict[str, Any] | None:
+    """Find the text segment belonging to the given message id."""
+    for segment in reversed(flow.snapshot_segments):
+        if segment["kind"] == "text" and segment["id"] == message_id:
+            return segment
+    return None
+
+
+def _track_tool_call_segment(flow: FlowState, tool_call_id: str) -> None:
+    """Record a tool call in the current tool segment, opening one if needed."""
+    if flow.snapshot_segments and flow.snapshot_segments[-1]["kind"] == "tool_calls":
+        flow.snapshot_segments[-1]["call_ids"].append(tool_call_id)
+    else:
+        flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": [tool_call_id]})
+
+
+def _track_reasoning_segment(flow: FlowState, message_id: str) -> None:
+    """Record where a reasoning block sits in the snapshot emission order."""
+    for segment in flow.snapshot_segments:
+        if segment["kind"] == "reasoning" and segment["id"] == message_id:
+            return
+    flow.snapshot_segments.append({"kind": "reasoning", "id": message_id})
+
+
 def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]:
     """Emit TextMessage events for TextContent."""
     if not content.text:
@@ -492,6 +527,7 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li
     if not flow.message_id:
         flow.message_id = generate_event_id()
         flow.accumulated_text = ""
+        _open_text_segment(flow, flow.message_id)
         events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant"))
     elif flow.accumulated_text and content.text == flow.accumulated_text:
         # Guard against full-message replay chunks that can appear after streaming deltas.
@@ -500,6 +536,9 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li
 
     events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
     flow.accumulated_text += content.text
+    segment = _text_segment_for(flow, flow.message_id)
+    if segment is not None:
+        segment["text"] = flow.accumulated_text
     return events
 
 
@@ -534,6 +573,7 @@ def _emit_tool_call(
         }
         flow.pending_tool_calls.append(tool_entry)
         flow.tool_calls_by_id[tool_call_id] = tool_entry
+        _track_tool_call_segment(flow, tool_call_id)
 
     elif tool_call_id:
         flow.tool_call_id = tool_call_id
@@ -809,6 +849,7 @@ def _emit_approval_request(
         flow.pending_tool_calls.append(confirm_entry)
         flow.tool_calls_by_id[confirm_id] = confirm_entry
         flow.tool_calls_ended.add(confirm_id)
+        _track_tool_call_segment(flow, confirm_id)
 
     flow.waiting_for_approval = True
     return events
@@ -870,6 +911,7 @@ def _emit_mcp_tool_call(content: Content, flow: FlowState) -> list[BaseEvent]:
     }
     flow.pending_tool_calls.append(tool_entry)
     flow.tool_calls_by_id[tool_call_id] = tool_entry
+    _track_tool_call_segment(flow, tool_call_id)
 
     return events
 
@@ -1020,6 +1062,7 @@ def _emit_text_reasoning(content: Content, flow: FlowState | None = None) -> lis
             if content.protected_data is not None:
                 reasoning_entry["encryptedValue"] = content.protected_data
             flow.reasoning_messages.append(reasoning_entry)
+            _track_reasoning_segment(flow, message_id)
         else:
             existing_entry["content"] = full_text
             if content.protected_data is not None:
diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py
index bd5853f439a..d8909d4a789 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_run.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_run.py
@@ -26,6 +26,7 @@
 from agent_framework_ag_ui._agent_run import (
     PendingApprovalEntry,
     PendingApprovalKey,
+    _build_messages_snapshot,
     _build_safe_metadata,
     _canonical_approval_resume_messages,
     _create_state_context_message,
@@ -407,6 +408,96 @@ def test_emit_text_skips_when_waiting_for_approval():
     assert len(events) == 0
 
 
+def _snapshot_kinds(event):
+    """Flatten a MessagesSnapshotEvent into (kind, message) pairs for order assertions."""
+    kinds = []
+    for message in event.messages:
+        dumped = message if isinstance(message, dict) else message.model_dump(exclude_none=True)
+        if dumped.get("role") == "assistant" and dumped.get("tool_calls"):
+            kinds.append(("tool_calls", dumped))
+        elif dumped.get("role") == "assistant":
+            kinds.append(("text", dumped))
+        elif dumped.get("role") == "tool":
+            kinds.append(("result", dumped))
+        elif dumped.get("role") == "reasoning":
+            kinds.append(("reasoning", dumped))
+    return kinds
+
+
+def test_snapshot_places_lead_in_text_before_tool_calls():
+    """Lead-in narration streamed before the tool calls stays in front of them."""
+    flow = FlowState()
+    _emit_text(Content.from_text("Let me research the docs first."), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["text", "tool_calls"]
+    assert kinds[0][1]["content"] == "Let me research the docs first."
+    assert kinds[0][1]["id"] == flow.message_id
+    assert kinds[1][1]["tool_calls"][0]["id"] == "call_1"
+
+
+def test_snapshot_preserves_stream_order_around_tool_results():
+    """A turn shaped text -> calls -> results -> text snapshots in that order."""
+    flow = FlowState()
+    _emit_text(Content.from_text("First, the plan."), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+    _emit_tool_result(Content.from_function_result(call_id="call_1", result="done"), flow)
+    _emit_text(Content.from_text("And the summary."), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["text", "tool_calls", "result", "text"]
+    assert kinds[0][1]["content"] == "First, the plan."
+    assert kinds[2][1].get("toolCallId", kinds[2][1].get("tool_call_id")) == "call_1"
+    assert kinds[3][1]["content"] == "And the summary."
+    assert kinds[3][1]["id"] != kinds[0][1]["id"]
+
+
+def test_snapshot_tool_only_message_reuses_stream_message_id():
+    """Tool-only turns keep the message id the stream opened with."""
+    flow = FlowState()
+    flow.message_id = "tool-only-msg"
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls"]
+    assert kinds[0][1]["id"] == "tool-only-msg"
+
+
+def test_snapshot_keeps_reasoning_in_emission_order():
+    """Reasoning blocks keep their streamed position instead of always trailing."""
+    flow = FlowState()
+    _emit_text_reasoning(Content.from_text("thinking out loud"), flow)
+    _emit_text(Content.from_text("Answer incoming."), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["reasoning", "text", "tool_calls"]
+
+
+def test_snapshot_without_segment_tracking_keeps_legacy_layout():
+    """Flows assembled without the emit helpers keep the old fixed grouping."""
+    flow = FlowState()
+    flow.accumulated_text = "late text"
+    tool_entry = {"id": "call_1", "type": "function", "function": {"name": "docs_fetch", "arguments": ""}}
+    flow.pending_tool_calls.append(tool_entry)
+    flow.tool_calls_by_id["call_1"] = tool_entry
+    flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_1", "content": "done"})
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result", "text"]
+
+
 def test_emit_text_skips_when_skip_text_flag():
     """Test _emit_text skips with skip_text flag."""
     flow = FlowState()
@@ -1619,7 +1710,7 @@ def test_reasoning_without_flow_does_not_error(self):
         assert isinstance(events[0], ReasoningStartEvent)
 
     def test_snapshot_reasoning_ordering(self):
-        """Reasoning messages appear after assistant text in snapshot."""
+        """Reasoning keeps its streamed position in the snapshot."""
         from agent_framework_ag_ui._agent_run import _build_messages_snapshot
 
         flow = FlowState()
@@ -1631,10 +1722,10 @@ def test_snapshot_reasoning_ordering(self):
 
         snapshot = _build_messages_snapshot(flow, [{"id": "u1", "role": "user", "content": "Hi"}])
 
-        # user -> assistant text -> reasoning
+        # reasoning streamed before the answer, so it snapshots before it too
         assert len(snapshot.messages) == 3
         roles = [_message_role(m) for m in snapshot.messages]
-        assert roles == ["user", "assistant", "reasoning"]
+        assert roles == ["user", "reasoning", "assistant"]
 
     def test_reasoning_accumulates_incremental_deltas(self):
         """Multiple reasoning deltas with the same id accumulate into one entry."""

From 24eb54fb39220734612715aea49086ec65f748d5 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Wed, 22 Jul 2026 14:42:42 +0800
Subject: [PATCH 2/3] Address moonbox3's review: cover the remaining snapshot
 gaps

- Preopened message ids (tool-only path) now open a text segment when
  the first text arrives, so their content can't drop out of the snapshot.
- A tool result closes the current tool-call segment, so
  call A -> result A -> call B snapshots as two pairs in stream order.
- emitted_call_ids only marks calls actually emitted, keeping stale
  segment ids eligible for the leftover fallback.
- The leftover path carries its tool results too instead of dropping them.
---
 .../ag-ui/agent_framework_ag_ui/_agent_run.py | 15 +++--
 .../agent_framework_ag_ui/_run_common.py      | 15 ++++-
 python/packages/ag-ui/tests/ag_ui/test_run.py | 66 +++++++++++++++++++
 3 files changed, 89 insertions(+), 7 deletions(-)

diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
index 777b6c8fa7c..a1fb39dba8e 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
@@ -1559,14 +1559,18 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict
             message_id = tool_open_id or generate_event_id()
             tool_open_id = None
             all_messages.append({"id": message_id, "role": "assistant", "tool_calls": [call.copy() for call in calls]})
-            segment_call_ids = set(segment["call_ids"])
-            emitted_call_ids.update(segment_call_ids)
-            all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in segment_call_ids)
+            # Only mark the calls we actually emitted; a stale segment id that
+            # never made it into tool_calls_by_id must stay eligible for the
+            # leftover path below rather than vanishing silently.
+            emitted_ids = {call["id"] for call in calls}
+            emitted_call_ids.update(emitted_ids)
+            all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in emitted_ids)
         elif kind == "reasoning":
             all_messages.extend(entry for entry in flow.reasoning_messages if entry.get("id") == segment["id"])
 
     leftover_calls = [tc for tc in flow.pending_tool_calls if tc.get("id") not in emitted_call_ids]
     if leftover_calls:
+        leftover_ids = {call.get("id") for call in leftover_calls if call.get("id") is not None}
         all_messages.append(
             {
                 "id": tool_open_id or generate_event_id(),
@@ -1574,7 +1578,10 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict
                 "tool_calls": [call.copy() for call in leftover_calls],
             }
         )
-        emitted_call_ids.update(call_id for call in leftover_calls if (call_id := call.get("id")) is not None)
+        # Their results ride along too; without this they would be marked
+        # emitted above and then excluded from the final append below.
+        all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") in leftover_ids)
+        emitted_call_ids.update(leftover_ids)
     all_messages.extend(result for result in flow.tool_results if result.get("toolCallId") not in emitted_call_ids)
 
 
diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
index b36007e528c..da7cbd7d36f 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
@@ -534,11 +534,16 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li
         logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id)
         return []
 
+    # The message may have been pre-opened by the tool-only path, which never
+    # goes through this function, so the first text arriving later has no
+    # segment yet; without one the snapshot would drop it.
+    segment = _text_segment_for(flow, flow.message_id)
+    if segment is None:
+        segment = _open_text_segment(flow, flow.message_id)
+
     events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
     flow.accumulated_text += content.text
-    segment = _text_segment_for(flow, flow.message_id)
-    if segment is not None:
-        segment["text"] = flow.accumulated_text
+    segment["text"] = flow.accumulated_text
     return events
 
 
@@ -722,6 +727,10 @@ def _emit_tool_result_common(
             "content": result_content,
         }
     )
+    # A result closes the current tool-call segment: a later call opens a new
+    # one, so `call A -> result A -> call B` snapshots as two call/result pairs
+    # in stream order instead of grouping B with A (moonbox3's replay concern).
+    flow.snapshot_segments.append({"kind": "tool_results"})
 
     if predictive_handler:
         predictive_handler.apply_pending_updates()
diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py
index d8909d4a789..217983c3d9e 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_run.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_run.py
@@ -498,6 +498,72 @@ def test_snapshot_without_segment_tracking_keeps_legacy_layout():
     assert [kind for kind, _ in kinds] == ["tool_calls", "result", "text"]
 
 
+def test_snapshot_includes_text_when_message_preopened_by_tool_only_path():
+    """Text that arrives after a tool-only preopen still lands in the snapshot."""
+    flow = FlowState()
+    # The tool-only detection in agent_run.py preopens message_id without
+    # going through _emit_text, so the first text has no segment yet.
+    flow.message_id = "preopened"
+    _emit_tool_call(Content.from_function_call(call_id="call_1", name="docs_fetch", arguments="{}"), flow)
+    _emit_text(Content.from_text("Let me check the docs."), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "text"]
+    assert kinds[1][1]["content"] == "Let me check the docs."
+    assert kinds[1][1]["id"] == "preopened"
+
+
+def test_snapshot_separates_calls_across_results():
+    """call A -> result A -> call B -> result B snapshots as two pairs, not grouped."""
+    flow = FlowState()
+    _emit_tool_call(Content.from_function_call(call_id="call_a", name="tool_a", arguments="{}"), flow)
+    _emit_tool_result(Content.from_function_result(call_id="call_a", result="a done"), flow)
+    _emit_tool_call(Content.from_function_call(call_id="call_b", name="tool_b", arguments="{}"), flow)
+    _emit_tool_result(Content.from_function_result(call_id="call_b", result="b done"), flow)
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result", "tool_calls", "result"]
+    assert kinds[0][1]["tool_calls"][0]["id"] == "call_a"
+    assert kinds[1][1].get("toolCallId", kinds[1][1].get("tool_call_id")) == "call_a"
+    assert kinds[2][1]["tool_calls"][0]["id"] == "call_b"
+    assert kinds[3][1].get("toolCallId", kinds[3][1].get("tool_call_id")) == "call_b"
+
+
+def test_snapshot_leftover_call_keeps_its_result():
+    """A pending call untracked by segments still appears with its result."""
+    flow = FlowState()
+    tool_entry = {"id": "call_x", "type": "function", "function": {"name": "tool_x", "arguments": ""}}
+    flow.pending_tool_calls.append(tool_entry)
+    flow.tool_calls_by_id["call_x"] = tool_entry
+    flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_x", "content": "done"})
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result"]
+    assert kinds[1][1].get("toolCallId", kinds[1][1].get("tool_call_id")) == "call_x"
+
+
+def test_snapshot_stale_segment_id_falls_back_to_leftover():
+    """A segment id missing from tool_calls_by_id stays eligible for the fallback."""
+    flow = FlowState()
+    flow.snapshot_segments.append({"kind": "tool_calls", "call_ids": ["call_x"]})
+    tool_entry = {"id": "call_x", "type": "function", "function": {"name": "tool_x", "arguments": ""}}
+    flow.pending_tool_calls.append(tool_entry)
+    # tool_calls_by_id intentionally lacks call_x, so the segment emits nothing
+    flow.tool_results.append({"id": "r1", "role": "tool", "toolCallId": "call_x", "content": "done"})
+
+    event = _build_messages_snapshot(flow, [])
+
+    kinds = _snapshot_kinds(event)
+    assert [kind for kind, _ in kinds] == ["tool_calls", "result"]
+    assert kinds[0][1]["tool_calls"][0]["id"] == "call_x"
+
+
 def test_emit_text_skips_when_skip_text_flag():
     """Test _emit_text skips with skip_text flag."""
     flow = FlowState()

From 6e5528d332ee8031887b2faee3f7ad97b72c2113 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Thu, 23 Jul 2026 18:56:45 +0800
Subject: [PATCH 3/3] Python: narrow leftover tool-call ids so pyright accepts
 the update

---
 python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
index a1fb39dba8e..267f434243f 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
@@ -1570,7 +1570,7 @@ def _append_segmented_snapshot_messages(flow: FlowState, all_messages: list[dict
 
     leftover_calls = [tc for tc in flow.pending_tool_calls if tc.get("id") not in emitted_call_ids]
     if leftover_calls:
-        leftover_ids = {call.get("id") for call in leftover_calls if call.get("id") is not None}
+        leftover_ids = {cid for call in leftover_calls if (cid := call.get("id")) is not None}
         all_messages.append(
             {
                 "id": tool_open_id or generate_event_id(),