fix(gateway): keep overflow stream chunks editable

This commit is contained in:
2001y 2026-06-14 12:28:16 +09:00 committed by Teknium
parent b932997f44
commit b86ca7cae5
2 changed files with 103 additions and 37 deletions

View File

@ -808,39 +808,63 @@ class GatewayStreamConsumer:
and self._message_id is None
):
# No existing message to edit (first message or after a
# segment break). Use truncate_message — the same
# helper the non-streaming path uses — to split with
# proper word/code-fence boundaries and chunk
# indicators like "(1/2)".
chunks = self.adapter.truncate_message(
self._accumulated, _safe_limit, len_fn=_len_fn,
)
# segment break). Seal only the overflowing head chunks
# as fixed messages, then keep the trailing chunk in
# _accumulated so the normal send/edit path below makes
# it the active preview. That lets chunk 2, 3, ... keep
# updating in-place as later streamed deltas arrive
# instead of posting every split as an immutable message.
chunks_delivered = False
reply_to = self._message_id or self._initial_reply_to_id
for chunk in chunks:
reply_to = self._initial_reply_to_id
while _len_fn(self._accumulated) > _safe_limit:
_cp_budget = _custom_unit_to_cp(
self._accumulated, _safe_limit, _len_fn,
)
split_at = self._accumulated.rfind("\n", 0, _cp_budget)
if split_at < _cp_budget // 2:
split_at = _cp_budget
chunk = self._accumulated[:split_at]
new_id = await self._send_new_chunk(
chunk,
reply_to,
final=got_done,
)
if new_id is not None and new_id != reply_to:
chunks_delivered = True
self._accumulated = ""
self._last_sent_text = ""
if new_id is None or new_id == reply_to:
# Failed to deliver the sealed head; keep the
# full accumulated text intact so the gateway's
# fallback path can still deliver it completely.
chunks_delivered = False
break
chunks_delivered = True
reply_to = new_id
self._accumulated = self._accumulated[split_at:].lstrip("\n")
# The head chunk is sealed. Clear the edit target
# so the remaining tail is sent as a fresh active
# chunk, then edited by subsequent deltas.
self._message_id = None
self._message_created_ts = None
self._last_sent_text = ""
self._last_edit_time = time.monotonic()
if got_done:
# Only claim final delivery if THESE chunks actually
# landed. ``_already_sent`` may be True from prior
# tool-progress edits or fallback-mode promotion (#10748)
# — that doesn't mean the final answer reached the user.
self._final_response_sent = chunks_delivered
if chunks_delivered:
tail_delivered = True
if self._accumulated:
tail_delivered = await self._send_or_edit(
self._accumulated, finalize=True,
)
# Only claim final delivery if the sealed chunks and
# final tail actually landed. ``_already_sent`` may
# be True from prior progress/fallback state (#10748).
self._final_response_sent = chunks_delivered and tail_delivered
if self._final_response_sent:
self._final_content_delivered = True
return
if got_segment_break:
self._message_id = None
self._fallback_final_send = False
self._fallback_prefix = ""
if not self._accumulated:
continue
# This iteration consumed a _FLUSH barrier and delivered
# the buffered prose via the chunk loop above, then takes
@ -861,8 +885,8 @@ class GatewayStreamConsumer:
self._accumulated, _safe_limit, _len_fn,
)
split_at = self._accumulated.rfind("\n", 0, _cp_budget)
if split_at < _safe_limit // 2:
split_at = _safe_limit
if split_at < _cp_budget // 2:
split_at = _cp_budget
chunk = self._accumulated[:split_at]
# finalize=True so the adapter applies platform-specific
# rich-text markup (e.g. Telegram MarkdownV2). This
@ -1151,8 +1175,8 @@ class GatewayStreamConsumer:
while len_fn(remaining) > limit:
_cp_budget = _custom_unit_to_cp(remaining, limit, len_fn)
split_at = remaining.rfind("\n", 0, _cp_budget)
if split_at < limit // 2:
split_at = limit
if split_at < _cp_budget // 2:
split_at = _cp_budget
chunks.append(remaining[:split_at])
remaining = remaining[split_at:].lstrip("\n")
if remaining:

View File

@ -1183,6 +1183,53 @@ class TestFinalContentDeliveredGuard:
)
class TestInitialOverflowRollingEdit:
@pytest.mark.asyncio
async def test_initial_overflow_keeps_last_chunk_as_edit_target(self):
"""When the first visible flush already overflows, only sealed head
chunks should be posted as fixed messages. The trailing chunk must
remain the active edit target so later streamed deltas update that
second message instead of overwriting or posting a new one."""
adapter = MagicMock()
msg_ids = iter(["msg_1", "msg_2"])
adapter.send = AsyncMock(
side_effect=lambda **kw: SimpleNamespace(
success=True,
message_id=next(msg_ids),
)
)
adapter.edit_message = AsyncMock(
return_value=SimpleNamespace(success=True, message_id="msg_2"),
)
adapter.MAX_MESSAGE_LENGTH = 700
config = StreamConsumerConfig(
edit_interval=0.01,
buffer_threshold=5,
cursor="",
)
consumer = GatewayStreamConsumer(adapter, "chat_123", config)
head = "A" * 650
tail = "B" * 25
consumer.on_delta(head)
task = asyncio.create_task(consumer.run())
await asyncio.sleep(0.08)
consumer.on_delta(tail)
await asyncio.sleep(0.08)
consumer.finish()
await task
assert adapter.send.call_count == 2
assert adapter.edit_message.call_count >= 1
edited_texts = [call.kwargs["content"] for call in adapter.edit_message.call_args_list]
assert any("A" * 20 in text and tail in text for text in edited_texts), (
"the second overflow chunk should be edited with its existing tail "
"plus later deltas, not overwritten by only the later delta"
)
assert consumer.final_response_sent is True
class TestEditOverflowSplitAndDeliver:
"""When edit_message split-and-delivers an oversized payload across the
original message + N continuations (Telegram >4096 UTF-16), the consumer
@ -1985,11 +2032,6 @@ class TestUtf16OverflowDetection:
adapter.edit_message = AsyncMock(
return_value=SimpleNamespace(success=True),
)
# truncate_message: emit two halves so we can assert the split fired
adapter.truncate_message = MagicMock(
side_effect=lambda text, limit, **kw: [text[:len(text)//2], text[len(text)//2:]],
)
config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5)
consumer = GatewayStreamConsumer(adapter, "chat_123", config)
@ -2010,17 +2052,17 @@ class TestUtf16OverflowDetection:
consumer.finish()
await task
# The fix: stream consumer detects UTF-16 overflow and calls
# truncate_message to split. Without the fix, len() would return
# 2200 (under 4096) and no split would fire — Telegram would then
# reject the send or render \x00 artifacts.
adapter.truncate_message.assert_called(), (
# The fix: stream consumer detects UTF-16 overflow using the adapter's
# length function. Without that, len() would return 2200 (under the
# limit) and Hermes would attempt a single over-limit Telegram send.
sent_texts = [call.kwargs["content"] for call in adapter.send.call_args_list]
assert len(sent_texts) == 2, (
"UTF-16 overflow not detected — emoji text bypassed split path"
)
# truncate_message must have been called with len_fn=utf16_len
call_kwargs = adapter.truncate_message.call_args[1]
assert call_kwargs.get("len_fn") is utf16_len, (
f"truncate_message called without utf16_len: {call_kwargs}"
max_units = 4096
assert all(utf16_len(text) <= max_units for text in sent_texts), (
f"split chunks still exceed Telegram UTF-16 limit: "
f"{[utf16_len(text) for text in sent_texts]}"
)
def test_codepoint_only_adapter_falls_back_to_len(self):