diff --git a/pyproject.toml b/pyproject.toml index 99031243..44cb52ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/tests/test_mcp_pool_auth_integration.py b/tests/test_mcp_pool_auth_integration.py index ea226c69..27d8ade6 100644 --- a/tests/test_mcp_pool_auth_integration.py +++ b/tests/test_mcp_pool_auth_integration.py @@ -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}" + ) diff --git a/tests/test_mcp_pool_auth_introspection.py b/tests/test_mcp_pool_auth_introspection.py index 351b4952..d8f94bca 100644 --- a/tests/test_mcp_pool_auth_introspection.py +++ b/tests/test_mcp_pool_auth_introspection.py @@ -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: diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 6c35e932..24270efb 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -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) diff --git a/uv.lock b/uv.lock index f4fa3521..7f0611d8 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },