ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): pool-reuse 401 — entry-owned carrier + race-and-cancel

Browse Source 
Two pre-existing defects in the Phase 6 pool dispatch path that only
manifest when a pooled session is reused for a second dispatch:

1. The per-dispatch _AuthCapture allocated in _dispatch_pool was wired
   into the httpx response hook only at first connect (via
   _connect_one_pool). On a reused session no fresh connect runs, so
   the hook continues writing to the original-connect's carrier while
   the new dispatch inspects an empty carrier — auth_401/403 silently
   misclassified to "other", refresh-and-retry never fires.

2. Even with the carrier on the entry (so the hook writes to a stable
   reachable object), session.call_tool itself hangs forever on
   upstream 4xx for reused sessions. Trace: SDK's spawned
   handle_request_async raises HTTPStatusError, the outer
   streamablehttp_client TaskGroup cancels post_writer, post_writer's
   finally aclose's read_stream_writer, BaseSession's _receive_loop
   exits and enters its CONNECTION_CLOSED-fanout finally. anyio's
   send_nowait skips waiting receivers with pending_cancellation; the
   dispatch task (created by run_coroutine_threadsafe for the reuse
   case) is NOT in any cancel-scope chain, so the send "delivers" but
   the receiver's Event is set on stale state — receive() never
   wakes. Test 21 doesn't hit this because its 401 happens during
   initialize, in the same task that opens streamablehttp_client, so
   the cancel scope DOES propagate.

Fix:
- Move _AuthCapture ownership to PoolEntryState (and asyncio.Event
  alongside, allocated lazily on the mcp-loop). The hook closes over
  entry.auth_capture at first connect and stays valid across
  dispatches; reset under open_lock before each call_tool.
- Race session.call_tool against the carrier's fired_event in
  _dispatch_pool_with_entry. If the event wins (hook captured 4xx
  before SDK propagated), cancel call_tool and raise an internal
  _CarrierAuthSignal — _classify_failure resolves to auth_401/403
  via the carrier's status, the dispatcher evicts the broken
  session, and the cross-task retry handshake reconnects on a fresh
  bearer.

Adds tests/test_mcp_pool_auth_integration.py::test_integration_pool_reuse_401_refresh_and_retry_succeeds
which drives the reuse path through real upstream + real SDK and is
the structural gate against this class regressing. Negative-tested
twice: revert PoolEntryState.auth_capture → test fails (carrier
empty); revert the race → test times out (SDK hang).

Also drops the @pytest.mark.asyncio decorator (replaced with
@pytest.mark.anyio) on four tests in test_mcp_pool_auth_introspection.py.
The project depends on anyio's pytest plugin (anyio is in deps);
pytest-asyncio is NOT a project dep and CI's test (3.13) failed on
those four. Local pytest happened to pick it up via system Python.

Found via Copilot review on PR #481.
This commit is contained in:
Patrick Buckley
2026-05-05 23:28:52 -07:00
parent db9260d8c4
commit 97086fc617
5 changed files with 327 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+1 -1
View File
@@ -24,7 +24,7 @@ classifiers = [
dependencies = [
"openai>=2.24",
"httpx>=0.28",
"mcp>=1.6",
"mcp>=1.27",
"starlette>=0.45",
"uvicorn>=0.34",
"sse-starlette>=2.0",
tests/test_mcp_pool_auth_integration.py
+127
View File
@@ -716,3 +716,130 @@ def test_integration_static_path_unaffected(
assert state_after.session is not session_before, (
"Reconnect did not actually replace the session"
)
# ---------------------------------------------------------------------------
# Test 28: pool reuse — 401 on a SECOND dispatch (carrier owned by entry)
# ---------------------------------------------------------------------------
def test_integration_pool_reuse_401_refresh_and_retry_succeeds(
upstream: Any, running_loop_mgr: Any, storage: SQLiteBackend
) -> None:
"""Reused pool sessions still capture 401 correctly.
Dispatch 1 hits a passthrough upstream (200) and populates
``entry.session``. Dispatch 2 reuses that session — no fresh
connect, so a per-dispatch ``_AuthCapture`` would never reach
the httpx response hook (the hook closes over the carrier passed
at first connect, which lives on the entry). A correctly-wired
entry-owned carrier is the only shape that lets dispatch 2's 401
surface to the dispatcher.
Two independent production bugs gate this test passing; both must
hold for reused-session 401 recovery to work end-to-end:
1. The carrier must live on the pool entry (not per-dispatch) so
the response hook bound at first connect writes to the same
object the dispatcher reads across reuse. Verified by reverting
``PoolEntryState.auth_capture`` to a per-dispatch
``_AuthCapture()`` allocation: the carrier-fired event never
reaches the dispatcher and the test times out.
2. The dispatcher must race ``call_tool`` against the carrier's
fired event. The SDK's ``_receive_loop`` runs in BaseSession's
TaskGroup nested inside ``streamablehttp_client``'s TaskGroup;
when an upstream 4xx fires, the outer TaskGroup cancels
``_receive_loop`` mid-finally before it can deliver
``CONNECTION_CLOSED`` to the response stream's waiting
receiver. anyio's ``send_nowait`` skips waiters with pending
cancellation — but our dispatch task (created via
``run_coroutine_threadsafe`` for the reused-session case) has
NO pending cancellation, so the send delivers but the receiver
never wakes (the waiter's Event is set on stale state). Result:
a forever-hung ``response_stream_reader.receive()``. Verified
by reverting the ``asyncio.wait({call_task, fired_task})``
race in ``_dispatch_pool_with_entry`` to a bare ``await
session.call_tool(...)``: the test times out.
This test is the structural gate against the per-dispatch carrier
pattern: it looks right in code review and passes single-dispatch
integration tests, but breaks silently on session reuse — and the
SDK-level hang the carrier fix exposes silently strands the
dispatcher even when the carrier is correct.
"""
url, behaviour = upstream
behaviour["mode"] = "never" # passthrough — dispatch 1 succeeds
mgr, _loop, _ = running_loop_mgr
cipher = make_mcp_token_cipher()
_seed_oauth_server(storage, name="pool-srv", url=url)
_seed_user_token(storage, cipher)
mgr.set_storage(storage)
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
async def _fake_classified(**kwargs: Any) -> TokenLookupResult:
if kwargs.get("force_refresh"):
return TokenLookupResult(kind="token", token="refreshed-bearer")
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
):
# Dispatch 1: passthrough success. Establishes the pooled session.
result1 = mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "first"}, user_id="user-1", timeout=15
)
assert "echoed:first" in result1
entry = mgr._user_pool_entries[("user-1", "pool-srv")]
session_after_first = entry.session
assert session_after_first is not None, (
"test setup: dispatch 1 did not populate entry.session; "
"subsequent dispatch will not exercise the reuse path"
)
# Reconfigure upstream to 401 once on the next call. Reset the
# auth-headers log so we can count dispatch-2's POSTs cleanly.
behaviour["post_auth_headers"] = []
behaviour["mode"] = "once_401"
behaviour["_fired"] = False
# Dispatch 2: same (user, server). The hook from dispatch 1's
# connect is still bound to entry.auth_capture. The 401 fires;
# the dispatcher's auth_401 path triggers refresh-and-retry.
result2 = mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "second"}, user_id="user-1", timeout=15
)
assert "echoed:second" in result2, (
f"reused-session 401 retry did not succeed. result: {result2!r}. "
"If this is JSON with mcp_consent_required, the dispatcher "
"fell through to consent_required emission; if a generic "
"tool error, the carrier was empty (auth branch unreachable)."
)
assert mgr._consecutive_failures.get("pool-srv", 0) == 0, (
"auth failures must not trip the per-server breaker"
)
# Dispatch 2 produces multiple POSTs: the original 401 with the
# rejected bearer, then the retry's full connect handshake
# (initialize + notifications/initialized + tools/list) followed by
# the actual tools/call — all under the refreshed bearer. The retry
# reconnects because the auth_401 handler evicted the broken session.
post_headers = behaviour.get("post_auth_headers", [])
assert len(post_headers) >= 2, (
f"expected >=2 POSTs after dispatch 2 (401 + retry); "
f"got {len(post_headers)}: {post_headers}"
)
# First POST is the original bearer that got 401'd.
assert post_headers[0] == "Bearer access-aaa", (
f"first POST was {post_headers[0]!r}; expected the original bearer"
)
# Every subsequent POST carries the refreshed bearer (the retry
# ran with force_refresh=True and reconnected with the new token).
refreshed = post_headers[1:]
assert all(h == "Bearer refreshed-bearer" for h in refreshed), (
f"retry POSTs carried unexpected bearer(s); observed: {post_headers}"
)
tests/test_mcp_pool_auth_introspection.py
+74 -51
View File
@@ -177,25 +177,37 @@ class TestClassifyFailureWithCapture:
# ---------------------------------------------------------------------------
# Tests 9-10: carrier per-dispatch isolation + hook only fires on 4xx
# Tests 9-10: carrier field-reset isolation + hook only fires on 4xx
# ---------------------------------------------------------------------------
class TestCarrierLifecycle:
def test_capture_resets_per_dispatch(self, running_loop_mgr, storage: SQLiteBackend) -> None:
"""Each ``_dispatch_pool`` allocates a fresh ``_AuthCapture`` so a
prior dispatch's 401 cannot leak into the next dispatch.
"""The carrier's fields are reset before each ``call_tool`` so a
prior dispatch's 401 cannot leak into the next dispatch's
classification.
The autouse ``_install_capture_intercept`` fixture stashes the
per-dispatch carrier on the manager. Two consecutive dispatches
must observe DIFFERENT carrier identities — if the dispatcher
reused one capture, the second dispatch would see the same
``id()``.
The carrier is owned by the pool entry (the httpx response hook
closes over ``entry.auth_capture`` at first connect; a per-
dispatch carrier would never reach the hook on a reused
session). Object identity is
therefore expected to be the SAME across dispatches; what
matters is that the FIELDS are reset under ``open_lock``
before ``call_tool`` runs.
Verified by reverting ``_dispatch_pool`` to allocate
``_AuthCapture`` once (e.g., set it on ``self._shared_capture``
in ``__init__`` and reuse) and confirming this test fails
because the two dispatches observe the same carrier identity.
Test shape:
1. Dispatch 1's stub ``call_tool`` populates the carrier with a
401 + WWW-Authenticate.
2. Dispatch 2's stub ``call_tool`` records the carrier's state
as observed AT CALL ENTRY — before this stub writes anything.
3. Assert the recorded state is reset (status=None,
www_authenticate=None) — proves dispatch 1's payload did
not leak.
Verified by reverting ``_dispatch_pool_with_entry`` to remove
the two reset assignments under ``open_lock`` and confirming
this test fails because dispatch 2 observes dispatch 1's
leaked 401.
"""
from unittest.mock import patch
@@ -208,12 +220,21 @@ class TestCarrierLifecycle:
mgr.set_storage(storage)
mgr.set_app_state(_make_app_state(storage, cipher=cipher))
observed_capture_ids: list[int] = []
# Dispatch 1's stub leaves a 401 on the carrier; dispatch 2's
# stub records carrier state at entry.
observed_states: list[tuple[int | None, str | None]] = []
call_index = [0]
async def _call_tool(name: str, args: dict[str, Any]) -> Any:
cap = getattr(mgr, "_test_active_capture", None)
if cap is not None:
observed_capture_ids.append(id(cap))
assert cap is not None, "test setup error: capture not stashed"
observed_states.append((cap.status, cap.www_authenticate))
if call_index[0] == 0:
# Simulate dispatch 1 seeing a 401 — leak this into
# the carrier so dispatch 2 must explicitly reset.
cap.status = 401
cap.www_authenticate = 'Bearer error="invalid_token"'
call_index[0] += 1
content = MagicMock()
content.text = "ok"
res = MagicMock()
@@ -239,10 +260,17 @@ class TestCarrierLifecycle:
mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5)
mgr.call_tool_sync("mcp__pool-srv__do_thing", {}, user_id="user-1", timeout=5)
assert len(observed_capture_ids) == 2
assert observed_capture_ids[0] != observed_capture_ids[1], (
"Both dispatches saw the same _AuthCapture object; the dispatcher "
"is reusing the carrier across calls instead of allocating fresh."
assert len(observed_states) == 2
# Dispatch 1 saw the freshly-constructed carrier (None/None) —
# initial state from PoolEntryState's default_factory.
assert observed_states[0] == (None, None), (
f"dispatch 1 saw stale carrier state: {observed_states[0]}"
)
# Dispatch 2 must observe the reset — dispatch 1's leaked 401
# is gone.
assert observed_states[1] == (None, None), (
f"dispatch 2 leaked dispatch 1's state: {observed_states[1]}; "
"_dispatch_pool_with_entry's reset of entry.auth_capture is broken"
)
def test_capture_dataclass_isolated_when_constructed_separately(self) -> None:
@@ -257,7 +285,7 @@ class TestCarrierLifecycle:
cap_b.status = 403
assert cap_a.status == 401
@pytest.mark.asyncio
@pytest.mark.anyio
async def test_response_hook_captures_4xx_only(self) -> None:
"""Hook records on 401 and 403 only; 200/201/202/500 are ignored.
@@ -310,7 +338,7 @@ class TestCarrierLifecycle:
class TestForceRefresh:
@pytest.mark.asyncio
@pytest.mark.anyio
async def test_force_refresh_bypasses_token_freshness_check(
self, storage: SQLiteBackend
) -> None:
@@ -358,7 +386,7 @@ class TestForceRefresh:
assert forced.kind == "token"
assert forced.token == "refreshed-token"
@pytest.mark.asyncio
@pytest.mark.anyio
async def test_force_refresh_collapses_concurrent_callers_via_lock(
self, storage: SQLiteBackend
) -> None:
@@ -439,7 +467,7 @@ class TestForceRefresh:
f"observed {call_count} round-trips."
)
@pytest.mark.asyncio
@pytest.mark.anyio
async def test_force_refresh_goes_through_pg_refresh_lock_path(
self, storage: SQLiteBackend
) -> None:
@@ -543,6 +571,7 @@ class TestDispatcherAuthFlows:
access_token: str,
*,
auth_capture: Any = None,
auth_fired_event: Any = None,
) -> Any:
entry = await self_inner._ensure_pool_entry(key)
sess = MagicMock()
@@ -586,20 +615,13 @@ class TestDispatcherAuthFlows:
return TokenLookupResult(kind="token", token="refreshed-bearer")
return TokenLookupResult(kind="token", token="access-aaa")
# The first call_tool populates the carrier with 401; the
# second call_tool returns success.
# The first call_tool populates the entry-owned carrier with
# 401 (via _populate_active_capture, simulating what the
# production response hook would record); the second call_tool
# returns success.
async def _call_tool(name: str, args: dict[str, Any]) -> Any:
nonlocal call_count
call_count += 1
# Find the calling _dispatch_pool's capture by walking the
# most recently allocated entry session — the test bridges
# via the dispatch path, so we mutate the per-dispatch
# carrier directly via the patched factory's hook semantics.
# Easier: use a sentinel exception that the dispatcher's
# capture pre-populates via the hook test path. But for
# the dispatcher unit test we drive the carrier through
# an exception-injection that ALSO populates the captured
# carrier — see the test fixtures below.
if call_count == 1:
# Simulate the SDK swallow path: carrier was populated
# by the hook, but the exception that surfaces is a
@@ -905,6 +927,7 @@ class TestBreakerInvariant:
access_token: str,
*,
auth_capture: Any = None,
auth_fired_event: Any = None,
) -> Any:
entry = await self_inner._ensure_pool_entry(key)
sess = MagicMock()
@@ -1225,34 +1248,30 @@ class TestRetryTimeoutBudget:
# ---------------------------------------------------------------------------
# Test helpers — bridge the test stand-in for the real httpx event hook.
#
# The dispatcher allocates a fresh ``_AuthCapture`` per dispatch and threads
# it into ``_dispatch_pool_with_entry`` via the ``auth_capture`` kwarg. To
# unit-test the dispatcher's reaction without a real upstream, the fake
# ``call_tool`` populates the carrier directly via this helper, simulating
# what the production response hook would do.
# The dispatcher consults ``entry.auth_capture`` (carrier owned by the pool
# entry, persisted across dispatches; the production response hook closes
# over it at first connect). To unit-test the dispatcher's reaction without
# a real upstream, the fake ``call_tool`` populates that carrier directly
# via this helper, simulating what the production response hook would do.
# ---------------------------------------------------------------------------
def _populate_active_capture(mgr: MCPClientManager, *, status: int, header: str) -> None:
"""Mutate the per-dispatch ``_AuthCapture`` stashed on the manager
"""Mutate the entry-owned ``_AuthCapture`` stashed on the manager
by the autouse ``_install_capture_intercept`` fixture.
Used by stub ``call_tool`` implementations to simulate what the
production response hook would record on a 4xx upstream response.
Raises ``RuntimeError`` if invoked outside the autouse fixture's
intercept window — the carrier is per-dispatch, not per-test, so
it only exists while a dispatch is in flight.
intercept window — the helper only resolves the carrier while a
dispatch is in flight.
"""
capture = getattr(mgr, "_test_active_capture", None)
if capture is None:
# Allocate a fresh sentinel that the dispatcher's auth-classification
# branch will see. The fake call_tool mutates THIS object; the
# dispatcher consults its own per-dispatch capture, but since we
# also patch _classify_failure to consult the test sentinel...
raise RuntimeError(
"test setup error: _populate_active_capture called before "
"_install_capture_intercept; the dispatcher's per-dispatch "
"_AuthCapture isn't observable from this stub."
"_install_capture_intercept; the entry's _AuthCapture isn't "
"observable from this stub."
)
capture.status = status
capture.www_authenticate = header
@@ -1261,7 +1280,7 @@ def _populate_active_capture(mgr: MCPClientManager, *, status: int, header: str)
@pytest.fixture(autouse=True)
def _install_capture_intercept(monkeypatch: pytest.MonkeyPatch) -> None:
"""Wrap ``_dispatch_pool_with_entry`` so the test's fake ``call_tool``
can populate the per-dispatch ``_AuthCapture`` via ``mgr._test_active_capture``.
can populate the entry-owned ``_AuthCapture`` via ``mgr._test_active_capture``.
The wrap is a no-op for tests that don't call ``_populate_active_capture``;
they simply never read the attribute. Dispatcher-asserting tests rely
@@ -1273,8 +1292,12 @@ def _install_capture_intercept(monkeypatch: pytest.MonkeyPatch) -> None:
original = mcp_client_mod.MCPClientManager._dispatch_pool_with_entry
async def _wrapped(self: MCPClientManager, **kwargs: Any) -> str:
# Stash the carrier so the test's call_tool stub can populate it.
self._test_active_capture = kwargs.get("auth_capture") # type: ignore[attr-defined]
# Stash the entry's persistent carrier so the test's call_tool
# stub can populate it.
entry = kwargs.get("entry")
self._test_active_capture = ( # type: ignore[attr-defined]
entry.auth_capture if entry is not None else None
)
try:
return await original(self, **kwargs)
finally:
turnstone/core/mcp_client.py
+124 -23
View File
@@ -139,7 +139,21 @@ class _PoolDispatchRetryRequested(BaseException): # noqa: N818
"""
def _make_capturing_http_factory(capture: _AuthCapture) -> McpHttpClientFactory:
class _CarrierAuthSignal(Exception): # noqa: N818
"""Internal signal raised when the response hook captures a 4xx mid-call.
Raised from ``_dispatch_pool_with_entry`` when the carrier's fired
event wins the race against ``session.call_tool``. The dispatcher's
``_classify_failure(exc, capture=...)`` resolves the actual auth
class (``auth_401`` / ``auth_403``) from the carrier's ``status``,
so this exception is just a structural placeholder — it never
surfaces to callers.
"""
def _make_capturing_http_factory(
capture: _AuthCapture, fired_event: asyncio.Event | None = None
) -> McpHttpClientFactory:
"""Return an ``httpx`` factory that records 4xx auth signals into ``capture``.
The hook is ``async`` because :class:`httpx.AsyncClient` invokes
@@ -176,6 +190,12 @@ def _make_capturing_http_factory(capture: _AuthCapture) -> McpHttpClientFactory:
capture.status = status
headers = response.headers.get_list("www-authenticate")
capture.www_authenticate = headers[0] if headers else None
if fired_event is not None:
# Wakes the dispatcher's race in ``_dispatch_pool_with_entry``.
# See the docstring on _CarrierAuthSignal for why call_tool
# cannot be relied on to propagate the failure on a reused
# session.
fired_event.set()
def _factory(
headers: dict[str, str] | None = None,
@@ -269,6 +289,14 @@ class PoolEntryState:
eviction skips entries with ``in_flight > 0`` so a long-running
``call_tool`` can release ``open_lock`` while still pinning the
entry's session against teardown.
``auth_capture`` is bound once at first connect (passed to the
httpx response hook factory); the hook closes over this object for
the life of the underlying ``httpx.AsyncClient``. Each dispatch
resets the carrier's fields under ``open_lock`` before
``call_tool`` and reads them after — a per-dispatch carrier would
fail silently on session reuse because the hook is bound at
connect time, not per dispatch.
"""
key: tuple[str, str] # (user_id, server_name)
@@ -284,6 +312,13 @@ class PoolEntryState:
prompts: list[dict[str, Any]] | None = None
last_used: float = 0.0
in_flight: int = 0
auth_capture: _AuthCapture = field(default_factory=_AuthCapture)
# Set by the response hook when the carrier captures a 4xx; awaited
# by ``_dispatch_pool_with_entry``'s race against ``call_tool``.
# Must be allocated on the mcp-loop (per :class:`asyncio.Event`'s
# loop-binding contract); ``_ensure_pool_entry`` runs on the loop
# so the dataclass default_factory is safe.
auth_fired_event: asyncio.Event = field(default_factory=asyncio.Event)
# ---------------------------------------------------------------------------
@@ -924,6 +959,7 @@ class MCPClientManager:
access_token: str,
*,
auth_capture: _AuthCapture | None = None,
auth_fired_event: asyncio.Event | None = None,
) -> PoolEntryState:
"""Connect a single per-(user, server) pool entry.
@@ -984,7 +1020,9 @@ class MCPClientManager:
client_kwargs: dict[str, Any] = {"url": url, "headers": headers}
if auth_capture is not None:
client_kwargs["httpx_client_factory"] = _make_capturing_http_factory(auth_capture)
client_kwargs["httpx_client_factory"] = _make_capturing_http_factory(
auth_capture, fired_event=auth_fired_event
)
stack = AsyncExitStack()
await stack.__aenter__()
@@ -2190,9 +2228,10 @@ class MCPClientManager:
Pool-path 401/403 handling: the SDK's ``post_writer`` swallows
``httpx.HTTPStatusError``; we recover the upstream auth signal
via a response-hook carrier on the per-dispatch
``httpx.AsyncClient``. A 401 triggers one refresh-and-retry
with ``force_refresh=True``; persistent 401 emits
via a response-hook carrier owned by the pool entry (bound to
the entry's persistent ``httpx.AsyncClient`` at first connect;
see ``PoolEntryState.auth_capture``). A 401 triggers one
refresh-and-retry with ``force_refresh=True``; persistent 401 emits
``mcp_consent_required``. A 403 with
``WWW-Authenticate: error="insufficient_scope"`` emits
``mcp_insufficient_scope`` with the parsed scope set. Other
@@ -2529,13 +2568,10 @@ class MCPClientManager:
key = (user_id, server_name)
entry = await self._ensure_pool_entry(key)
# First attempt — fresh capture per dispatch so a prior
# dispatch's 401 cannot leak into this one's classification.
# Even though ``open_lock`` is held across ``call_tool``,
# allocating per-call (rather than per-entry) protects against
# future concurrent-multiplex regressions if the lock scope
# ever shrinks.
capture = _AuthCapture()
# See PoolEntryState.auth_capture for why the carrier is
# entry-owned; _dispatch_pool_with_entry resets it under
# open_lock before call_tool, we read it after.
capture = entry.auth_capture
try:
result = await self._dispatch_pool_with_entry(
entry=entry,
@@ -2544,7 +2580,6 @@ class MCPClientManager:
access_token=access_token,
original_name=original_name,
arguments=arguments,
auth_capture=capture,
)
except BaseException as exc:
classification = self._classify_failure(exc, capture=capture)
@@ -2660,18 +2695,22 @@ class MCPClientManager:
access_token: str,
original_name: str,
arguments: dict[str, Any],
auth_capture: _AuthCapture,
) -> str:
"""Hold ``entry.open_lock`` across connect-or-reuse AND ``call_tool``.
Lock held across ``call_tool`` because the per-dispatch
Lock held across ``call_tool`` because the entry's
``_AuthCapture`` is keyed off the httpx event hook; releasing
would let a concurrent same-(user, server) dispatch overwrite
the carrier mid-flight, attributing one caller's auth failure
to another. Holding the lock serialises same-(user, server)
calls — acceptable because that contention scenario is rare
in practice (ChatSession dispatches sequentially and per-
(user, server) parallelism is not a production requirement).
the lock would let a concurrent same-(user, server) dispatch
overwrite the carrier mid-flight and attribute one caller's
auth failure to another. Holding the lock serialises
same-(user, server) calls — acceptable because that contention
scenario is rare in practice (ChatSession dispatches
sequentially and per-(user, server) parallelism is not a
production requirement).
Reset of ``entry.auth_capture`` and ``entry.auth_fired_event``
happens under the lock before ``call_tool`` — see
:class:`PoolEntryState` for why the carrier is entry-owned.
``entry.in_flight`` accounting is preserved for the eviction
interlock — it's belt-and-braces here since ``open_lock.locked()``
@@ -2681,18 +2720,80 @@ class MCPClientManager:
async with entry.open_lock:
entry.last_used = time.monotonic()
self._user_pool_last_used[key] = entry.last_used
entry.auth_capture.status = None
entry.auth_capture.www_authenticate = None
entry.auth_fired_event.clear()
session = entry.session
if session is None:
# Lazy connect — also covers post-eviction recovery.
fresh = await self._connect_one_pool(
key, cfg, access_token, auth_capture=auth_capture
key,
cfg,
access_token,
auth_capture=entry.auth_capture,
auth_fired_event=entry.auth_fired_event,
)
session = fresh.session
if session is None:
raise RuntimeError(f"Pool connect for {key!r} produced no session")
entry.in_flight += 1
try:
result = await session.call_tool(original_name, arguments)
# Race ``call_tool`` against the carrier's fired event.
# Without this race, an upstream 4xx on a REUSED session
# never propagates back through ``call_tool``: the SDK's
# ``_receive_loop`` is in BaseSession's TaskGroup, nested
# inside ``streamablehttp_client``'s TaskGroup. When
# the spawned ``handle_request_async`` task raises
# ``HTTPStatusError``, the outer TaskGroup cancels
# ``_receive_loop`` mid-finally, before it can deliver
# ``CONNECTION_CLOSED`` to the response stream's waiting
# receiver. anyio's ``send_nowait`` skips waiters with
# pending cancellation; here the dispatcher's task has
# NO pending cancellation (it was created by a fresh
# ``run_coroutine_threadsafe`` and is not in the
# streamablehttp_client cancel-scope chain), so the
# send delivers but the receiver never wakes — the
# waiter's Event is set on a stale state. Result: a
# forever-hung ``response_stream_reader.receive()``.
# The carrier-fired event lets us short-circuit before
# the SDK's hang manifests.
call_task = asyncio.create_task(session.call_tool(original_name, arguments))
fired_task = asyncio.create_task(entry.auth_fired_event.wait())
try:
done, pending = await asyncio.wait(
{call_task, fired_task},
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
# ``_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
# ``CancelledError`` we asked for and any
# ``BaseExceptionGroup`` the SDK's anyio
# TaskGroup may wrap on teardown.
for task in (call_task, fired_task):
if not task.done():
task.cancel()
for task in (call_task, fired_task):
with contextlib.suppress(BaseException):
await task
if call_task in done:
result = call_task.result()
else:
# Hook captured 4xx before call_tool returned. The
# SDK won't propagate the failure through
# call_tool, so eagerly tear down the session and
# raise a sentinel that the dispatcher's
# ``_classify_failure`` will resolve via the
# carrier (which holds the captured status).
raise _CarrierAuthSignal()
finally:
entry.in_flight -= 1
return _decode_tool_result(result)
uv.lock
Generated
+1 -1
View File
@@ -2623,7 +2623,7 @@ requires-dist = [
{ name = "httpx", specifier = ">=0.28" },
{ name = "httpx-sse", specifier = ">=0.4" },
{ name = "lacme", marker = "extra == 'tls'", specifier = ">=1.0.5" },
{ name = "mcp", specifier = ">=1.6" },
{ name = "mcp", specifier = ">=1.27" },
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.14" },
{ name = "numpy", marker = "extra == 'sandbox'", specifier = ">=2.0" },
{ name = "openai", specifier = ">=2.24" },