fix(oauth): narrow context.lock scope in async_auth_flow

Closed
PE
4mo ago
+332-18

Narrows the lock scope inside OAuthClientProvider.async_auth_flow so the actual HTTP request yield runs outside any lock. Concurrent requests through the same provider — in particular, a POST issued while a GET SSE long-poll is in flight — no longer serialize on context.lock. A new refresh_lock preserves single-flight semantics for token refresh so the change does not introduce a refresh race.

Motivation and Context

OAuthClientProvider.async_auth_flow previously held self.context.lock for the entire flow body, including the response = yield request that hands the request off to httpx. For requests that complete quickly this is fine, but a Streamable HTTP transport opens a long-lived GET SSE stream for server-initiated messages alongside the POST stream for client-initiated RPCs. Both go through the same OAuthClientProvider instance, so the GET SSE long-poll's yield request holds context.lock for the entire SSE lifetime. Any concurrent tools/call POST then blocks on the lock until the SSE returns, producing a multi-second stall in OAuth-authenticated MCP connections.

This commit splits the single coarse lock into purpose-specific scopes:

  • Phase 1 (context.lock): initialize state, capture protocol_version, and decide whether a refresh is needed. Short-held; no HTTP I/O.
  • Phase 2 (refresh_lock, new): single-flight token refresh. The refresh yield happens outside any lock. A double-check inside context.lock ensures concurrent waiters do not redundantly refresh after another coroutine completed one.
  • Phase 3 (no lock): add the auth header and yield the actual request. GET SSE long-polls and other long-running requests no longer block unrelated traffic.
  • Phase 4 (context.lock): 401 / 403 full OAuth re-auth path. Conservatively kept under lock because this path is rare and its yielded sub-requests target the AS, not the resource server. A future change can narrow this further with the same pattern.

refresh_lock is added to OAuthContext (alongside the existing lock) so concurrent expirations trigger at most one refresh request. The double-check pattern inside refresh_lock makes the second waiter observe the freshly-updated token and skip its own refresh.

How Has This Been Tested?

Five new tests in tests/client/test_auth.py exercise the refactor:

  • test_concurrent_request_not_blocked_by_pending_long_running_request — the core regression test. Drives one async_auth_flow generator to its yield (simulating a GET SSE long-poll), then opens a second flow on the same provider and asserts it reaches its own yield within anyio.fail_after(5). Pre-fix this hangs indefinitely.
  • test_refresh_lock_double_check_skips_redundant_refresh — a flow started after another finished its refresh observes the fresh token in Phase 1 and skips Phase 2 entirely; no second refresh request is sent.
  • test_double_check_inside_refresh_lock_skips_second_refresh — monkeypatches is_token_valid to return False in Phase 1 and True in the Phase 2 double-check, exercising the branch where a second coroutine's refresh is correctly elided inside refresh_lock.
  • test_refresh_with_failed_status_clears_tokens — covers the failed-refresh path (status_code != 200) including the _initialized = False reset.
  • test_refresh_with_invalid_json_clears_tokens — covers the ValidationError branch when the refresh response body is malformed.

./scripts/test is clean: 1177 passed / 98 skipped / 1 xfailed, 100.00% coverage on src/mcp/client/auth/oauth2.py, strict-no-cover reports no leftover # pragma: no cover violations. Four previously-needed pragmas in _handle_refresh_response and the Phase 2 yield are now covered by the new tests and removed.

Breaking Changes

None. Public API is unchanged. Concurrency semantics narrow only — previously-serialized requests now proceed in parallel; token refresh remains serialized via refresh_lock.

Types of changes

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation update

Checklist

  • I have read the MCP Documentation
  • My code follows the repository's style guidelines
  • New and existing tests pass locally
  • I have added appropriate error handling
  • I have added or updated documentation as needed

Additional context

Targets main (v2 development). The same bug exists on v1.x; if a backport is wanted, that can land as a follow-up PR.

Fixes #1326