fix(tui): stop empty-id assistant deltas from merging across turns (#4201)

_apply_assistant_delta matched AssistantDelta.id anywhere in the transcript,
which is correct for a genuine per-message id but not for an empty one.
runtime._as_str() coerces a missing/None chunk id to "", and that value is
shared by every id-less chunk from every turn, not just the current one.
Once one assistant row had id="", every later, unrelated AssistantDelta that
also carried id="" (e.g. from a provider that never stamps per-chunk ids)
matched that same stale row instead of starting a fresh one, silently
folding a second turn's answer backward into the first turn's bubble.

Route empty-id deltas to a dedicated path that tracks the current turn's
row by position (streaming_anonymous_row_index, reset on RunStarted/
RunEnded/ClearRows) instead of by id, mirroring the existing empty-id guards
in _apply_tool_started/_apply_tool_result but adapted for assistant text:
unlike a tool call, an id-less assistant delta still needs to be displayed,
so it starts a new row rather than being dropped. Multiple id-less chunks
legitimately arrive within one turn (per-token streaming), so they keep
coalescing into that row -- but only while it is still the transcript tail;
once a tool card is appended after it (the same way a genuine id naturally
changes across a tool round-trip), the next empty-id delta starts fresh
instead of reaching backward past the tool card.

Add regression coverage for the cross-turn merge, same-turn coalescing,
the tool-call-interleaved edge case, and non-interference with the
existing id-keyed path.
This commit is contained in:
Daoyuan Li 2026-07-15 07:27:00 -07:00 committed by GitHub
parent 3247f61750
commit 6ef0aa2aea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 252 additions and 9 deletions

View File

@ -140,6 +140,15 @@ class ViewState:
# Id of the message currently being generated this turn. Only this row renders # Id of the message currently being generated this turn. Only this row renders
# as plain text while streaming; everything else (history) stays Markdown. # as plain text while streaming; everything else (history) stays Markdown.
streaming_id: str | None = None streaming_id: str | None = None
# Row index of the *anonymous* (empty-id) assistant row receiving deltas this
# turn, if any. A genuine id is a reliable cross-chunk key (see
# `_apply_assistant_delta`'s whole-transcript id scan), but an empty id ("" —
# see `runtime._as_str`) is shared by every id-less chunk from every turn, so
# it cannot be matched the same way: scanning for `row.id == ""` would fold a
# brand new turn's text into whatever earlier turn's row happened to be
# id-less too. This index instead pins "this turn's" anonymous row by
# position, reset alongside `streaming_id` at the start/end of every turn.
streaming_anonymous_row_index: int | None = None
def initial_state(rows: tuple[Row, ...] = ()) -> ViewState: def initial_state(rows: tuple[Row, ...] = ()) -> ViewState:
@ -164,13 +173,14 @@ def reduce(state: ViewState, action: Action) -> ViewState:
if isinstance(action, RunStarted): if isinstance(action, RunStarted):
# New turn: no message is actively streaming yet (the client re-emits # New turn: no message is actively streaming yet (the client re-emits
# prior messages first; those must not be treated as the active one). # prior messages first; those must not be treated as the active one).
return replace(state, streaming=True, streaming_id=None) return replace(state, streaming=True, streaming_id=None, streaming_anonymous_row_index=None)
if isinstance(action, RunEnded): if isinstance(action, RunEnded):
return replace( return replace(
state, state,
streaming=False, streaming=False,
streaming_id=None, streaming_id=None,
streaming_anonymous_row_index=None,
usage=action.usage if action.usage is not None else state.usage, usage=action.usage if action.usage is not None else state.usage,
) )
@ -193,20 +203,28 @@ def reduce(state: ViewState, action: Action) -> ViewState:
return replace(state, title=action.title) return replace(state, title=action.title)
if isinstance(action, ClearRows): if isinstance(action, ClearRows):
return replace(state, rows=(), title=None, streaming_id=None) return replace(state, rows=(), title=None, streaming_id=None, streaming_anonymous_row_index=None)
return state return state
def _apply_assistant_delta(state: ViewState, action: AssistantDelta) -> ViewState: def _apply_assistant_delta(state: ViewState, action: AssistantDelta) -> ViewState:
"""Update the assistant row with this id (anywhere in the transcript), or """Update the assistant row for this delta, or start a new one.
start a new one.
On a thread with history, the client re-emits every prior message on each A genuine (non-empty) id is matched anywhere in the transcript, not just
new turn (its dedup is per-turn), and a re-emitted *older* message can arrive the most recent assistant row: on a thread with history, the client
after a newer one has started so we must match by id across the whole re-emits every prior message on each new turn (its dedup is per-turn), and
transcript, not just the most recent assistant row, or prior answers get a re-emitted *older* message can arrive after a newer one has started so
duplicated. matching only the tail row would duplicate prior answers.
An empty id ("" some providers/paths never stamp per-chunk ids, see
``runtime._as_str``) is NOT a reliable key for that same scan: unlike a
genuine id, it is shared by every id-less chunk from *every* turn, so
matching `row.id == ""` across the whole transcript would fold a brand
new turn's text into whatever earlier turn's row happened to be id-less
too silently vanishing the new turn's answer into a stale row. Empty-id
deltas are therefore routed to `_apply_assistant_delta_anonymous`, which
tracks "this turn's" row by position instead of by id.
Updates also merge by content rather than blindly concatenating, to absorb Updates also merge by content rather than blindly concatenating, to absorb
full re-sends / cumulative snapshots vs. genuine incremental deltas: full re-sends / cumulative snapshots vs. genuine incremental deltas:
@ -215,6 +233,8 @@ def _apply_assistant_delta(state: ViewState, action: AssistantDelta) -> ViewStat
* accumulated starts with new text -> stale/shorter re-send: keep * accumulated starts with new text -> stale/shorter re-send: keep
* otherwise -> a real delta: append * otherwise -> a real delta: append
""" """
if not action.id:
return _apply_assistant_delta_anonymous(state, action)
rows = list(state.rows) rows = list(state.rows)
for i, row in enumerate(rows): for i, row in enumerate(rows):
@ -234,6 +254,46 @@ def _apply_assistant_delta(state: ViewState, action: AssistantDelta) -> ViewStat
return _mark_streaming(_append(state, AssistantRow(text=action.text, id=action.id)), action.id) return _mark_streaming(_append(state, AssistantRow(text=action.text, id=action.id)), action.id)
def _apply_assistant_delta_anonymous(state: ViewState, action: AssistantDelta) -> ViewState:
"""Handle an ``AssistantDelta`` whose id is empty (see `_apply_assistant_delta`).
Multiple id-less chunks legitimately arrive for a single turn a provider
that never stamps per-chunk ids still streams token by token, e.g.
``"Hel"`` then ``"lo"`` so the first empty-id delta of a turn starts a
new row, and later empty-id deltas keep appending to that row (tracked by
``state.streaming_anonymous_row_index``, reset on every
``RunStarted``/``RunEnded``/``ClearRows``, not by id, so a later turn
always starts its own new row instead of matching the previous turn's
leftover id-less row the bug this split exists to avoid).
The tracked row is only reused while it is still the LAST row in the
transcript. A genuine id naturally changes across a tool round-trip
(LangGraph gives the post-tool continuation a new AIMessage id), which is
why an interleaved ``ToolStarted``/``ToolResult`` already starts a new row
in the id-keyed path (see `test_assistant_delta_with_new_id_after_tool_
creates_separate_row`). An empty id has no such natural signal it is
always ``""`` before and after the tool call so this function uses row
*position* as the substitute: once anything else has been appended (a
tool card, in practice), the anonymous row is no longer the tail, and the
next empty-id delta starts a fresh row rather than reaching backward past
the tool card into stale text.
"""
index = state.streaming_anonymous_row_index
if index is not None and index == len(state.rows) - 1:
row = state.rows[index]
if isinstance(row, AssistantRow) and not row.error:
# Same no-op / merge semantics as the id-keyed path above.
if row.text == action.text and len(action.text) > 1:
return state
rows = list(state.rows)
merged = _merge_stream_text(row.text, action.text)
rows[index] = replace(row, text=merged)
return _mark_streaming_anonymous(replace(state, rows=tuple(rows)), index)
new_state = _append(state, AssistantRow(text=action.text, id=action.id))
return _mark_streaming_anonymous(new_state, len(new_state.rows) - 1)
def _mark_streaming(state: ViewState, message_id: str) -> ViewState: def _mark_streaming(state: ViewState, message_id: str) -> ViewState:
"""Record the actively-streaming message id (only while a run is active).""" """Record the actively-streaming message id (only while a run is active)."""
if state.streaming: if state.streaming:
@ -241,6 +301,24 @@ def _mark_streaming(state: ViewState, message_id: str) -> ViewState:
return state return state
def _mark_streaming_anonymous(state: ViewState, index: int) -> ViewState:
"""Record the active turn's anonymous-row index (only while a run is active).
Deliberately leaves ``streaming_id`` at ``None`` rather than ``""``: unlike
a genuine id, ``""`` would be shared by every anonymous row across every
turn, so using it as the render layer's "is this the row actively
streaming" key (``render.render_transcript``) would flag every past
anonymous row as actively streaming too, the moment a new one starts. The
cost is purely cosmetic an anonymous row never gets the
raw-text-while-streaming treatment other rows get, only the id-less
fallback path is affected in exchange for not reintroducing a cross-turn
ambiguity into the render layer that this fix removes from ``rows``.
"""
if state.streaming:
return replace(state, streaming_id=None, streaming_anonymous_row_index=index)
return state
def _merge_stream_text(existing: str, incoming: str) -> str: def _merge_stream_text(existing: str, incoming: str) -> str:
if not existing: if not existing:
return incoming return incoming

View File

@ -143,3 +143,36 @@ def test_stream_actions_surfaces_exception_as_error_then_ends():
actions = list(stream_actions(_BoomClient(), "go")) actions = list(stream_actions(_BoomClient(), "go"))
assert any(isinstance(a, AssistantError) and "model down" in a.text for a in actions) assert any(isinstance(a, AssistantError) and "model down" in a.text for a in actions)
assert isinstance(actions[-1], RunEnded) assert isinstance(actions[-1], RunEnded)
def test_stream_actions_two_turns_with_none_ids_produce_separate_rows():
"""Some providers/paths never stamp per-chunk ids: the raw chunk carries
an explicit ``id: None``, which ``_as_str`` coerces to ``""``. Two
separate turns from such a provider must not fold into one row -- see
``_apply_assistant_delta_anonymous`` in view_state.py. Drives the real
translate()/stream_actions() bridge, not just the reducer directly."""
first_turn = _FakeClient(
[
StreamEvent(type="messages-tuple", data={"type": "ai", "content": "First turn answer.", "id": None}),
StreamEvent(type="end", data={"usage": None}),
]
)
second_turn = _FakeClient(
[
StreamEvent(type="messages-tuple", data={"type": "ai", "content": "Second turn answer.", "id": None}),
StreamEvent(type="end", data={"usage": None}),
]
)
state = initial_state()
for action in stream_actions(first_turn, "first question"):
state = reduce(state, action)
for action in stream_actions(second_turn, "second question"):
state = reduce(state, action)
assistants = [r for r in state.rows if r.kind == "assistant"]
# Pre-fix: both turns' AssistantDelta carry id="" and the second turn's
# text is folded into the first turn's row instead of starting a new one.
assert len(assistants) == 2
assert assistants[0].text == "First turn answer."
assert assistants[1].text == "Second turn answer."

View File

@ -241,3 +241,135 @@ def test_merge_stream_text_newline_split_across_chunks():
def test_merge_stream_text_genuine_delta_append(): def test_merge_stream_text_genuine_delta_append():
"""Normal deltas that don't overlap still append.""" """Normal deltas that don't overlap still append."""
assert _merge_stream_text("Hello ", "world") == "Hello world" assert _merge_stream_text("Hello ", "world") == "Hello world"
# ---------------------------------------------------------------------------
# Empty/missing-id assistant deltas: some providers/paths never stamp
# per-chunk ids (runtime._as_str coerces a missing id to ""). Matching by id
# like the normal path would fold EVERY id-less turn into whichever id-less
# row happened to exist first, since "" is shared across turns -- unlike a
# genuine id. These pin the fix: an empty id always keys off the CURRENT
# turn (never a stale row from an earlier turn), while still coalescing
# multiple id-less chunks that legitimately arrive within one turn.
# ---------------------------------------------------------------------------
def test_assistant_delta_empty_id_starts_new_row_per_turn_not_merged_with_prior_turn():
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="First turn answer."))
state = reduce(state, RunEnded())
state = reduce(state, UserSubmitted("second question"))
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="Second turn answer."))
state = reduce(state, RunEnded())
assistants = [r for r in state.rows if r.kind == "assistant"]
# Pre-fix: both turns share id="" so the second folds into the first via
# the whole-transcript id scan, losing "First turn answer." entirely.
assert len(assistants) == 2
assert assistants[0].text == "First turn answer."
assert assistants[1].text == "Second turn answer."
def test_assistant_delta_empty_id_coalesces_multiple_chunks_within_same_turn():
"""An id-less provider still streams token by token; chunks within ONE
turn must accumulate into a single row, not fragment into many."""
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="Hel"))
state = reduce(state, AssistantDelta(id="", text="lo"))
state = reduce(state, AssistantDelta(id="", text=" world"))
state = reduce(state, RunEnded())
assistants = [r for r in state.rows if r.kind == "assistant"]
assert len(assistants) == 1
assert assistants[0].text == "Hello world"
def test_assistant_delta_empty_id_starts_fresh_row_after_interleaved_tool_call():
"""An empty id has no signal to distinguish "same message, paused for a
tool call" from "a new message that happens to also be id-less" -- unlike
a genuine id, which naturally changes across a tool round-trip (a new
AIMessage gets a new id; see
test_assistant_delta_with_new_id_after_tool_creates_separate_row). Once a
tool card has been appended, the previous anonymous row is no longer the
transcript tail, so the next empty-id delta must start a NEW row rather
than reach backward past the tool card and silently prepend text that
arrived after the tool ran."""
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="Let me check. "))
state = reduce(state, ToolStarted(tool_call_id="t1", tool_name="bash", args={}))
state = reduce(state, ToolResult(tool_call_id="t1", content="ok", is_error=False))
state = reduce(state, AssistantDelta(id="", text="Done."))
state = reduce(state, RunEnded())
kinds = [r.kind for r in state.rows]
assert kinds == ["assistant", "tool", "assistant"]
assistants = [r for r in state.rows if r.kind == "assistant"]
assert [a.text for a in assistants] == ["Let me check. ", "Done."]
def test_assistant_delta_empty_id_coalesces_consecutive_chunks_before_a_tool_call():
"""Multiple id-less chunks with NOTHING interleaved (the realistic
per-token streaming case) still coalesce into one row up until a tool
card breaks the streak."""
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="Let me "))
state = reduce(state, AssistantDelta(id="", text="check. "))
state = reduce(state, ToolStarted(tool_call_id="t1", tool_name="bash", args={}))
state = reduce(state, ToolResult(tool_call_id="t1", content="ok", is_error=False))
state = reduce(state, RunEnded())
kinds = [r.kind for r in state.rows]
assert kinds == ["assistant", "tool"]
assistants = [r for r in state.rows if r.kind == "assistant"]
assert assistants[0].text == "Let me check. "
def test_assistant_delta_empty_id_does_not_disturb_legitimate_id_sequence():
"""A normal, non-empty id sequence must keep coalescing correctly even
after the transcript has already seen an earlier, unrelated empty-id
turn (proves the two code paths -- id-keyed vs. anonymous -- don't
interfere with each other)."""
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="anonymous turn"))
state = reduce(state, RunEnded())
state = reduce(state, UserSubmitted("question"))
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="m1", text="Hel"))
state = reduce(state, AssistantDelta(id="m1", text="lo"))
state = reduce(state, RunEnded())
assistants = [r for r in state.rows if r.kind == "assistant"]
assert [a.text for a in assistants] == ["anonymous turn", "Hello"]
def test_assistant_delta_empty_id_resend_within_turn_is_noop():
"""Same multi-char no-op re-send semantics apply to the anonymous path."""
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="Hey there!"))
state = reduce(state, AssistantDelta(id="", text="Hey there!"))
assistants = [r for r in state.rows if r.kind == "assistant"]
assert len(assistants) == 1
assert assistants[0].text == "Hey there!"
def test_clear_rows_resets_anonymous_streaming_index():
"""A stale anonymous-row index must not resurrect after ClearRows."""
state = initial_state()
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="before clear"))
state = reduce(state, ClearRows())
state = reduce(state, RunStarted())
state = reduce(state, AssistantDelta(id="", text="after clear"))
assistants = [r for r in state.rows if r.kind == "assistant"]
assert len(assistants) == 1
assert assistants[0].text == "after clear"