fix: telegram polling recovery after network failures (#7468)

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:

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

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()