Skip to content

fix(telegram): close reconnect races that leave adapter half-destroyed #128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hashbender wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(telegram): close reconnect races that leave adapter half-destroyed #128

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 28 additions & 61 deletions gateway/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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





Expand Down Expand Up @@ -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 <prompt>"
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,
)
Comment on lines 8456 to 8468

@tenki-reviewer tenki-reviewer Bot Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 /queue command silently drops media attachments, reply context, and event metadata (bug)

The refactored /queue handler in GatewayRunner._handle_message (gateway/run.py:8462-8468) creates a bare MessageEvent that discards all media (media_urls, media_types), reply context (reply_to_message_id, reply_to_text, reply_to_author_id, reply_to_author_name, reply_to_is_own_message), auto_skill, raw_message, internal flag, and timestamp from the original event. This reverses a prior fix that intentionally preserved these fields. Users who /queue a photo, document, or replied-to message silently lose the attachment when the queued turn runs. Additionally, the has_media check was removed so queuing with only media (no caption) is now rejected outright.

💡 Suggestion: Restore the media/reply field copying from the original event when constructing the queued MessageEvent. Reintroduce the has_media check so /queue with an attachment but no text caption is still allowed (the old code's documented intent).

📋 Prompt for AI Agents

In gateway/run.py at the /queue handler around line 8456, replace lines 8457-8468 with: (1) compute has_media = bool(getattr(event, 'media_urls', None)), (2) guard with 'if not queued_text and not has_media: return usage line', (3) construct queued_event with message_type=event.message_type if has_media else MessageType.TEXT, and copy the following fields from the original event: raw_message, media_urls, media_types, reply_to_message_id, reply_to_text, reply_to_author_id, reply_to_author_name, reply_to_is_own_message, auto_skill, internal, timestamp. Keep source, message_id, and channel_prompt as-is.

self._enqueue_fifo(_quick_key, queued_event, adapter)
depth = self._queue_depth(_quick_key, adapter=self.adapters.get(source.platform))
Expand Down
136 changes: 33 additions & 103 deletions plugins/platforms/telegram/adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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():
Expand All @@ -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:
Expand All @@ -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()
Comment on lines +2974 to +2978

@tenki-reviewer tenki-reviewer Bot Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Telegram adapter disconnect leaks text-batch and polling-error tasks into torn-down session (bug)

The refactored TelegramAdapter.disconnect() (adapter.py:2932-2983) cancels media-group tasks and photo-batch tasks but omits two task families: (1) text-batch tasks in _pending_text_batch_tasks are never cancelled or cleared, and their asyncio.sleep() delay means they can fire handle_message() on an adapter where self._app and self._bot are already None; (2) self._polling_error_task is never cancelled, so an in-flight recovery cycle will keep retrying start_polling() on a torn-down adapter until MAX_NETWORK_RETRIES exhausts (~435s of wasted sleep). Additionally, _should_drop_delayed_delivery() guards were removed from all enqueue/flush methods, so new tasks can still be spawned during teardown.

💡 Suggestion: Add text-batch task cancellation (mirroring the photo-batch pattern at lines 2974-2978) and polling-error task cancellation (mirroring the heartbeat pattern at lines 2936-2942) to disconnect(). Clear the corresponding dicts after cancellation.

📋 Prompt for AI Agents

In disconnect() at plugins/platforms/telegram/adapter.py, add two cleanup blocks: (A) Before the app-shutdown block (around line 2962), cancel and clear _polling_error_task: if self._polling_error_task and not self._polling_error_task.done(): self._polling_error_task.cancel(); self._polling_error_task = None. (B) After the photo-batch cleanup block (around line 2978), add text-batch cancellation: iterate self._pending_text_batch_tasks.values(), cancel each non-done task, then clear both _pending_text_batch_tasks and _pending_text_batches.


self._mark_disconnected()
self._app = None
self._bot = None
logger.info("[%s] Disconnected from Telegram", self.name)
Expand Down Expand Up @@ -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 "")
Expand Down Expand Up @@ -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 ""),
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)

@tenki-reviewer tenki-reviewer Bot Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Media-group flush task unconditionally pops replacement task from tracking dict (bug)

In _flush_media_group_event (adapter.py:7322), the finally block performs self._media_group_tasks.pop(media_group_id, None) without checking whether the stored task is still the current task. When _queue_media_group_event cancels a prior task and creates a replacement, the cancelled task's finally block executes after the replacement is stored, removes the replacement's entry, and orphanes it. The orphaned task is invisible to disconnect()'s cancellation loop (list(self._media_group_tasks.values())) and fires handle_message() on a potentially torn-down adapter. The sibling flush methods _flush_text_batch (line 6956) and _flush_photo_batch (line 6987) retain the identity guard — this inconsistency suggests an oversight.

💡 Suggestion: Restore the current_task identity check before popping from _media_group_tasks, matching the pattern still used by _flush_text_batch and _flush_photo_batch.

Suggested change
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)
📋 Prompt for AI Agents

In _flush_media_group_event at plugins/platforms/telegram/adapter.py, capture current_task = asyncio.current_task() before the try block (around line 7313) and change line 7322 from self._media_group_tasks.pop(media_group_id, None) to 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:
"""
Expand Down
Loading
Loading