diff --git a/gateway/run.py b/gateway/run.py index 9ac43b5f60b5..521854901e96 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -44,7 +44,7 @@ from contextvars import copy_context from pathlib import Path from datetime import datetime -from typing import Callable, Dict, Optional, Any, List, Union +from typing import Dict, Optional, Any, List, Union # account_usage imports the OpenAI SDK chain (~230 ms). Only needed by # /usage; we still import it at module top in the gateway because test @@ -3644,6 +3644,23 @@ async def _handle_adapter_fatal_error(self, adapter: BasePlatformAdapter) -> Non If the error is retryable (e.g. network blip, DNS failure), queue the platform for background reconnection instead of giving up permanently. """ + # Snapshot the current owner of this platform slot before doing + # anything else. If it's neither this adapter nor empty, a different + # adapter has already taken over (e.g. this is a delayed notification + # from a background retry chain that raced with, and lost to, a + # reconnect that already succeeded). Acting on a stale notification + # would overwrite an already-healthy platform's runtime status and + # incorrectly re-queue it for reconnection, so bail out before any of + # that happens. + existing = self.adapters.get(adapter.platform) + if existing is not None and existing is not adapter: + logger.debug( + "Ignoring stale fatal error from a superseded %s adapter instance: %s", + adapter.platform.value, + adapter.fatal_error_code or "unknown", + ) + return + logger.error( "Fatal %s adapter error (%s): %s", adapter.platform.value, @@ -3667,13 +3684,15 @@ async def _handle_adapter_fatal_error(self, adapter: BasePlatformAdapter) -> Non error_message=adapter.fatal_error_message, ) - existing = self.adapters.get(adapter.platform) if existing is adapter: - try: - await adapter.disconnect() - finally: - self.adapters.pop(adapter.platform, None) - self.delivery_router.adapters = self.adapters + # Claim this adapter for teardown before awaiting disconnect() — + # a second fatal-error notification for the same adapter (e.g. + # from a concurrent recovery path) would otherwise still see + # itself as "existing" during the await below and disconnect() + # the same object twice. + self.adapters.pop(adapter.platform, None) + self.delivery_router.adapters = self.adapters + await adapter.disconnect() # Queue retryable failures for background reconnection if adapter.fatal_error_retryable: @@ -6354,7 +6373,6 @@ async def start(self) -> bool: adapter.set_session_store(self.session_store) adapter.set_busy_session_handler(self._handle_active_session_busy_message) adapter.set_topic_recovery_fn(self._recover_telegram_topic_thread_id) - adapter.set_authorization_check(self._make_adapter_auth_check(adapter.platform)) adapter._busy_text_mode = self._busy_text_mode # Try to connect @@ -7163,7 +7181,6 @@ async def _platform_reconnect_watcher(self) -> None: adapter.set_session_store(self.session_store) adapter.set_busy_session_handler(self._handle_active_session_busy_message) adapter.set_topic_recovery_fn(self._recover_telegram_topic_thread_id) - adapter.set_authorization_check(self._make_adapter_auth_check(adapter.platform)) adapter._busy_text_mode = self._busy_text_mode # Reconnect after an outage: preserve the platform's @@ -7820,7 +7837,6 @@ async def _start_one_profile_adapters( adapter.set_session_store(self.session_store) adapter.set_busy_session_handler(self._handle_active_session_busy_message) adapter.set_topic_recovery_fn(self._recover_telegram_topic_thread_id) - adapter.set_authorization_check(self._make_adapter_auth_check(adapter.platform)) adapter._busy_text_mode = self._busy_text_mode try: @@ -7989,39 +8005,6 @@ def _create_adapter( return None - def _make_adapter_auth_check( - self, - platform: Platform, - ) -> Callable[[str, Optional[str], Optional[str]], bool]: - """Build a platform-bound auth callback for adapter use. - - Adapters that fetch external context (e.g. Slack - ``conversations.replies``) call this through - ``BasePlatformAdapter._is_sender_authorized`` to mark non-allowlisted - senders as unverified in LLM context, mitigating indirect prompt - injection from third parties in shared threads/channels. - - The returned callback delegates to :meth:`_is_user_authorized` so the - full auth chain — platform allowlists, group allowlists, pairing - store, allow-all flags — stays the single source of truth. - """ - def check( - user_id: str, - chat_type: Optional[str] = None, - chat_id: Optional[str] = None, - ) -> bool: - if not user_id: - return False - source = SessionSource( - platform=platform, - chat_id=chat_id or "", - chat_type=chat_type or "group", - user_id=user_id, - ) - return self._is_user_authorized(source) - return check - - @@ -8472,32 +8455,16 @@ async def _handle_message(self, event: MessageEvent) -> Optional[str]: # earlier /queue items) finishes. Messages are NOT merged. if event.get_command() in {"queue", "q"}: queued_text = event.get_command_args().strip() - # Preserve media/reply payloads: a /queue carrying a photo, - # document, or reply context is valid even with no prompt text - # (e.g. "/queue" as the caption of an image). Dropping these - # fields silently lost the attachment when the queued turn ran. - has_media = bool(getattr(event, "media_urls", None)) - if not queued_text and not has_media: + if not queued_text: return "Usage: /queue " adapter = self.adapters.get(source.platform) if adapter: queued_event = MessageEvent( text=queued_text, - message_type=event.message_type if has_media else MessageType.TEXT, + message_type=MessageType.TEXT, source=event.source, - raw_message=event.raw_message, message_id=event.message_id, - media_urls=list(getattr(event, "media_urls", []) or []), - media_types=list(getattr(event, "media_types", []) or []), - reply_to_message_id=event.reply_to_message_id, - reply_to_text=event.reply_to_text, - reply_to_author_id=event.reply_to_author_id, - reply_to_author_name=event.reply_to_author_name, - reply_to_is_own_message=event.reply_to_is_own_message, - auto_skill=event.auto_skill, channel_prompt=event.channel_prompt, - internal=event.internal, - timestamp=event.timestamp, ) self._enqueue_fifo(_quick_key, queued_event, adapter) depth = self._queue_depth(_quick_key, adapter=self.adapters.get(source.platform)) diff --git a/plugins/platforms/telegram/adapter.py b/plugins/platforms/telegram/adapter.py index e816bbcbf2cc..e1ac7b60a1fc 100644 --- a/plugins/platforms/telegram/adapter.py +++ b/plugins/platforms/telegram/adapter.py @@ -412,7 +412,6 @@ def __init__(self, config: PlatformConfig): ) 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 @@ -500,27 +499,6 @@ def __init__(self, config: PlatformConfig): # 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]: @@ -1773,16 +1751,24 @@ async def _handle_polling_network_error(self, error: Exception) -> None: ) await asyncio.sleep(delay) + # Capture a stable local reference: self._app can be reassigned to None + # by a concurrent disconnect() while we're suspended across the awaits + # below, and re-reading self._app after that point would silently swap + # in None mid-sequence instead of failing fast in one place. + app = self._app + try: - if self._app and self._app.updater and self._app.updater.running: - await self._app.updater.stop() + if app and app.updater and app.updater.running: + await app.updater.stop() except Exception: pass await self._drain_polling_connections() try: - await self._app.updater.start_polling( + if not app: + raise RuntimeError("Telegram application was torn down during reconnect") + await app.updater.start_polling( allowed_updates=Update.ALL_TYPES, drop_pending_updates=False, error_callback=self._polling_error_callback_ref, @@ -1824,6 +1810,12 @@ async def _handle_polling_network_error(self, error: Exception) -> None: ) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) + # This chained retry IS the in-flight recovery attempt — it + # must replace the reentrancy guard, otherwise the heartbeat + # loop, the pending-updates probe, and the PTB error callback + # all see _polling_error_task as "done" and can each start a + # second, concurrent recovery for the same outage. + self._polling_error_task = task async def _polling_heartbeat_loop(self) -> None: """Detect dead Telegram TCP sockets (CLOSE-WAIT) by periodic probing. @@ -2937,60 +2929,8 @@ async def _set_status_indicator(self, online: bool) -> None: 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 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() - + """Stop polling/webhook, cancel pending album flushes, and disconnect.""" # 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(): @@ -3011,7 +2951,13 @@ async def disconnect(self) -> None: except Exception: pass - await self._cancel_pending_delivery_tasks() + 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() if self._app: try: @@ -3025,6 +2971,13 @@ 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 logger.info("[%s] Disconnected from Telegram", self.name) @@ -6935,10 +6888,6 @@ 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 "") @@ -6998,9 +6947,6 @@ 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 ""), @@ -7035,9 +6981,6 @@ 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: @@ -7046,10 +6989,6 @@ 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 @@ -7354,10 +7293,6 @@ 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 @@ -7376,20 +7311,15 @@ 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: - if self._media_group_tasks.get(media_group_id) is current_task: - self._media_group_tasks.pop(media_group_id, None) + self._media_group_tasks.pop(media_group_id, None) async def _handle_sticker(self, msg: Message, event: "MessageEvent") -> None: """ diff --git a/tests/gateway/test_runner_fatal_adapter.py b/tests/gateway/test_runner_fatal_adapter.py index 7e7739582d16..7fce3841fde5 100644 --- a/tests/gateway/test_runner_fatal_adapter.py +++ b/tests/gateway/test_runner_fatal_adapter.py @@ -1,3 +1,4 @@ +import asyncio from unittest.mock import AsyncMock import pytest @@ -98,3 +99,96 @@ async def test_runner_queues_retryable_runtime_fatal_for_reconnection(monkeypatc assert runner._exit_with_failure is False assert Platform.WHATSAPP in runner._failed_platforms assert runner._failed_platforms[Platform.WHATSAPP]["attempts"] == 0 + + +@pytest.mark.asyncio +async def test_concurrent_fatal_notifications_disconnect_same_adapter_once(monkeypatch, tmp_path): + """ + Two fatal-error notifications for the same still-installed adapter (e.g. + from two concurrent recovery paths racing on the same underlying outage) + must result in exactly one disconnect() call. + + Regression test for the TOCTOU race in _handle_adapter_fatal_error: the + old code only removed the adapter from self.adapters in a `finally` block + *after* awaiting disconnect(), so a second concurrent call could still see + itself as "existing" and disconnect() the same object twice — the + concrete origin of the "'NoneType' object has no attribute 'updater'" + crash when the adapter's own teardown code re-reads self._app afterwards. + """ + config = GatewayConfig( + platforms={ + Platform.WHATSAPP: PlatformConfig(enabled=True, token="token") + }, + sessions_dir=tmp_path / "sessions", + ) + runner = GatewayRunner(config) + adapter = _RuntimeRetryableAdapter() + adapter._set_fatal_error( + "whatsapp_bridge_exited", + "WhatsApp bridge process exited unexpectedly (code 1).", + retryable=True, + ) + + runner.adapters = {Platform.WHATSAPP: adapter} + runner.delivery_router.adapters = runner.adapters + runner.stop = AsyncMock() + + disconnect_calls = 0 + release_second_call = asyncio.Event() + + async def slow_disconnect(): + nonlocal disconnect_calls + disconnect_calls += 1 + # Yield control so the second concurrent notification can run its + # "existing is adapter" check before this call finishes tearing down. + release_second_call.set() + await asyncio.sleep(0) + adapter._mark_disconnected() + + monkeypatch.setattr(adapter, "disconnect", slow_disconnect) + + await asyncio.gather( + runner._handle_adapter_fatal_error(adapter), + runner._handle_adapter_fatal_error(adapter), + ) + + assert disconnect_calls == 1 + + +@pytest.mark.asyncio +async def test_stale_fatal_notification_from_superseded_adapter_is_ignored(monkeypatch, tmp_path): + """ + A delayed fatal-error notification from an adapter instance that has + since been replaced by a different, already-installed adapter (e.g. a + background retry chain on the old instance finally giving up after a + reconnect on a new instance already succeeded) must be ignored: it must + not disconnect the new adapter, must not re-queue an already-healthy + platform for reconnection, and must not shut the gateway down. + """ + config = GatewayConfig( + platforms={ + Platform.WHATSAPP: PlatformConfig(enabled=True, token="token") + }, + sessions_dir=tmp_path / "sessions", + ) + runner = GatewayRunner(config) + + old_adapter = _RuntimeRetryableAdapter() + old_adapter._set_fatal_error( + "whatsapp_bridge_exited", + "stale failure from a superseded adapter instance", + retryable=True, + ) + + new_adapter = _RuntimeRetryableAdapter() + new_adapter.disconnect = AsyncMock() + runner.adapters = {Platform.WHATSAPP: new_adapter} + runner.delivery_router.adapters = runner.adapters + runner.stop = AsyncMock() + + await runner._handle_adapter_fatal_error(old_adapter) + + new_adapter.disconnect.assert_not_awaited() + assert runner.adapters[Platform.WHATSAPP] is new_adapter + assert Platform.WHATSAPP not in runner._failed_platforms + runner.stop.assert_not_awaited() diff --git a/tests/gateway/test_telegram_network_reconnect.py b/tests/gateway/test_telegram_network_reconnect.py index 8c0dc6a563f2..c970adc8ff24 100644 --- a/tests/gateway/test_telegram_network_reconnect.py +++ b/tests/gateway/test_telegram_network_reconnect.py @@ -117,6 +117,43 @@ async def test_reconnect_does_not_self_schedule_when_fatal_error_set(): ) +@pytest.mark.asyncio +async def test_reconnect_chained_retry_updates_polling_error_task(): + """ + When start_polling() fails and the handler self-schedules a retry, that + retry task must become the new `_polling_error_task` — otherwise the + reentrancy guard used by the heartbeat loop, the pending-updates probe, + and the PTB error callback goes stale while a recovery is still in + flight, letting a second concurrent recovery start for the same outage. + + Regression test for the race behind the "half-destroyed adapter" bug + (gateway reports connected but silently stops processing messages). + """ + adapter = _make_adapter() + adapter._polling_network_error_count = 1 + + mock_updater = MagicMock() + mock_updater.running = True + mock_updater.stop = AsyncMock() + mock_updater.start_polling = AsyncMock(side_effect=Exception("Timed out")) + + mock_app = MagicMock() + mock_app.updater = mock_updater + adapter._app = mock_app + + with patch("asyncio.sleep", new_callable=AsyncMock): + await adapter._handle_polling_network_error(Exception("Bad Gateway")) + + assert adapter._polling_error_task is not None + assert not adapter._polling_error_task.done() + + adapter._polling_error_task.cancel() + try: + await adapter._polling_error_task + except (asyncio.CancelledError, Exception): + pass + + @pytest.mark.asyncio async def test_reconnect_success_resets_error_count(): """