fix(telegram): cancel delayed deliveries on disconnect

Buffered text/photo/media-group flushes and the polling-error recovery
task sit behind an asyncio.sleep(). On disconnect they kept running and
dispatched handle_message() into a torn-down session, producing stale or
duplicate deliveries. disconnect() only cancelled media-group and photo
batch tasks — text batches and the polling-error task leaked.

Set a _drop_delayed_deliveries flag from _mark_disconnected/_set_fatal_error
(cleared by _mark_connected) and check it in all enqueue+flush paths so a
flush that wins the race against teardown drops instead of dispatching.
_cancel_pending_delivery_tasks() now cancels+clears all four task maps,
skipping the current task. Media-group flush finally-block guarded so a
cancelled stale flush cannot erase a replacement task handle.
This commit is contained in:
CRWuTJ 2026-06-30 17:06:20 -07:00 committed by Teknium
parent 7de485703b
commit 8ad15ff7dd
2 changed files with 263 additions and 17 deletions

View file

@ -412,6 +412,7 @@ class TelegramAdapter(BasePlatformAdapter):
)
self._pending_text_batches: Dict[str, MessageEvent] = {}
self._pending_text_batch_tasks: Dict[str, asyncio.Task] = {}
self._drop_delayed_deliveries = False
self._polling_error_task: Optional[asyncio.Task] = None
self._polling_conflict_count: int = 0
self._polling_network_error_count: int = 0
@ -499,6 +500,27 @@ class TelegramAdapter(BasePlatformAdapter):
# same key edit the same message instead of appending new ones (#30045).
self._status_message_ids: Dict[tuple, str] = {}
def _mark_connected(self) -> None:
self._drop_delayed_deliveries = False
super()._mark_connected()
def _mark_disconnected(self) -> None:
self._drop_delayed_deliveries = True
super()._mark_disconnected()
def _set_fatal_error(self, code: str, message: str, *, retryable: bool) -> None:
self._drop_delayed_deliveries = True
super()._set_fatal_error(code, message, retryable=retryable)
def _should_drop_delayed_delivery(self) -> bool:
"""True once teardown/fatal-error started — delayed flushes must drop.
Buffered text/photo/media-group flushes sit behind an asyncio.sleep().
If disconnect wins the race, dispatching them spawns an agent on a
torn-down session, producing stale/duplicate deliveries.
"""
return bool(getattr(self, "_drop_delayed_deliveries", False))
def _notification_kwargs(
self, metadata: Optional[Dict[str, Any]]
) -> Dict[str, Any]:
@ -2915,8 +2937,60 @@ class TelegramAdapter(BasePlatformAdapter):
self.name, text, e,
)
async def _cancel_pending_delivery_tasks(self) -> None:
"""Cancel every delayed-delivery task family before disconnect completes.
Covers media-group, photo-batch and text-batch flush tasks plus the
polling-error recovery task. Each sits behind an ``asyncio.sleep()``;
if teardown leaves them running they dispatch ``handle_message`` into a
torn-down session. Skips the current task so the coroutine driving
teardown does not cancel itself.
"""
current_task = asyncio.current_task()
pending_tasks: list[asyncio.Task] = []
awaitable_tasks: list[asyncio.Task] = []
seen: set[int] = set()
def collect(task: Optional[asyncio.Task]) -> None:
if not task or task.done() or task is current_task:
return
marker = id(task)
if marker in seen:
return
seen.add(marker)
pending_tasks.append(task)
if asyncio.isfuture(task) or asyncio.iscoroutine(task):
awaitable_tasks.append(task)
for task in list(self._media_group_tasks.values()):
collect(task)
for task in list(self._pending_photo_batch_tasks.values()):
collect(task)
for task in list(self._pending_text_batch_tasks.values()):
collect(task)
collect(self._polling_error_task)
for task in pending_tasks:
task.cancel()
if awaitable_tasks:
await asyncio.gather(*awaitable_tasks, return_exceptions=True)
self._media_group_tasks.clear()
self._media_group_events.clear()
self._pending_photo_batch_tasks.clear()
self._pending_photo_batches.clear()
self._pending_text_batch_tasks.clear()
self._pending_text_batches.clear()
if self._polling_error_task is not current_task:
self._polling_error_task = None
async def disconnect(self) -> None:
"""Stop polling/webhook, cancel pending album flushes, and disconnect."""
"""Stop polling/webhook, cancel pending delayed deliveries, and disconnect."""
# Mark disconnected first so the drop guard short-circuits any flush
# that wins the race against teardown and prevents new delayed tasks
# from being scheduled by late update handlers.
self._mark_disconnected()
# Cancel the heartbeat before tearing down the app so the probe task
# cannot fire get_me() into a half-shutdown bot client.
if self._polling_heartbeat_task and not self._polling_heartbeat_task.done():
@ -2937,13 +3011,7 @@ class TelegramAdapter(BasePlatformAdapter):
except Exception:
pass
pending_media_group_tasks = list(self._media_group_tasks.values())
for task in pending_media_group_tasks:
task.cancel()
if pending_media_group_tasks:
await asyncio.gather(*pending_media_group_tasks, return_exceptions=True)
self._media_group_tasks.clear()
self._media_group_events.clear()
await self._cancel_pending_delivery_tasks()
if self._app:
try:
@ -2957,13 +3025,6 @@ class TelegramAdapter(BasePlatformAdapter):
logger.warning("[%s] Error during Telegram disconnect: %s", self.name, e, exc_info=True)
self._release_platform_lock()
for task in self._pending_photo_batch_tasks.values():
if task and not task.done():
task.cancel()
self._pending_photo_batch_tasks.clear()
self._pending_photo_batches.clear()
self._mark_disconnected()
self._app = None
self._bot = None
logger.info("[%s] Disconnected from Telegram", self.name)
@ -6874,6 +6935,10 @@ class TelegramAdapter(BasePlatformAdapter):
concatenates them and waits for a short quiet period before
dispatching the combined message.
"""
if self._should_drop_delayed_delivery():
logger.debug("[Telegram] Dropping text batch enqueue after disconnect started")
return
key = self._text_batch_key(event)
existing = self._pending_text_batches.get(key)
chunk_len = len(event.text or "")
@ -6933,6 +6998,9 @@ class TelegramAdapter(BasePlatformAdapter):
event = self._pending_text_batches.pop(key, None)
if not event:
return
if self._should_drop_delayed_delivery():
logger.debug("[Telegram] Dropping text batch flush after disconnect started")
return
logger.info(
"[Telegram] Flushing text batch %s (%d chars)",
key, len(event.text or ""),
@ -6967,6 +7035,9 @@ class TelegramAdapter(BasePlatformAdapter):
event = self._pending_photo_batches.pop(batch_key, None)
if not event:
return
if self._should_drop_delayed_delivery():
logger.debug("[Telegram] Dropping photo batch flush after disconnect started")
return
logger.info("[Telegram] Flushing photo batch %s with %d image(s)", batch_key, len(event.media_urls))
await self.handle_message(event)
finally:
@ -6975,6 +7046,10 @@ class TelegramAdapter(BasePlatformAdapter):
def _enqueue_photo_event(self, batch_key: str, event: MessageEvent) -> None:
"""Merge photo events into a pending batch and schedule flush."""
if self._should_drop_delayed_delivery():
logger.debug("[Telegram] Dropping photo batch enqueue after disconnect started")
return
existing = self._pending_photo_batches.get(batch_key)
if existing is None:
self._pending_photo_batches[batch_key] = event
@ -7279,6 +7354,10 @@ class TelegramAdapter(BasePlatformAdapter):
new user message and interrupts the first. We debounce briefly and merge the
attachments into a single MessageEvent.
"""
if self._should_drop_delayed_delivery():
logger.debug("[Telegram] Dropping media group enqueue after disconnect started")
return
existing = self._media_group_events.get(media_group_id)
if existing is None:
self._media_group_events[media_group_id] = event
@ -7297,15 +7376,20 @@ class TelegramAdapter(BasePlatformAdapter):
)
async def _flush_media_group_event(self, media_group_id: str) -> None:
current_task = asyncio.current_task()
try:
await asyncio.sleep(self.MEDIA_GROUP_WAIT_SECONDS)
event = self._media_group_events.pop(media_group_id, None)
if event is not None:
if self._should_drop_delayed_delivery():
logger.debug("[Telegram] Dropping media group flush after disconnect started")
return
await self.handle_message(event)
except asyncio.CancelledError:
return
finally:
self._media_group_tasks.pop(media_group_id, None)
if self._media_group_tasks.get(media_group_id) is current_task:
self._media_group_tasks.pop(media_group_id, None)
async def _handle_sticker(self, msg: Message, event: "MessageEvent") -> None:
"""

View file

@ -7,7 +7,7 @@ from the same session and aggregate them before dispatching.
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, patch
import pytest
@ -23,9 +23,25 @@ def _make_adapter():
config = PlatformConfig(enabled=True, token="test-token")
adapter = object.__new__(TelegramAdapter)
adapter._platform = Platform.TELEGRAM
adapter.platform = Platform.TELEGRAM
adapter.config = config
adapter._running = True
adapter._fatal_error_code = None
adapter._fatal_error_message = None
adapter._fatal_error_retryable = True
adapter._drop_delayed_deliveries = False
adapter._pending_text_batches = {}
adapter._pending_text_batch_tasks = {}
adapter._pending_photo_batches = {}
adapter._pending_photo_batch_tasks = {}
adapter._media_group_events = {}
adapter._media_group_tasks = {}
adapter._polling_error_task = None
adapter._polling_heartbeat_task = None
adapter._app = None
adapter._bot = None
adapter._set_status_indicator = AsyncMock()
adapter._release_platform_lock = lambda: None
adapter._text_batch_delay_seconds = 0.1 # fast for tests
adapter._active_sessions = {}
adapter._pending_messages = {}
@ -164,3 +180,149 @@ class TestTextBatching:
adapter.handle_message.assert_called_once()
dispatched = adapter.handle_message.call_args[0][0]
assert dispatched.source.thread_id == "222"
@pytest.mark.asyncio
async def test_disconnect_cancels_pending_text_batch_without_dispatch(self):
"""Disconnect should not let buffered text flush into a stale run."""
adapter = _make_adapter()
adapter._enqueue_text_event(_make_event("stale text"))
await adapter.disconnect()
await asyncio.sleep(0.2)
adapter.handle_message.assert_not_called()
assert adapter._pending_text_batches == {}
assert adapter._pending_text_batch_tasks == {}
@pytest.mark.asyncio
async def test_disconnected_adapter_drops_pending_text_flush_before_dispatch(self):
"""A pending text flush should drop its event if teardown wins the race."""
adapter = _make_adapter()
adapter._enqueue_text_event(_make_event("stale text"))
adapter._mark_disconnected()
await asyncio.sleep(0.2)
adapter.handle_message.assert_not_called()
assert adapter._pending_text_batches == {}
assert adapter._pending_text_batch_tasks == {}
@pytest.mark.asyncio
async def test_disconnected_adapter_drops_late_text_batch_enqueue(self):
"""Late update handlers should not schedule batches after teardown starts."""
adapter = _make_adapter()
adapter._mark_disconnected()
adapter._enqueue_text_event(_make_event("late text"))
await asyncio.sleep(0.2)
adapter.handle_message.assert_not_called()
assert adapter._pending_text_batches == {}
assert adapter._pending_text_batch_tasks == {}
@pytest.mark.asyncio
async def test_disconnected_adapter_drops_pending_photo_flush_before_dispatch(self):
"""A pending photo batch should not dispatch after disconnect starts."""
adapter = _make_adapter()
adapter._media_batch_delay_seconds = 0.1
event = _make_event("photo caption")
event.media_urls = ["/tmp/photo.jpg"]
event.media_types = ["image/jpeg"]
adapter._enqueue_photo_event("chat:photo-burst", event)
adapter._mark_disconnected()
await asyncio.sleep(0.2)
adapter.handle_message.assert_not_called()
assert adapter._pending_photo_batches == {}
assert adapter._pending_photo_batch_tasks == {}
@pytest.mark.asyncio
async def test_disconnected_adapter_drops_pending_media_group_flush_before_dispatch(self):
"""A pending media group should not dispatch after disconnect starts."""
from plugins.platforms.telegram.adapter import TelegramAdapter
adapter = _make_adapter()
event = _make_event("album caption")
event.media_urls = ["/tmp/photo.jpg"]
event.media_types = ["image/jpeg"]
with patch.object(TelegramAdapter, "MEDIA_GROUP_WAIT_SECONDS", 0.1):
await adapter._queue_media_group_event("album-1", event)
adapter._mark_disconnected()
await asyncio.sleep(0.2)
adapter.handle_message.assert_not_called()
assert adapter._media_group_events == {}
assert adapter._media_group_tasks == {}
@pytest.mark.asyncio
async def test_stale_media_group_flush_does_not_clear_newer_task(self):
"""A cancelled album flush must not erase the replacement task handle."""
from plugins.platforms.telegram.adapter import TelegramAdapter
adapter = _make_adapter()
first = _make_event("first album caption")
first.media_urls = ["/tmp/first.jpg"]
first.media_types = ["image/jpeg"]
second = _make_event("second album caption")
second.media_urls = ["/tmp/second.jpg"]
second.media_types = ["image/jpeg"]
with patch.object(TelegramAdapter, "MEDIA_GROUP_WAIT_SECONDS", 1.0):
await adapter._queue_media_group_event("album-race", first)
first_task = adapter._media_group_tasks["album-race"]
await asyncio.sleep(0)
await adapter._queue_media_group_event("album-race", second)
replacement_task = adapter._media_group_tasks["album-race"]
assert replacement_task is not first_task
await asyncio.sleep(0)
assert adapter._media_group_tasks.get("album-race") is replacement_task
replacement_task.cancel()
await asyncio.gather(replacement_task, return_exceptions=True)
@pytest.mark.asyncio
async def test_cancel_pending_delivery_tasks_skips_current_polling_error_task(self):
"""The teardown helper must not cancel the coroutine doing cleanup."""
adapter = _make_adapter()
current_task = asyncio.current_task()
stale_task = asyncio.create_task(asyncio.sleep(60))
adapter._pending_text_batches["text"] = _make_event("text")
adapter._pending_text_batch_tasks["text"] = stale_task
adapter._polling_error_task = current_task
await adapter._cancel_pending_delivery_tasks()
assert stale_task.done()
assert stale_task.cancelled()
assert not current_task.cancelled()
assert adapter._pending_text_batches == {}
assert adapter._pending_text_batch_tasks == {}
assert adapter._polling_error_task is current_task
@pytest.mark.asyncio
async def test_disconnect_cancels_all_pending_delivery_task_maps(self):
"""Photo/media/polling delayed tasks are awaited and queues are cleared."""
adapter = _make_adapter()
tasks = [asyncio.create_task(asyncio.sleep(60)) for _ in range(4)]
adapter._pending_text_batches["text"] = _make_event("text")
adapter._pending_text_batch_tasks["text"] = tasks[0]
adapter._pending_photo_batches["photo"] = _make_event("photo")
adapter._pending_photo_batch_tasks["photo"] = tasks[1]
adapter._media_group_events["media"] = _make_event("media")
adapter._media_group_tasks["media"] = tasks[2]
adapter._polling_error_task = tasks[3]
await adapter.disconnect()
assert all(task.done() for task in tasks)
assert adapter._pending_text_batches == {}
assert adapter._pending_text_batch_tasks == {}
assert adapter._pending_photo_batches == {}
assert adapter._pending_photo_batch_tasks == {}
assert adapter._media_group_events == {}
assert adapter._media_group_tasks == {}
assert adapter._polling_error_task is None