fix(gateway): flush undelivered tail before segment reset to preserve streamed text (#8124)

When a streaming edit fails mid-stream (flood control, transport error)
and a tool boundary arrives before the fallback threshold is reached,
the pre-boundary tail in `_accumulated` was silently discarded by
`_reset_segment_state`. The user saw a frozen partial message and
missing words on the other side of the tool call.

Flush the undelivered tail as a continuation message before the reset,
computed relative to the last successfully-delivered prefix so we don't
duplicate content the user already saw.
This commit is contained in:
konsisumer 2026-04-18 07:10:05 +02:00 committed by Teknium
parent e017131403
commit 1d1e1277e4
2 changed files with 105 additions and 2 deletions

View file

@ -502,11 +502,13 @@ class TestSegmentBreakOnToolBoundary:
@pytest.mark.asyncio
async def test_segment_break_clears_failed_edit_fallback_state(self):
"""A tool boundary after edit failure must not duplicate the next segment."""
"""A tool boundary after edit failure must flush the undelivered tail
without duplicating the prefix the user already saw (#8124)."""
adapter = MagicMock()
send_results = [
SimpleNamespace(success=True, message_id="msg_1"),
SimpleNamespace(success=True, message_id="msg_2"),
SimpleNamespace(success=True, message_id="msg_3"),
]
adapter.send = AsyncMock(side_effect=send_results)
adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=False, error="flood_control:6"))
@ -526,7 +528,60 @@ class TestSegmentBreakOnToolBoundary:
await task
sent_texts = [call[1]["content"] for call in adapter.send.call_args_list]
assert sent_texts == ["Hello ▉", "Next segment"]
# The undelivered "world" tail must reach the user, and the next
# segment must not duplicate "Hello" that was already visible.
assert sent_texts == ["Hello ▉", "world", "Next segment"]
@pytest.mark.asyncio
async def test_segment_break_after_mid_stream_edit_failure_preserves_tail(self):
"""Regression for #8124: when an earlier edit succeeded but later edits
fail (persistent flood control) and a tool boundary arrives before the
fallback threshold is reached, the pre-boundary tail must still be
delivered not silently dropped by the segment reset."""
adapter = MagicMock()
# msg_1 for the initial partial, msg_2 for the flushed tail,
# msg_3 for the post-boundary segment.
send_results = [
SimpleNamespace(success=True, message_id="msg_1"),
SimpleNamespace(success=True, message_id="msg_2"),
SimpleNamespace(success=True, message_id="msg_3"),
]
adapter.send = AsyncMock(side_effect=send_results)
# First two edits succeed, everything after fails with flood control
# — simulating Telegram's "edit once then get rate-limited" pattern.
edit_results = [
SimpleNamespace(success=True), # "Hello world ▉" — succeeds
SimpleNamespace(success=False, error="flood_control:6.0"), # "Hello world more ▉" — flood triggered
SimpleNamespace(success=False, error="flood_control:6.0"), # finalize edit at segment break
SimpleNamespace(success=False, error="flood_control:6.0"), # cursor-strip attempt
]
adapter.edit_message = AsyncMock(side_effect=edit_results + [edit_results[-1]] * 10)
adapter.MAX_MESSAGE_LENGTH = 4096
config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5, cursor="")
consumer = GatewayStreamConsumer(adapter, "chat_123", config)
consumer.on_delta("Hello")
task = asyncio.create_task(consumer.run())
await asyncio.sleep(0.08)
consumer.on_delta(" world")
await asyncio.sleep(0.08)
consumer.on_delta(" more")
await asyncio.sleep(0.08)
consumer.on_delta(None) # tool boundary
consumer.on_delta("Here is the tool result.")
consumer.finish()
await task
sent_texts = [call[1]["content"] for call in adapter.send.call_args_list]
# "more" must have been delivered, not dropped.
all_text = " ".join(sent_texts)
assert "more" in all_text, (
f"Pre-boundary tail 'more' was silently dropped: sends={sent_texts}"
)
# Post-boundary text must also reach the user.
assert "Here is the tool result." in all_text
@pytest.mark.asyncio
async def test_no_message_id_enters_fallback_mode(self):