fix(mcp): apply PR #489 review feedback + de-flake pool reuse 401 retry

PR #489 review feedback (Copilot + github-code-quality):
- closeSettingsPanel now closes nested revoke modal first on close-button
  path (Escape was already handled by the parent keydown trap deferring
  to the inner trap; missing-modal-on-close-button was an orphan-modal
  hazard).
- _refreshConsentBadge now updates the settings button's aria-label +
  title dynamically with the pending-consent count for screen readers
  (badge stays aria-hidden — the count is in the label).
- _MAX_INSUFFICIENT_SCOPE_REPORTED promoted to public
  MAX_INSUFFICIENT_SCOPE_REPORTED in mcp_http_parsers; drops cross-module
  private import in mcp_oauth's /start handler.
- Stale test comment in test_session_mcp_dispatch_error.py corrected:
  _exec_read_resource does not log with exc_info=True (bearer-leak
  invariant).
- Rejected the protocol-method ellipsis warning: rest of _protocol.py
  uses ... consistently per Protocol convention.

Lint:
- ruff format applied to test_mcp_pool_auth_integration.py and
  test_mcp_pool_auth_resource_integration.py (combined `with` grammar —
  pure formatting).

Flake fix — test_integration_pool_reuse_401_refresh_and_retry_succeeds
on Python 3.11 / resource-constrained CI:

Same cross-task scope hazard f6a3b66 fixed at the close side, surfacing
at the connect side. asyncio.wait_for at mcp_client.py:1206 wraps
streamablehttp_client.__aenter__ in a fresh asyncio.Task. That fresh
task enters anyio cancel scopes, completes, and dies. The eventual
stack.aclose() during eviction or auth_401 retry runs from a different
task and tries to exit scopes whose entering task is dead — anyio
raises RuntimeError, the wedged anyio state blocks the retry's stack
teardown + reconnect, and the call exceeds the 15s budget on slow
workers.

Fix: replace asyncio.wait_for with `async with asyncio.timeout(...)` so
the streamablehttp_client.__aenter__ runs in the dispatch task itself,
no fresh-task scope ownership. Aligns with invariant 18 (asyncio.timeout
not asyncio.wait_for for any SDK / AS / pool-loop await crossing anyio
scopes).

Static path (_connect_one) at lines 905 and 1000 deliberately retains
asyncio.wait_for — auth_type ∈ {none, static} is byte-identical
(invariant 1) and the narrow connect-once / no-eviction-then-reuse
pattern doesn't trigger the cross-task hazard. Anchor comments pin
both directions: a future migration there would break invariant 1; a
future revert at 1206 would re-introduce the flake.

The cited test is the symptom (non-deterministically times out under
load), not a structural gate (no deterministic asyncio.timeout
assertion exists). The comment block at line 1206 records this so a
maintainer who reverts and finds green on a fast machine doesn't
conclude the fix is unneeded.

Verified on Python 3.11.14 (/tmp/venv311) and 3.13.7 (.venv): ruff
format clean, ruff check clean, mypy clean. 368 unit tests + 30 pool
integration tests pass on both interpreters; the previously-flaky test
passed 20× in isolation on 3.11.

Multi-stage /review (4 finders × verify × dedupe): bug/security/perf
returned zero findings; quality returned 3 confirmed minor/nit items
all of which are applied here (q-1 anchor comments at 905+1000, q-2
symptom-vs-gate clarification at 1206, q-3 module-docstring sentence
in mcp_http_parsers).
This commit is contained in:
Patrick Buckley
2026-05-07 13:39:49 -07:00
parent 5a3f46a1fa
commit 4a3e3607be
8 changed files with 187 additions and 85 deletions
+1 -1
View File
@@ -391,7 +391,7 @@ class TestAuthorize:
app = _build_app(storage=storage, http_client=http_client_mock, token_store=token_store)
client = TestClient(app, raise_server_exceptions=False)
# 33 tokens — one over the cap (``_MAX_INSUFFICIENT_SCOPE_REPORTED = 32``).
# 33 tokens — one over the cap (``MAX_INSUFFICIENT_SCOPE_REPORTED = 32``).
scopes = "%20".join(f"s{i}" for i in range(33))
with _public_addr_patch():
resp = client.get(
+40 -35
View File
@@ -399,13 +399,14 @@ def test_integration_401_with_refresh_failure_emits_consent_required(
return TokenLookupResult(kind="refresh_failed")
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15
)
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
# Structured-error envelopes flow back via ``RuntimeError(json_str)``
# so the session-layer ``except Exception`` handler routes the
@@ -444,13 +445,14 @@ def test_integration_403_insufficient_scope_emits_structured_error(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15
)
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_insufficient_scope"
@@ -489,13 +491,14 @@ def test_integration_403_no_insufficient_scope_emits_generic_forbidden(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15
)
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_tool_call_forbidden"
@@ -563,13 +566,14 @@ def test_integration_403_multi_www_authenticate_drops_injected_scopes(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15
)
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_insufficient_scope", (
@@ -615,13 +619,14 @@ def test_integration_401_retry_ceiling(
return TokenLookupResult(kind="token", token=f"refreshed-{refresh_count}")
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
mgr.call_tool_sync(
"mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15
)
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.call_tool_sync("mcp__pool-srv__echo", {"payload": "x"}, user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_consent_required"
@@ -357,10 +357,13 @@ def test_resource_read_persistent_401_emits_consent_required(
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,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
@@ -392,10 +395,13 @@ def test_resource_read_403_insufficient_scope_emits_structured_error(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
@@ -430,10 +436,13 @@ def test_resource_read_403_generic_forbidden(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
@@ -476,10 +485,13 @@ def test_resource_read_breaker_unaffected_by_auth_failures(
# repopulates this; the test simulates that out-of-band.
for _ in range(10):
_seed_pool_resource_map(mgr, "user-1", "pool-srv", "res://hello")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=15)
payload = json.loads(str(exc_info.value))
assert payload["error"]["code"] == "mcp_consent_required"
@@ -505,10 +517,13 @@ def test_resource_read_missing_token_emits_consent_required(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="missing")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=10)
payload = json.loads(str(exc_info.value))
@@ -529,10 +544,13 @@ def test_resource_read_decrypt_failure_emits_token_undecryptable(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="decrypt_failure")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=10)
payload = json.loads(str(exc_info.value))
@@ -556,10 +574,13 @@ def test_resource_read_http_url_emits_url_insecure(
async def _fake_classified(**_kwargs: Any) -> TokenLookupResult:
return TokenLookupResult(kind="token", token="access-aaa")
with patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
), pytest.raises(RuntimeError) as exc_info:
with (
patch(
"turnstone.core.mcp_client.get_user_access_token_classified",
side_effect=_fake_classified,
),
pytest.raises(RuntimeError) as exc_info,
):
mgr.read_resource_sync("res://hello", user_id="user-1", timeout=5)
payload = json.loads(str(exc_info.value))
+4 -3
View File
@@ -165,9 +165,10 @@ class TestExecReadResourceDispatchError:
"call_id": "rc_1",
"resource_uri": "https://example.com/r",
}
# The exec site logs ``log.warning(... exc_info=True)`` on failure.
# Patch the logger so the test doesn't emit noise to the captured
# stderr — assertions don't depend on log output.
# The exec site emits a ``log.warning`` (no ``exc_info`` — bearer-leak
# invariant) on failure. Patch the logger so the test doesn't emit
# noise to the captured stderr — assertions don't depend on log
# output.
with patch("turnstone.core.session.log"):
session._exec_read_resource(item)
+45 -12
View File
@@ -44,6 +44,7 @@ from mcp.shared._httpx_utils import (
from turnstone.core.config import load_config
from turnstone.core.log import get_logger
from turnstone.core.mcp_http_parsers import (
MAX_INSUFFICIENT_SCOPE_REPORTED,
parse_www_authenticate_error,
parse_www_authenticate_scope,
)
@@ -113,18 +114,12 @@ def _validate_oauth_user_url(url: str) -> None:
# 1:1 against the new entry point when we migrate.
# Defensive cap on the number of scopes we report in
# ``mcp_insufficient_scope`` audit/error payloads. Real ASes return
# single-digit scope counts; the cap stops a malicious upstream from
# bloating either surface via a thousand-token ``scope=`` value.
_MAX_INSUFFICIENT_SCOPE_REPORTED = 32
# Defensive cap on the number of tools we accept from any single MCP
# server's ``tools/list`` response. Real servers expose at most a few
# dozen tools; a misconfigured or hostile upstream returning thousands
# would amplify both memory (one OpenAI tool dict per entry) and
# downstream BM25 reindex cost. Mirrors the ``_MAX_ERROR_LEN`` /
# ``_MAX_INSUFFICIENT_SCOPE_REPORTED`` defensive ceilings: we truncate
# ``MAX_INSUFFICIENT_SCOPE_REPORTED`` defensive ceilings: we truncate
# rather than reject so partial visibility beats zero visibility, and
# emit a warning so operators can investigate.
_MAX_TOOLS_PER_SERVER = 1000
@@ -907,6 +902,14 @@ class MCPClientManager:
# orphaned cancel-scope tasks spinning at 100% CPU.
await self._tcp_probe(name, cfg["url"])
# ``wait_for`` retained — the static path is byte-identical
# (invariant 1) and its narrow connect-once / no-eviction-then-
# reuse pattern doesn't trigger the cross-task scope-exit that
# f6a3b66 fixed for ``_safe_close_stack`` and that the pool
# path's ``_connect_one_pool`` migrated at line 1206 below.
# Migrating here would change the static path's bytes; reverting
# 1206 to match would re-introduce the pool-path flake. Pin both
# directions.
read, write, _ = await asyncio.wait_for(
stack.enter_async_context(
streamablehttp_client(url=cfg["url"], headers=cfg.get("headers"))
@@ -1002,6 +1005,8 @@ class MCPClientManager:
state.stack = stack
try:
# ``wait_for`` retained — see line 905 for the static-path
# exemption from the pool-path migration (invariant 1 byte-identical).
await asyncio.wait_for(session.initialize(), timeout=self._CONNECT_TIMEOUT)
except asyncio.CancelledError:
state.stack = None
@@ -1208,10 +1213,38 @@ class MCPClientManager:
await stack.__aenter__()
try:
await self._tcp_probe(key, url)
read, write, _ = await asyncio.wait_for(
stack.enter_async_context(streamablehttp_client(**client_kwargs)),
timeout=self._CONNECT_TIMEOUT,
)
# ``asyncio.timeout`` (NOT ``asyncio.wait_for``) — same anyio /
# Python 3.11 reasoning as ``session.initialize`` below and
# the ``_safe_close_stack`` fix in f6a3b66. ``wait_for``
# wraps the inner coroutine in a fresh :class:`asyncio.Task`
# which enters ``streamablehttp_client``'s anyio cancel
# scopes. When that fresh task completes (connect succeeded)
# its scopes are recorded on the stack but the entering
# task is dead; the eventual ``stack.aclose()`` during
# eviction (or auth_401 retry) tries to exit those scopes
# from a different task and anyio raises
# ``RuntimeError('Attempted to exit cancel scope in a
# different task...')`` — the same cross-task hazard
# f6a3b66 fixed at the close side.
#
# Surfacing: ``test_integration_pool_reuse_401_refresh_and_retry_succeeds``
# times out (non-deterministically) on Python 3.11 /
# resource-constrained CI when reverted — the wedged anyio
# state from dispatch 1's fresh-task connect blocks the
# auth_401 retry's stack teardown + reconnect within the
# 15s call budget. The test is the symptom, NOT a structural
# gate (its docstring at tests/test_mcp_pool_auth_integration.py:806-836
# asserts carrier-on-entry + race-against-fired-event, neither
# of which exercises this scope-ownership invariant). The
# structural argument is ``_safe_close_stack`` at line 834-846
# plus invariant 18 (``asyncio.timeout`` not ``asyncio.wait_for``
# for any SDK / AS / pool-loop await crossing anyio scopes).
# Reverting and finding the test green on a fast machine
# does NOT validate the revert.
async with asyncio.timeout(self._CONNECT_TIMEOUT):
read, write, _ = await stack.enter_async_context(
streamablehttp_client(**client_kwargs)
)
entry.streams = (read, write)
except asyncio.CancelledError:
task = asyncio.current_task()
@@ -4342,7 +4375,7 @@ class MCPClientManager:
error_token = parse_www_authenticate_error(header)
if error_token == "insufficient_scope":
scopes = parse_www_authenticate_scope(header)
scopes = scopes[:_MAX_INSUFFICIENT_SCOPE_REPORTED]
scopes = scopes[:MAX_INSUFFICIENT_SCOPE_REPORTED]
await emit_oauth_failure_audit(
app_state=self._app_state,
user_id=user_id,
+18
View File
@@ -23,6 +23,12 @@ returns a normalised ``{key.lower(): value}`` dict. The thin extraction
wrappers (``parse_www_authenticate_scope`` / ``parse_www_authenticate_error``)
preserve the original return shapes so call sites only need to swap the
import.
Also hosts ``MAX_INSUFFICIENT_SCOPE_REPORTED`` — the shared defensive cap
on scope-list lengths consumed by both the WWW-Authenticate parser
(``mcp_client``) and the ``/v1/api/mcp/oauth/start?scopes=`` step-up
handler (``mcp_oauth``). Living here avoids one module importing a
private name from the other.
"""
from __future__ import annotations
@@ -185,6 +191,18 @@ def parse_www_authenticate_bearer(header: str) -> dict[str, str]:
return out
# Defensive cap on the number of scopes reported in
# ``mcp_insufficient_scope`` audit/error payloads and accepted from the
# ``/v1/api/mcp/oauth/start?scopes=`` step-up query param. Real ASes
# return single-digit scope counts; the cap stops a malicious upstream
# (or buggy client) from bloating either surface via a thousand-token
# scope list. Lives here so the WWW-Authenticate parser (consumer:
# ``mcp_client``) and the ``/start`` handler (consumer: ``mcp_oauth``)
# share a single source of truth without one importing a private name
# from the other.
MAX_INSUFFICIENT_SCOPE_REPORTED = 32
def is_valid_scope_token(token: str) -> bool:
"""Return True iff ``token`` is a valid RFC 6749 §3.3 ``scope-token``.
+6 -6
View File
@@ -39,6 +39,7 @@ from turnstone.core.audit import record_audit
from turnstone.core.log import get_logger
from turnstone.core.mcp_crypto import MCPTokenDecryptError
from turnstone.core.mcp_http_parsers import (
MAX_INSUFFICIENT_SCOPE_REPORTED,
is_valid_scope_token,
parse_www_authenticate_bearer,
)
@@ -1928,9 +1929,10 @@ async def _handle_mcp_oauth_authorize_inner(request: Request) -> Response:
# Validated against the RFC 6749 §3.3 grammar so a malicious or buggy
# client can't smuggle CR/LF/tab/control bytes through the AS round-
# trip into downstream log or notification paths. The cap matches
# the per-call ceiling used in the WWW-Authenticate parser
# (``mcp_client._MAX_INSUFFICIENT_SCOPE_REPORTED``); over-capped
# input is rejected loudly so callers don't silently lose state.
# the per-call ceiling used in the WWW-Authenticate parser via the
# shared ``MAX_INSUFFICIENT_SCOPE_REPORTED`` constant in
# ``mcp_http_parsers``; over-capped input is rejected loudly so
# callers don't silently lose state.
#
# Splitting on a single space (NOT ``str.split()``) is intentional:
# Python's whitespace split would silently strip embedded CR/LF/tab,
@@ -1939,10 +1941,8 @@ async def _handle_mcp_oauth_authorize_inner(request: Request) -> Response:
requested_scopes_raw = request.query_params.get("scopes", "")
requested_scopes: list[str] = []
if requested_scopes_raw:
from turnstone.core.mcp_client import _MAX_INSUFFICIENT_SCOPE_REPORTED
candidates = [tok for tok in requested_scopes_raw.split(" ") if tok]
if len(candidates) > _MAX_INSUFFICIENT_SCOPE_REPORTED:
if len(candidates) > MAX_INSUFFICIENT_SCOPE_REPORTED:
return JSONResponse({"error": "Invalid scope token"}, status_code=400)
for tok in candidates:
if not is_valid_scope_token(tok):
+24
View File
@@ -5162,8 +5162,14 @@ function _refreshConsentBadge() {
if (!btn) return;
var existing = btn.querySelector(".settings-consent-badge");
var n = _pendingConsentServers.size;
// Keep the visible badge and the accessible name in lockstep so screen-
// reader users get the same pending-consent signal that sighted users
// get from the red dot. The badge itself stays aria-hidden because the
// count is already reflected in the button's aria-label/title.
if (n === 0) {
if (existing) existing.remove();
btn.setAttribute("aria-label", "MCP server connections");
btn.setAttribute("title", "MCP server connections");
return;
}
if (!existing) {
@@ -5173,6 +5179,14 @@ function _refreshConsentBadge() {
btn.appendChild(existing);
}
existing.textContent = String(n);
var label =
"MCP server connections (" +
n +
" pending consent" +
(n === 1 ? "" : "s") +
")";
btn.setAttribute("aria-label", label);
btn.setAttribute("title", label);
}
/**
@@ -5824,6 +5838,16 @@ function openSettingsPanel() {
}
function closeSettingsPanel() {
// If the nested revoke confirmation is still up, tear it down first
// — otherwise hiding the parent panel would leave an orphan modal
// overlay floating with its own keydown trap still attached. The
// Escape-key path inside the parent's keydown trap defers to the
// inner trap; this branch is the close-button path that doesn't go
// through that trap.
var inner = document.getElementById("revoke-mcp-overlay");
if (inner && inner.style.display !== "none") {
cancelRevokeMcp();
}
var overlay = document.getElementById("settings-overlay");
if (overlay) overlay.style.display = "none";
if (_settingsTrap) {