fix(telegram): cancel delayed deliveries on disconnect

gateway/platforms/telegram.py

@@ -239,6 +239,7 @@ def __init__(self, config: PlatformConfig):
         self._text_batch_split_delay_seconds = float(os.getenv("HERMES_TELEGRAM_TEXT_BATCH_SPLIT_DELAY_SECONDS", "2.0"))
         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
@@ -252,6 +253,21 @@ def __init__(self, config: PlatformConfig):
         # Approval button state: message_id → session_key
         self._approval_state: Dict[int, 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:
+        return bool(getattr(self, "_drop_delayed_deliveries", False))
+
     @staticmethod
     def _is_callback_user_authorized(user_id: str) -> bool:
         """Return whether a Telegram inline-button caller may perform gated actions."""
@@ -914,15 +930,50 @@ def _polling_error_callback(error: Exception) -> None:
             logger.error("[%s] Failed to connect to Telegram: %s", self.name, e, exc_info=True)
             return False
 
-    async def disconnect(self) -> None:
-        """Stop polling/webhook, cancel pending album flushes, and disconnect."""
-        pending_media_group_tasks = list(self._media_group_tasks.values())
-        for task in pending_media_group_tasks:
+    async def _cancel_pending_delivery_tasks(self) -> None:
+        """Cancel delayed Telegram delivery tasks before disconnect completes."""
+        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 pending_media_group_tasks:
-            await asyncio.gather(*pending_media_group_tasks, return_exceptions=True)
+        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 delayed deliveries, and disconnect."""
+        self._mark_disconnected()
+        await self._cancel_pending_delivery_tasks()
 
         if self._app:
             try:
@@ -936,12 +987,6 @@ async def disconnect(self) -> None:
                 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
@@ -2513,6 +2558,10 @@ def _enqueue_text_event(self, event: MessageEvent) -> None:
         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 "")
@@ -2557,6 +2606,9 @@ async def _flush_text_batch(self, key: str) -> None:
             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 ""),
@@ -2591,6 +2643,9 @@ async def _flush_photo_batch(self, batch_key: str) -> None:
             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:
@@ -2599,6 +2654,10 @@ async def _flush_photo_batch(self, batch_key: str) -> None:
 
     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
@@ -2821,6 +2880,10 @@ async def _queue_media_group_event(self, media_group_id: str, event: MessageEven
         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
@@ -2839,15 +2902,20 @@ async def _queue_media_group_event(self, media_group_id: str, event: MessageEven
         )
 
     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:
         """

tests/gateway/test_telegram_text_batching.py

@@ -21,9 +21,21 @@ 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._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._app = None
+    adapter._bot = None
     adapter._text_batch_delay_seconds = 0.1  # fast for tests
     adapter._active_sessions = {}
     adapter._pending_messages = {}
@@ -119,3 +131,150 @@ async def test_batch_cleans_up_after_flush(self):
 
         assert len(adapter._pending_text_batches) == 0
         assert len(adapter._pending_text_batch_tasks) == 0
+
+
+    @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 gateway.platforms.telegram 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 gateway.platforms.telegram 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