mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-19 15:18:03 +00:00
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:
parent
7de485703b
commit
8ad15ff7dd
2 changed files with 263 additions and 17 deletions
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue