diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index a08e169f2f9..25af5bff0fb 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -540,18 +540,26 @@ class GatewayStreamConsumer: if isinstance(self.adapter, _BasePlatformAdapter) else len ) + # Resolve native draft streaming before choosing the overflow budget. + # Rich-capable adapters (Telegram rich messages) can raise the budget + # above the legacy edit limit ONLY when native drafts are active: draft + # frames/final sends can use the rich endpoint, while editMessageText + # still has the 4096-ish legacy cap for progressive edits. If we let + # edit-based streaming accumulate to the rich cap, Telegram's adapter + # has to split the same full prefix on every edit and users see many + # duplicate "(1/2)" chunks before the tail arrives. + self._use_draft_streaming = self._resolve_draft_streaming() # Rich-capable adapters (Telegram rich messages) raise this above the - # legacy per-message limit so a reply that fits one rich send/draft - # isn't fragmented at 4096 while streaming. See _raw_message_limit. + # legacy per-message limit so a reply that fits one rich draft/final + # send isn't fragmented at 4096 while draft-streaming. See + # _raw_message_limit. _raw_limit = self._raw_message_limit() _safe_limit = max(500, _raw_limit - _len_fn(self.cfg.cursor) - 100) - # Resolve native draft streaming once per run. When enabled the - # consumer routes mid-stream frames through adapter.send_draft and - # leaves _message_id=None so the existing got_done path delivers the - # final answer as a regular sendMessage (drafts have no message_id - # to edit). - self._use_draft_streaming = self._resolve_draft_streaming() + # When native draft streaming is enabled the consumer routes mid-stream + # frames through adapter.send_draft and leaves _message_id=None so the + # existing got_done path delivers the final answer as a regular + # sendMessage (drafts have no message_id to edit). if self._use_draft_streaming: type(self)._draft_id_counter += 1 self._draft_id = type(self)._draft_id_counter @@ -1319,7 +1327,10 @@ class GatewayStreamConsumer: base = getattr(self.adapter, "MAX_MESSAGE_LENGTH", 4096) # isinstance gate: MagicMock adapters return mock objects (truthy, not # ints) for arbitrary attribute access — keep them on the base limit. - if isinstance(self.adapter, _BasePlatformAdapter): + # Also keep edit-based streaming on the legacy edit limit. A higher + # rich-message cap is safe only for native draft streaming because + # draft/final sends can use rich endpoints; progressive edits cannot. + if isinstance(self.adapter, _BasePlatformAdapter) and self._use_draft_streaming: try: cap = self.adapter.streaming_overflow_limit() except Exception as e: diff --git a/tests/gateway/test_stream_consumer.py b/tests/gateway/test_stream_consumer.py index cd49d3d7478..a53f58aea00 100644 --- a/tests/gateway/test_stream_consumer.py +++ b/tests/gateway/test_stream_consumer.py @@ -7,6 +7,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest from gateway.stream_consumer import GatewayStreamConsumer, StreamConsumerConfig +from gateway.config import Platform, PlatformConfig +from gateway.platforms.base import BasePlatformAdapter, SendResult # ── _clean_for_display unit tests ──────────────────────────────────────── @@ -2363,3 +2365,50 @@ class TestStripOrphanCloseTags: assert tag not in consumer._accumulated assert "trailing prose" in consumer._accumulated assert "more" in consumer._accumulated + + +class _RichCapEditAdapter(BasePlatformAdapter): + """Minimal adapter whose rich streaming cap exceeds its edit cap.""" + + MAX_MESSAGE_LENGTH = 4096 + + def __init__(self): + super().__init__(PlatformConfig(enabled=True), Platform.TELEGRAM) + + async def connect(self) -> bool: + return True + + async def disconnect(self) -> None: + return None + + async def send(self, chat_id: str, content: str, reply_to=None, metadata=None) -> SendResult: + return SendResult(success=True, message_id="msg_1") + + async def edit_message(self, chat_id: str, message_id: str, content: str, *, finalize: bool = False, metadata=None) -> SendResult: + return SendResult(success=True, message_id=message_id) + + async def get_chat_info(self, chat_id: str) -> dict: + return {} + + def streaming_overflow_limit(self): + return 32768 + + +class TestRichCapEditStreamingOverflow: + def test_edit_transport_ignores_rich_overflow_cap(self): + """Regression: Telegram rich cap must not drive edit-based streaming. + + If edit transport accumulates to the 32k rich cap, Telegram's 4096-char + edit path repeatedly split-delivers the same full prefix; users see many + duplicate ``(1/2)`` chunks before the tail. The rich cap is only safe + for native draft streaming, where frames/final delivery use rich send + endpoints instead of progressive edits. + """ + consumer = GatewayStreamConsumer(_RichCapEditAdapter(), "chat") + consumer._use_draft_streaming = False + assert consumer._raw_message_limit() == _RichCapEditAdapter.MAX_MESSAGE_LENGTH + + def test_draft_transport_can_use_rich_overflow_cap(self): + consumer = GatewayStreamConsumer(_RichCapEditAdapter(), "chat") + consumer._use_draft_streaming = True + assert consumer._raw_message_limit() == 32768