diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index 62953533a684..d24812c73740 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -814,36 +814,49 @@ class GatewayStreamConsumer: # 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._initial_reply_to_id - while _len_fn(self._accumulated) > _safe_limit: - _cp_budget = _custom_unit_to_cp( + chunks = self._truncate_for_stream( + self._accumulated, _safe_limit, _len_fn, + ) + if len(chunks) <= 1: + # A malformed/legacy adapter result must not leave + # this overflow branch with an unsplittable payload. + chunks = self._split_text_chunks( 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] + chunks_delivered = False + reply_to = self._initial_reply_to_id + all_heads_delivered = len(chunks) > 1 + for chunk in chunks[:-1]: new_id = await self._send_new_chunk( chunk, reply_to, final=got_done, ) if new_id is None or new_id == reply_to: - # Failed to deliver the sealed head; keep the + # Failed to deliver a sealed head; keep the # full accumulated text intact so the gateway's # fallback path can still deliver it completely. + all_heads_delivered = False 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 + + if all_heads_delivered: + self._accumulated = chunks[-1] + # The head chunks are 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 = "" + else: + # A prior head may have landed before a later head + # failed. Do not edit that sealed message with the + # unsplit full payload; let the fallback path retry. + self._message_id = None + self._message_created_ts = None + self._last_sent_text = "" self._last_edit_time = time.monotonic() if got_done: @@ -1164,7 +1177,8 @@ class GatewayStreamConsumer: @staticmethod def _split_text_chunks( - text: str, limit: int, + text: str, + limit: int, len_fn: "Callable[[str], int]" = len, ) -> list[str]: """Split text into reasonably sized chunks for fallback sends.""" @@ -1183,6 +1197,33 @@ class GatewayStreamConsumer: chunks.append(remaining) return chunks + def _truncate_for_stream( + self, + text: str, + limit: int, + len_fn: "Callable[[str], int]", + ) -> list[str]: + """Use the adapter's canonical splitter for streaming overflow. + + Platform adapters may add word-boundary, code-fence, table, or + platform-specific formatting rules. The consumer must not replace + those rules with newline-only slicing. Non-base test doubles and + legacy adapters retain the historical two-argument call shape. + """ + truncate = getattr(self.adapter, "truncate_message", None) + if not callable(truncate): + return self._split_text_chunks(text, limit, len_fn) + + if isinstance(self.adapter, _BasePlatformAdapter): + chunks = truncate(text, limit, len_fn=len_fn) + else: + chunks = truncate(text, limit) + if not isinstance(chunks, (list, tuple)) or not all( + isinstance(chunk, str) for chunk in chunks + ): + return self._split_text_chunks(text, limit, len_fn) + return list(chunks) + async def _send_fallback_final(self, text: str) -> None: """Send the final continuation after streaming edits stop working. diff --git a/tests/gateway/test_stream_consumer.py b/tests/gateway/test_stream_consumer.py index 58293b563a6c..db28bdcf04b5 100644 --- a/tests/gateway/test_stream_consumer.py +++ b/tests/gateway/test_stream_consumer.py @@ -1229,6 +1229,58 @@ class TestInitialOverflowRollingEdit: ) assert consumer.final_response_sent is True + @pytest.mark.asyncio + async def test_initial_overflow_uses_adapter_fence_aware_split(self): + """Initial rolling sends must preserve the adapter's fence contract.""" + adapter = TestUtf16OverflowDetection()._make_telegram_like_adapter() + from gateway.platforms.base import utf16_len + + msg_ids = iter(["msg_1", "msg_2", "msg_3"]) + 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_3"), + ) + raw_limit = 700 + setattr(adapter, "MAX_MESSAGE_LENGTH", raw_limit) + splitter = MagicMock(side_effect=adapter.truncate_message) + adapter.truncate_message = splitter + + config = StreamConsumerConfig( + edit_interval=0.01, + buffer_threshold=5, + cursor=" ▉", + ) + consumer = GatewayStreamConsumer(adapter, "chat_fenced", config) + fenced = "```python\n" + ("print('x')\n" * 100) + "```" + safe_limit = raw_limit - utf16_len(config.cursor) - 100 + expected_chunks = adapter.truncate_message( + fenced, safe_limit, len_fn=adapter.message_len_fn, + ) + splitter.reset_mock() + + consumer.on_delta(fenced) + task = asyncio.create_task(consumer.run()) + await asyncio.sleep(0.08) + consumer.on_delta("\nTail after the fenced stream.") + await asyncio.sleep(0.08) + consumer.finish() + await task + + sent_texts = [call.kwargs["content"] for call in adapter.send.call_args_list] + edited_texts = [call.kwargs["content"] for call in adapter.edit_message.call_args_list] + assert splitter.call_count >= 1 + assert all(text.count("```") % 2 == 0 for text in sent_texts + edited_texts) + assert len(sent_texts) == len(expected_chunks) + assert sent_texts[:-1] == expected_chunks[:-1] + assert sent_texts[-1].startswith(expected_chunks[-1]) + assert any("Tail after the fenced stream." in text for text in edited_texts) + assert all(utf16_len(text) <= safe_limit for text in sent_texts) + class TestEditOverflowSplitAndDeliver: """When edit_message split-and-delivers an oversized payload across the