fix(mcp): asyncio.timeout (not wait_for) for safe-close-stack on Python 3.11

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.
This commit is contained in:
Patrick Buckley
2026-05-06 11:19:10 -07:00
parent 97086fc617
commit f6a3b66ea4
+20 -16
View File
@@ -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.