From f6a3b66ea4cc89cdd885380f468e0e09fbcb834a Mon Sep 17 00:00:00 2001
From: Patrick Buckley <eous@users.noreply.github.com>
Date: Wed, 6 May 2026 11:19:10 -0700
Subject: [PATCH] fix(mcp): asyncio.timeout (not wait_for) for safe-close-stack
 on Python 3.11
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Python 3.11's ``asyncio.wait_for`` wraps its inner coroutine in a fresh
``asyncio.Task`` via ``ensure_future``. When the inner is
``stack.aclose()`` on an ``AsyncExitStack`` containing
``streamablehttp_client(...)`` (anyio cancel scopes entered in the
calling task), the fresh task's attempt to exit those scopes raises
``RuntimeError('Attempted to exit cancel scope in a different task
than it was entered in')``. Python 3.12+ rewrote ``wait_for`` to use
``asyncio.timeout`` internally — runs in the current task — so 3.13
ran the same code path successfully.

Symptom on 3.11: integration tests where ``session.initialize()``
returns 4xx (e.g., 403 insufficient_scope tests) hit
``_connect_one_pool``'s ``except Exception:`` handler →
``_safe_teardown_on_connect_failure`` → ``_safe_close_stack`` → cross-
task RuntimeError. The ``concurrent.futures._base.CancelledError``
that surfaces in ``future.result(timeout=...)`` is the cascade
fallout from the asyncio loop's exception handler reacting to the
unretrieved-task-exception.

Fix: use ``asyncio.timeout`` instead of ``asyncio.wait_for`` for the
5s aclose bound. Equivalent semantics, current-task execution, works
on 3.11+. The 5s guard against ``aclose()`` hanging on a broken stack
is preserved.

Verified on Python 3.11.14 (full suite 5427 passed) and 3.13.7 (full
suite 5427 passed); all 9 integration tests pass on both.

Pre-existing bug — surfaced only after the marker fix in 5c9850c
let CI's test (3.11) actually run the 4xx tests.
---
 turnstone/core/mcp_client.py | 36 ++++++++++++++++++++----------------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py
index 24270efb..86b387bc 100644
--- a/turnstone/core/mcp_client.py
+++ b/turnstone/core/mcp_client.py
@@ -665,17 +665,23 @@ class MCPClientManager:
         the auth-retry path's eager teardown would propagate the SDK's
         own collected fallout.
 
-        The 5s ``asyncio.wait_for`` is a deliberate guard against
-        ``aclose()`` hanging on a broken stack (e.g., a never-completing
-        anyio task during teardown). The auth_401 retry runs on a
-        fresh :class:`asyncio.Task` scheduled via
-        :func:`asyncio.run_coroutine_threadsafe` (see
-        :meth:`_dispatch_pool_sync`), so this helper is never invoked
-        from inside an active cancellation context — ``wait_for``'s
-        own cancellation observation does not abort the aclose early.
+        The 5s timeout uses ``asyncio.timeout`` (NOT ``asyncio.wait_for``)
+        because Python 3.11's ``asyncio.wait_for`` wraps its inner
+        coroutine in a fresh :class:`asyncio.Task`. When that fresh
+        task runs ``stack.aclose()``, it tries to exit anyio cancel
+        scopes that were entered in the CALLING task — anyio rejects
+        the cross-task scope-exit with
+        ``RuntimeError('Attempted to exit cancel scope in a different
+        task than it was entered in')`` and the aclose fails.
+        ``asyncio.timeout`` runs the inner code in the current task
+        (matches Python 3.12+'s ``wait_for`` rewrite), preserving
+        scope-exit identity. The 5s bound stays as protection against
+        ``aclose()`` hanging on a broken stack (a never-completing
+        anyio task during teardown).
         """
         try:
-            await asyncio.wait_for(stack.aclose(), timeout=5)
+            async with asyncio.timeout(5):
+                await stack.aclose()
         except (Exception, asyncio.CancelledError, BaseExceptionGroup):
             log.debug("Error closing AsyncExitStack; ignoring", exc_info=True)
 
@@ -2765,16 +2771,14 @@ class MCPClientManager:
                         return_when=asyncio.FIRST_COMPLETED,
                     )
                 finally:
-                    # Cancel-and-await: matches the codebase's standard
-                    # cancel pattern (cf. shutdown helper). Awaiting
-                    # cancelled tasks here pins the broken session's
-                    # streams against the auth_401 retry's
+                    # Cancel-and-await both losers. Awaiting cancelled
+                    # tasks here pins the broken session's streams
+                    # against the auth_401 retry's
                     # ``_safe_close_stack`` teardown — without it the
                     # cancelled ``call_task`` could keep touching the
                     # SDK's stream state concurrently with the new
-                    # ``_connect_one_pool``'s aclose, racing exactly
-                    # the way the rest of this fix was written to
-                    # avoid. ``BaseException`` covers both the
+                    # ``_connect_one_pool``'s aclose.
+                    # ``BaseException`` covers both the
                     # ``CancelledError`` we asked for and any
                     # ``BaseExceptionGroup`` the SDK's anyio
                     # TaskGroup may wrap on teardown.