From 4aa499ff9f3fcc0c38ce61da46805a4dcc8f612e Mon Sep 17 00:00:00 2001 From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com> Date: Sat, 11 Jul 2026 10:56:57 +0530 Subject: [PATCH] fix(telegram): harden flood fallback recovery Keep empty-tail recovery scoped to the current stream segment and bound fallback flood retries. Preserve Telegram's server retry hint without blocking final delivery through a long cooldown. --- gateway/stream_consumer.py | 56 ++++++++++++---- plugins/platforms/telegram/adapter.py | 6 +- tests/gateway/test_telegram_final_delivery.py | 67 ++++++++++++++++++- 3 files changed, 113 insertions(+), 16 deletions(-) diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index d345d8f27fd..fca8bf43847 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -160,6 +160,10 @@ class GatewayStreamConsumer: # reply that was split across the platform's edit limit while streaming # doesn't leave stale fragments above the final message. self._preview_message_ids: "set[str]" = set() + # IDs from only the active text segment. A tool boundary preserves + # the run-wide set for fresh-final bookkeeping, but a failure recovery + # must never delete an earlier finalized preamble/commentary message. + self._segment_preview_message_ids: "set[str]" = set() self._already_sent = False self._edit_supported = True # Disabled when progressive edits are no longer usable self._last_edit_time = 0.0 @@ -173,6 +177,10 @@ class GatewayStreamConsumer: # Telegram overflow delivery. In that case the already-visible prefix # is intentional content, not a stale preview to delete. self._fallback_preserve_partial_messages = False + # Keep fallback recovery responsive. Telegram's adapter already bounds + # edit retries at five seconds; a final-delivery fallback must not hold + # the stream task through a longer flood cooldown before retrying. + self._max_fallback_flood_retry_seconds = 5.0 self._flood_strikes = 0 # Consecutive flood-control edit failures self._current_edit_interval = self.cfg.edit_interval # Adaptive backoff self._final_response_sent = False @@ -341,6 +349,7 @@ class GatewayStreamConsumer: self._fallback_final_send = False self._fallback_prefix = "" self._fallback_preserve_partial_messages = False + self._segment_preview_message_ids = set() # #29346: a tool/segment boundary means what we delivered was an interim # preamble, not the final answer — clear the flags so a premature setter # can't fool the gateway. Safe: got_done returns before any reset, and @@ -1073,15 +1082,15 @@ class GatewayStreamConsumer: ) if result.success: break - if attempt == 0 and self._is_flood_error(result): - delay = float(getattr(result, "retry_after", None) or 3.0) + retry_delay = self._fallback_flood_retry_delay(result) + if attempt == 0 and retry_delay is not None: logger.debug( "Flood control on fallback send, retrying in %.1fs", - delay, + retry_delay, ) - await asyncio.sleep(delay) + await asyncio.sleep(retry_delay) else: - break # non-flood error or second attempt failed + break # non-flood error, long flood wait, or second failure if not result or not result.success: if sent_any_chunk: @@ -1151,7 +1160,10 @@ class GatewayStreamConsumer: gateway can safely retry, and ``ambiguous`` when a timeout may have reached the platform already. """ - stale_ids = set(self._preview_message_ids) + # Tool/segment boundaries intentionally preserve the run-wide preview + # IDs for normal fresh-final cleanup. This recovery replaces only the + # active final segment, so never delete an earlier finalized preamble. + stale_ids = set(self._segment_preview_message_ids) if self._message_id and self._message_id != "__no_edit__": stale_ids.add(str(self._message_id)) @@ -1173,13 +1185,13 @@ class GatewayStreamConsumer: if getattr(result, "success", False): break - if attempt == 0 and self._is_flood_error(result): - delay = float(getattr(result, "retry_after", None) or 3.0) + retry_delay = self._fallback_flood_retry_delay(result) + if attempt == 0 and retry_delay is not None: logger.debug( "Flood control on empty fallback final send; retrying in %.1fs", - delay, + retry_delay, ) - await asyncio.sleep(delay) + await asyncio.sleep(retry_delay) continue return ( "ambiguous" @@ -1202,7 +1214,7 @@ class GatewayStreamConsumer: exc, ) - self._preview_message_ids = set() + self._segment_preview_message_ids = set() self._message_id = new_message_id or "__no_edit__" self._already_sent = True self._final_response_sent = True @@ -1222,6 +1234,22 @@ class GatewayStreamConsumer: name = result_or_exc.__class__.__name__.lower() return "timeout" in error or "timed out" in error or "timeout" in name + def _fallback_flood_retry_delay(self, result: Any) -> float | None: + """Return a bounded retry delay for a fallback send, if safe to retry.""" + if not self._is_flood_error(result): + return None + try: + delay = float(getattr(result, "retry_after", None) or 3.0) + except (TypeError, ValueError): + delay = 3.0 + if delay > self._max_fallback_flood_retry_seconds: + logger.debug( + "Flood control requests %.1fs; leaving final delivery to the gateway", + delay, + ) + return None + return max(0.0, delay) + def _is_flood_error(self, result) -> bool: """Check if a SendResult failure is due to flood control / rate limiting.""" err = getattr(result, "error", "") or "" @@ -1441,9 +1469,11 @@ class GatewayStreamConsumer: return base def _track_preview_id(self, message_id: Optional[str]) -> None: - """Record a real preview message id for fresh-final cleanup.""" + """Record a real preview message id for finalization cleanup.""" if message_id and message_id != "__no_edit__": - self._preview_message_ids.add(str(message_id)) + message_id = str(message_id) + self._preview_message_ids.add(message_id) + self._segment_preview_message_ids.add(message_id) def _track_preview_ids_from_result(self, result: Any) -> None: """Record every message id a send/edit result exposes: the primary id diff --git a/plugins/platforms/telegram/adapter.py b/plugins/platforms/telegram/adapter.py index bd4308feb12..5ea820761c4 100644 --- a/plugins/platforms/telegram/adapter.py +++ b/plugins/platforms/telegram/adapter.py @@ -4075,7 +4075,11 @@ class TelegramAdapter(BasePlatformAdapter): self.name, wait, ) if wait > 5.0: - return SendResult(success=False, error=f"flood_control:{wait}") + return SendResult( + success=False, + error=f"flood_control:{wait}", + retry_after=float(wait), + ) await asyncio.sleep(wait) try: await self._bot.edit_message_text( diff --git a/tests/gateway/test_telegram_final_delivery.py b/tests/gateway/test_telegram_final_delivery.py index e3a7bbb9871..31f7a230e56 100644 --- a/tests/gateway/test_telegram_final_delivery.py +++ b/tests/gateway/test_telegram_final_delivery.py @@ -5,8 +5,10 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from gateway.config import PlatformConfig from gateway.platforms.base import SendResult from gateway.stream_consumer import GatewayStreamConsumer, StreamConsumerConfig +from plugins.platforms.telegram.adapter import TelegramAdapter def _adapter() -> MagicMock: @@ -145,7 +147,7 @@ async def test_empty_tail_commit_honors_retry_after(monkeypatch): SendResult( success=False, error="Flood control exceeded", - retry_after=7.0, + retry_after=3.0, ), SendResult(success=True, message_id="final-1"), ] @@ -160,10 +162,71 @@ async def test_empty_tail_commit_honors_retry_after(monkeypatch): await consumer._send_fallback_final("Final answer") assert adapter.send.await_count == 2 - sleep.assert_awaited_once_with(7.0) + sleep.assert_awaited_once_with(3.0) assert consumer.final_content_delivered is True +@pytest.mark.asyncio +async def test_empty_tail_recovery_keeps_prior_segment_messages(): + """Recovery replaces only its current preview, not earlier preambles.""" + adapter = _adapter() + adapter.send.return_value = SendResult(success=True, message_id="final-1") + consumer = GatewayStreamConsumer(adapter, "chat-1") + + consumer._track_preview_id("preamble-1") + consumer._reset_segment_state() + consumer._track_preview_id("preview-1") + consumer._message_id = "preview-1" + consumer._last_sent_text = "Final answer" + consumer._fallback_final_send = True + + await consumer._send_fallback_final("Final answer") + + adapter.delete_message.assert_awaited_once_with("chat-1", "preview-1") + assert "preamble-1" in consumer._preview_message_ids + + +@pytest.mark.asyncio +async def test_empty_tail_commit_skips_long_flood_retry(monkeypatch): + adapter = _adapter() + adapter.send.return_value = SendResult( + success=False, + error="flood_control:30.0", + retry_after=30.0, + ) + sleep = AsyncMock() + monkeypatch.setattr("gateway.stream_consumer.asyncio.sleep", sleep) + + consumer = GatewayStreamConsumer(adapter, "chat-1") + consumer._message_id = "preview-1" + consumer._last_sent_text = "Final answer" + consumer._fallback_final_send = True + + await consumer._send_fallback_final("Final answer") + + adapter.send.assert_awaited_once() + sleep.assert_not_awaited() + assert consumer.final_response_sent is False + assert consumer.final_content_delivered is False + + +@pytest.mark.asyncio +async def test_telegram_long_flood_result_keeps_retry_after(): + """The real adapter contract preserves the server delay for consumers.""" + class FloodError(Exception): + retry_after = 30.0 + + adapter = TelegramAdapter(PlatformConfig(enabled=True, token="test-token")) + adapter._bot = MagicMock() + adapter._bot.edit_message_text = AsyncMock(side_effect=FloodError("Retry after 30")) + + result = await adapter.edit_message("123", "456", "Final answer", finalize=False) + + assert result.success is False + assert result.error == "flood_control:30.0" + assert result.retry_after == 30.0 + + @pytest.mark.asyncio async def test_ambiguous_empty_tail_timeout_preserves_duplicate_suppression(): adapter = _adapter()