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.
This commit is contained in:
kshitijk4poor 2026-07-11 10:56:57 +05:30 committed by kshitij
parent 04898631cb
commit 4aa499ff9f
3 changed files with 113 additions and 16 deletions

View file

@ -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

View file

@ -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(

View file

@ -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()