From 82881294752bae8e297bcd5d64767054db10a4e5 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Sun, 30 Aug 2026 22:00:04 +0530
Subject: [PATCH] fix(telegram): bound polling drain with wall-clock deadline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_drain_polling_connections still bounded its shutdown()/initialize() with
asyncio.wait_for (#66377), while its sibling the general-pool drain moved
to _await_with_thread_deadline (#98094). httpcore's pool close runs under
AsyncShieldCancellation, so a cancellation-resistant close keeps wait_for
pending forever even after its timeout fires — the tracked
_polling_error_task wedges and every escalation gate behind it stalls.

Use the same wall-clock deadline helper (cancel + abandon, no cancel-await)
on both polling-drain awaits, and add a regression test whose close
swallows cancellation — the shape the existing cancellable-hang test
cannot catch.
---
 plugins/platforms/telegram/adapter.py         | 14 ++++-
 .../test_telegram_network_reconnect.py        | 57 +++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/plugins/platforms/telegram/adapter.py b/plugins/platforms/telegram/adapter.py
index c782bfe8be..1eccbebd09 100644
--- a/plugins/platforms/telegram/adapter.py
+++ b/plugins/platforms/telegram/adapter.py
@@ -2496,15 +2496,23 @@ class TelegramAdapter(BasePlatformAdapter):
             return
         try:
             # Bounded: a wedged CLOSE-WAIT socket can make this close hang
-            # forever and freeze the reconnect ladder (#66377).
-            await asyncio.wait_for(polling_req.shutdown(), timeout=_DRAIN_TIMEOUT)
+            # forever and freeze the reconnect ladder (#66377). The wall-clock
+            # deadline helper — not asyncio.wait_for — because httpcore's pool
+            # close runs under AsyncShieldCancellation and a cancellation-
+            # resistant close wedges wait_for itself forever (#58236/#63309);
+            # same primitive as the general-pool drain (#98094).
+            await _await_with_thread_deadline(
+                polling_req.shutdown(), timeout=_DRAIN_TIMEOUT
+            )
         except Exception:
             logger.debug(
                 "[%s] Polling request shutdown failed/timed out (non-fatal)",
                 self.name, exc_info=True,
             )
         try:
-            await asyncio.wait_for(polling_req.initialize(), timeout=_DRAIN_TIMEOUT)
+            await _await_with_thread_deadline(
+                polling_req.initialize(), timeout=_DRAIN_TIMEOUT
+            )
             logger.debug(
                 "[%s] Polling request pool drained before reconnect", self.name
             )
diff --git a/tests/gateway/test_telegram_network_reconnect.py b/tests/gateway/test_telegram_network_reconnect.py
index 2a61e7139f..607f2c78a8 100644
--- a/tests/gateway/test_telegram_network_reconnect.py
+++ b/tests/gateway/test_telegram_network_reconnect.py
@@ -326,6 +326,63 @@ async def test_reconnect_stop_deadline_does_not_wait_for_cancel_cleanup(monkeypa
         await asyncio.gather(recovery, return_exceptions=True)
 
 
+@pytest.mark.asyncio
+async def test_reconnect_drain_survives_cancellation_resistant_close(monkeypatch):
+    """A cancellation-resistant polling-pool close must not wedge the ladder.
+
+    Distinct from ``test_reconnect_continues_if_drain_hangs``: that test's
+    ``_hang`` is cancellable, so the pre-existing ``asyncio.wait_for`` bound
+    also releases. httpcore's pool close runs under ``AsyncShieldCancellation``
+    (#58236/#63309) — a close that shields its cleanup keeps ``wait_for``
+    pending forever even after the timeout fires, wedging the tracked
+    ``_polling_error_task`` and every escalation gate behind it. The drain
+    must use the wall-clock deadline helper (abandon, not cancel-await),
+    same primitive as the general-pool drain (#98094).
+    """
+    adapter = _make_adapter()
+    adapter._polling_network_error_count = 1
+
+    mock_app, mock_polling_req = _make_mock_app()
+    release_close = asyncio.Event()
+    close_cancelled = asyncio.Event()
+
+    async def _shielded_close():
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            close_cancelled.set()
+            # Cancellation-resistant cleanup: swallows the cancel and waits.
+            await release_close.wait()
+            raise
+
+    mock_polling_req.shutdown = AsyncMock(side_effect=_shielded_close)
+    mock_polling_req.initialize = AsyncMock()
+    adapter._app = mock_app
+
+    monkeypatch.setattr(tg_adapter, "_DRAIN_TIMEOUT", 0.05, raising=False)
+    with patch("asyncio.sleep", new_callable=AsyncMock):
+        recovery = asyncio.create_task(
+            adapter._handle_polling_network_error(Exception("Timed out"))
+        )
+        done, _ = await asyncio.wait({recovery}, timeout=2)
+
+    try:
+        assert recovery in done, (
+            "reconnect remained blocked waiting for cancellation-shielded "
+            "polling-pool shutdown() cleanup"
+        )
+        assert close_cancelled.is_set(), "drain must have cancelled the close"
+        # Ladder still advanced: polling pool rebuilt and polling restarted.
+        mock_polling_req.initialize.assert_awaited_once()
+        mock_app.updater.start_polling.assert_awaited_once()
+        assert adapter._polling_error_task is None or adapter._polling_error_task.done()
+    finally:
+        release_close.set()
+        if not recovery.done():
+            recovery.cancel()
+        await asyncio.gather(recovery, return_exceptions=True)
+
+
 @pytest.mark.asyncio
 async def test_heartbeat_force_escalates_wedged_recovery_task(monkeypatch):
     """#66377: the heartbeat is an independent, cause-agnostic watchdog.