diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index d52d44b3798e7..62953533a684a 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -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: diff --git a/tests/gateway/test_stream_consumer.py b/tests/gateway/test_stream_consumer.py index f1b8ffa399cd6..58293b563a6ce 100644 --- a/tests/gateway/test_stream_consumer.py +++ b/tests/gateway/test_stream_consumer.py @@ -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):