diff --git a/astrbot/core/platform/sources/telegram/tg_adapter.py b/astrbot/core/platform/sources/telegram/tg_adapter.py index 46887344d..a37e5cfd4 100644 --- a/astrbot/core/platform/sources/telegram/tg_adapter.py +++ b/astrbot/core/platform/sources/telegram/tg_adapter.py @@ -3,12 +3,13 @@ import os import re import sys import uuid +from contextlib import suppress from typing import cast from apscheduler.schedulers.asyncio import AsyncIOScheduler from telegram import BotCommand, Update from telegram.constants import ChatType -from telegram.error import Forbidden, InvalidToken +from telegram.error import Forbidden, InvalidToken, NetworkError from telegram.ext import ApplicationBuilder, ContextTypes, ExtBot, filters from telegram.ext import MessageHandler as TelegramMessageHandler @@ -66,6 +67,7 @@ class TelegramPlatformAdapter(Platform): file_base_url = "https://api.telegram.org/file/bot" self.base_url = base_url + self.file_base_url = file_base_url self.enable_command_register = self.config.get( "telegram_command_register", @@ -77,23 +79,12 @@ class TelegramPlatformAdapter(Platform): ) self.last_command_hash = None - self.application = ( - ApplicationBuilder() - .token(self.config["telegram_token"]) - .base_url(base_url) - .base_file_url(file_base_url) - .build() - ) - message_handler = TelegramMessageHandler( - filters=filters.ALL, # receive all messages - callback=self.message_handler, - ) - self.application.add_handler(message_handler) - self.client = self.application.bot - logger.debug(f"Telegram base url: {self.client.base_url}") - self.scheduler = AsyncIOScheduler() self._terminating = False + self._loop: asyncio.AbstractEventLoop | None = None + self._polling_recovery_requested = asyncio.Event() + self._consecutive_polling_failures = 0 + self._last_polling_failure_at = 0.0 raw_delay = self.config.get("telegram_polling_restart_delay", 5.0) try: delay = float(raw_delay) @@ -113,6 +104,10 @@ class TelegramPlatformAdapter(Platform): ) delay = 0.1 self._polling_restart_delay = delay + self._polling_recovery_threshold = 3 + self._polling_failure_window = 60.0 + self._application_started = False + self._build_application() # Media group handling # Cache structure: {media_group_id: {"created_at": datetime, "items": [(update, context), ...]}} @@ -124,6 +119,85 @@ class TelegramPlatformAdapter(Platform): "telegram_media_group_max_wait", 10.0 ) # max seconds - hard cap to prevent indefinite delay + def _build_application(self) -> None: + self.application = ( + ApplicationBuilder() + .token(self.config["telegram_token"]) + .base_url(self.base_url) + .base_file_url(self.file_base_url) + .build() + ) + message_handler = TelegramMessageHandler( + filters=filters.ALL, + callback=self.message_handler, + ) + self.application.add_handler(message_handler) + self.client = self.application.bot + logger.debug(f"Telegram base url: {self.client.base_url}") + + async def _start_application(self) -> None: + await self.application.initialize() + await self.application.start() + + if self.enable_command_register: + await self.register_commands() + + self._application_started = True + + async def _shutdown_application( + self, + *, + delete_commands: bool, + ) -> None: + self._application_started = False + + updater = self.application.updater + if updater is not None: + with suppress(Exception): + await updater.stop() + + if delete_commands and self.enable_command_register: + with suppress(Exception): + await self.client.delete_my_commands() + + with suppress(Exception): + await self.application.stop() + + shutdown = getattr(self.application, "shutdown", None) + if shutdown is not None: + with suppress(Exception): + await shutdown() + + async def _recreate_application(self) -> None: + if self._terminating: + self._polling_recovery_requested.clear() + return + + logger.warning( + "Telegram polling hit repeated network errors; rebuilding the " + "Telegram application and HTTP client.", + ) + await self._shutdown_application(delete_commands=False) + self._build_application() + self._consecutive_polling_failures = 0 + self._last_polling_failure_at = 0.0 + self._polling_recovery_requested.clear() + + def _start_command_scheduler(self) -> None: + if not self.enable_command_refresh or not self.enable_command_register: + return + if self.scheduler.running: + return + + self.scheduler.add_job( + self.register_commands, + "interval", + seconds=self.config.get("telegram_command_register_interval", 300), + id="telegram_command_register", + misfire_grace_time=60, + ) + self.scheduler.start() + @override async def send_by_session( self, @@ -145,41 +219,42 @@ class TelegramPlatformAdapter(Platform): @override async def run(self) -> None: - await self.application.initialize() - await self.application.start() - - if self.enable_command_register: - await self.register_commands() - - if self.enable_command_refresh and self.enable_command_register: - self.scheduler.add_job( - self.register_commands, - "interval", - seconds=self.config.get("telegram_command_register_interval", 300), - id="telegram_command_register", - misfire_grace_time=60, - ) - self.scheduler.start() - - if not self.application.updater: - logger.error("Telegram Updater is not initialized. Cannot start polling.") - return + self._loop = asyncio.get_running_loop() + self._start_command_scheduler() while not self._terminating: try: + if not self._application_started: + await self._start_application() + + self._polling_recovery_requested.clear() + updater = self.application.updater + if updater is None: + logger.error( + "Telegram Updater is not initialized. Cannot start polling." + ) + self._application_started = False + await asyncio.sleep(self._polling_restart_delay) + continue logger.info("Starting Telegram polling...") - await self.application.updater.start_polling( - error_callback=self._on_polling_error - ) + await updater.start_polling(error_callback=self._on_polling_error) logger.info("Telegram Platform Adapter is running.") - while self.application.updater.running and not self._terminating: # noqa: ASYNC110 + while updater.running and not self._terminating: # noqa: ASYNC110 + if self._polling_recovery_requested.is_set(): + await self._recreate_application() + break await asyncio.sleep(1) + else: + if not self._terminating: + logger.warning( + "Telegram polling loop exited unexpectedly, " + f"retrying in {self._polling_restart_delay}s." + ) + continue if not self._terminating: - logger.warning( - "Telegram polling loop exited unexpectedly, " - f"retrying in {self._polling_restart_delay}s." - ) + logger.info("Telegram polling restarted with a fresh client.") + continue except asyncio.CancelledError: raise except (Forbidden, InvalidToken) as e: @@ -193,6 +268,9 @@ class TelegramPlatformAdapter(Platform): f"{type(e).__name__}: {e!s}. " f"Retrying in {self._polling_restart_delay}s.", ) + with suppress(Exception): + await self._shutdown_application(delete_commands=False) + self._build_application() if not self._terminating: await asyncio.sleep(self._polling_restart_delay) @@ -202,6 +280,33 @@ class TelegramPlatformAdapter(Platform): f"Telegram polling request failed: {type(error).__name__}: {error!s}", exc_info=error, ) + if not isinstance(error, NetworkError): + return + + if self._loop is None: + return + + now = self._loop.time() + if now - self._last_polling_failure_at > self._polling_failure_window: + self._consecutive_polling_failures = 0 + self._last_polling_failure_at = now + self._consecutive_polling_failures += 1 + + if self._consecutive_polling_failures < self._polling_recovery_threshold: + return + + logger.warning( + "Telegram polling encountered %s network failures within %.1fs; " + "scheduling client rebuild.", + self._consecutive_polling_failures, + self._polling_failure_window, + ) + if self._loop.is_closed(): + return + try: + self._loop.call_soon_threadsafe(self._polling_recovery_requested.set) + except RuntimeError: + return async def register_commands(self) -> None: """收集所有注册的指令并注册到 Telegram""" @@ -634,15 +739,8 @@ class TelegramPlatformAdapter(Platform): self._terminating = True if self.scheduler.running: self.scheduler.shutdown() - - await self.application.stop() - - if self.enable_command_register: - await self.client.delete_my_commands() - - # 保险起见先判断是否存在updater对象 - if self.application.updater is not None: - await self.application.updater.stop() + self._polling_recovery_requested.set() + await self._shutdown_application(delete_commands=True) logger.info("Telegram adapter has been closed.") except Exception as e: diff --git a/tests/fixtures/mocks/telegram.py b/tests/fixtures/mocks/telegram.py index 079292951..e5e31f178 100644 --- a/tests/fixtures/mocks/telegram.py +++ b/tests/fixtures/mocks/telegram.py @@ -9,6 +9,18 @@ from unittest.mock import AsyncMock, MagicMock import pytest +class MockTelegramNetworkError(Exception): + """Mock telegram.error.NetworkError used in tests.""" + + +class MockTelegramForbidden(Exception): + """Mock telegram.error.Forbidden used in tests.""" + + +class MockTelegramInvalidToken(Exception): + """Mock telegram.error.InvalidToken used in tests.""" + + def create_mock_telegram_modules(): """创建 Telegram 相关的 mock 模块。 @@ -28,6 +40,9 @@ def create_mock_telegram_modules(): mock_telegram.constants.ChatAction.UPLOAD_PHOTO = "upload_photo" mock_telegram.error = MagicMock() mock_telegram.error.BadRequest = Exception + mock_telegram.error.Forbidden = MockTelegramForbidden + mock_telegram.error.InvalidToken = MockTelegramInvalidToken + mock_telegram.error.NetworkError = MockTelegramNetworkError mock_telegram.ReactionTypeCustomEmoji = MagicMock mock_telegram.ReactionTypeEmoji = MagicMock @@ -73,7 +88,9 @@ def mock_telegram_modules(): monkeypatch.setitem(sys.modules, "telegram.constants", mocks["telegram"].constants) monkeypatch.setitem(sys.modules, "telegram.error", mocks["telegram"].error) monkeypatch.setitem(sys.modules, "telegram.ext", mocks["telegram.ext"]) - monkeypatch.setitem(sys.modules, "telegramify_markdown", mocks["telegramify_markdown"]) + monkeypatch.setitem( + sys.modules, "telegramify_markdown", mocks["telegramify_markdown"] + ) monkeypatch.setitem(sys.modules, "apscheduler", mocks["apscheduler"]) monkeypatch.setitem( sys.modules, "apscheduler.schedulers", mocks["apscheduler"].schedulers @@ -120,16 +137,16 @@ class MockTelegramBuilder: from tests.fixtures.helpers import NoopAwaitable app = MagicMock() - app.bot = MagicMock() - app.bot.username = "test_bot" - app.bot.base_url = "https://api.telegram.org/bottest_token_123/" + app.bot = MockTelegramBuilder.create_bot() app.initialize = AsyncMock() app.start = AsyncMock() app.stop = AsyncMock() + app.shutdown = AsyncMock() app.add_handler = MagicMock() app.updater = MagicMock() app.updater.start_polling = MagicMock(return_value=NoopAwaitable()) app.updater.stop = AsyncMock() + app.updater.running = False return app @staticmethod diff --git a/tests/test_telegram_adapter.py b/tests/test_telegram_adapter.py index 12f8a821b..8290899d1 100644 --- a/tests/test_telegram_adapter.py +++ b/tests/test_telegram_adapter.py @@ -7,11 +7,16 @@ import pytest import astrbot.api.message_components as Comp from tests.fixtures.helpers import ( + NoopAwaitable, create_mock_file, create_mock_update, make_platform_config, ) -from tests.fixtures.mocks.telegram import create_mock_telegram_modules +from tests.fixtures.mocks.telegram import ( + MockTelegramBuilder, + MockTelegramNetworkError, + create_mock_telegram_modules, +) _TELEGRAM_PLATFORM_ADAPTER = None _TELEGRAM_PLATFORM_EVENT = None @@ -212,3 +217,170 @@ async def test_telegram_final_segment_splits_long_plaintext_when_markdown_fails( assert len(second_call["text"]) == 18 assert "parse_mode" not in first_call assert "parse_mode" not in second_call + + +@pytest.mark.asyncio +async def test_telegram_polling_error_requests_rebuild_after_threshold(): + TelegramPlatformAdapter = _load_telegram_adapter() + adapter = TelegramPlatformAdapter( + make_platform_config("telegram"), + {}, + asyncio.Queue(), + ) + adapter._loop = asyncio.get_running_loop() + + assert not adapter._polling_recovery_requested.is_set() + + for _ in range(adapter._polling_recovery_threshold): + adapter._on_polling_error(MockTelegramNetworkError("proxy disconnected")) + + await asyncio.sleep(0) + + assert adapter._polling_recovery_requested.is_set() + + +@pytest.mark.asyncio +async def test_telegram_run_rebuilds_application_after_repeated_polling_errors(): + TelegramPlatformAdapter = _load_telegram_adapter() + module_globals = TelegramPlatformAdapter.__init__.__globals__ + app_one = MockTelegramBuilder.create_application() + app_one.updater.running = True + app_two = MockTelegramBuilder.create_application() + app_two.updater.running = True + created_apps = [app_one, app_two] + + builder = MagicMock() + builder.token.return_value = builder + builder.base_url.return_value = builder + builder.base_file_url.return_value = builder + builder.build.side_effect = created_apps + + adapter = None + + def start_polling_side_effect(*args, **kwargs): + nonlocal adapter + error_callback = kwargs["error_callback"] + assert adapter is not None + + async def _emit_errors(): + await asyncio.sleep(0) + for _ in range(adapter._polling_recovery_threshold): + error_callback(MockTelegramNetworkError("proxy disconnected")) + + asyncio.create_task(_emit_errors()) + return NoopAwaitable() + + app_one.updater.start_polling.side_effect = start_polling_side_effect + + async def second_start_polling(*args, **kwargs): + assert adapter is not None + adapter._terminating = True + + app_two.updater.start_polling.side_effect = second_start_polling + + with patch.dict( + module_globals, + { + "ApplicationBuilder": MagicMock(return_value=builder), + "AsyncIOScheduler": MagicMock( + return_value=MockTelegramBuilder.create_scheduler() + ), + }, + ): + adapter = TelegramPlatformAdapter( + make_platform_config("telegram"), + {}, + asyncio.Queue(), + ) + await adapter.run() + + assert builder.build.call_count == 2 + app_one.updater.stop.assert_awaited() + app_one.bot.delete_my_commands.assert_not_awaited() + app_one.stop.assert_awaited() + app_one.shutdown.assert_awaited() + app_two.initialize.assert_awaited() + app_two.start.assert_awaited() + + +@pytest.mark.asyncio +async def test_telegram_recreate_application_is_skipped_during_termination(): + TelegramPlatformAdapter = _load_telegram_adapter() + adapter = TelegramPlatformAdapter( + make_platform_config("telegram"), + {}, + asyncio.Queue(), + ) + adapter._terminating = True + adapter._polling_recovery_requested.set() + + await adapter._recreate_application() + + assert not adapter._polling_recovery_requested.is_set() + + +@pytest.mark.asyncio +async def test_telegram_run_rebuilds_fresh_application_after_recreate_init_failure(): + TelegramPlatformAdapter = _load_telegram_adapter() + module_globals = TelegramPlatformAdapter.__init__.__globals__ + app_one = MockTelegramBuilder.create_application() + app_one.updater.running = True + app_two = MockTelegramBuilder.create_application() + app_three = MockTelegramBuilder.create_application() + app_three.updater.running = True + created_apps = [app_one, app_two, app_three] + + builder = MagicMock() + builder.token.return_value = builder + builder.base_url.return_value = builder + builder.base_file_url.return_value = builder + builder.build.side_effect = created_apps + + adapter = None + + def first_start_polling(*args, **kwargs): + nonlocal adapter + error_callback = kwargs["error_callback"] + assert adapter is not None + + async def _emit_errors(): + await asyncio.sleep(0) + for _ in range(adapter._polling_recovery_threshold): + error_callback(MockTelegramNetworkError("proxy disconnected")) + + asyncio.create_task(_emit_errors()) + return NoopAwaitable() + + app_one.updater.start_polling.side_effect = first_start_polling + app_two.initialize.side_effect = TimeoutError("init timeout") + + async def final_start_polling(*args, **kwargs): + assert adapter is not None + adapter._terminating = True + + app_three.updater.start_polling.side_effect = final_start_polling + + with patch.dict( + module_globals, + { + "ApplicationBuilder": MagicMock(return_value=builder), + "AsyncIOScheduler": MagicMock( + return_value=MockTelegramBuilder.create_scheduler() + ), + }, + ): + adapter = TelegramPlatformAdapter( + make_platform_config( + "telegram", + telegram_polling_restart_delay=0.1, + ), + {}, + asyncio.Queue(), + ) + await adapter.run() + + assert builder.build.call_count == 3 + app_two.stop.assert_awaited() + app_two.shutdown.assert_awaited() + app_three.initialize.assert_awaited() + app_three.start.assert_awaited()