Release context.lock before the protected request is sent

src/mcp/client/auth/oauth2.py

@@ -579,12 +579,14 @@ async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None
 
     async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
         """httpx2 auth flow integration."""
+        # Capture protocol version from request headers
+        protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER)
+
         async with self.context.lock:
             if not self._initialized:
                 await self._initialize()
 
-            # Capture protocol version from request headers
-            self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER)
+            self.context.protocol_version = protocol_version
 
             if not self.context.is_token_valid() and self.context.can_refresh_token():
                 # Try to refresh token
@@ -598,9 +600,17 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
             if self.context.is_token_valid():
                 self._add_auth_header(request)
 
-            response = yield request
+        # Released before the request goes out: the lock serialises token acquisition, and
+        # holding it for the lifetime of the response would stall every other request on
+        # this provider until the response ends - unbounded for the standalone GET SSE stream.
+        response = yield request
+
+        if response.status_code == 401:
+            async with self.context.lock:
+                # Another request may have stamped its own version while this one was in
+                # flight; re-authorization has to use the version this request carried.
+                self.context.protocol_version = protocol_version
 
-            if response.status_code == 401:
                 # Perform full OAuth flow
                 try:
                     # OAuth flow must be inline due to generator constraints
@@ -751,8 +761,12 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
 
                 # Retry with new tokens
                 self._add_auth_header(request)
-                yield request
-            elif response.status_code == 403:
+
+            yield request
+        elif response.status_code == 403:
+            async with self.context.lock:
+                self.context.protocol_version = protocol_version
+
                 # Step 1: Extract error field from WWW-Authenticate header
                 error = extract_field_from_www_auth(response, "error")
 
@@ -782,4 +796,5 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
 
                 # Retry with new tokens
                 self._add_auth_header(request)
-                yield request
+
+            yield request

tests/client/test_auth.py

@@ -6,6 +6,7 @@
 from unittest import mock
 from urllib.parse import parse_qs, quote, unquote, urlparse
 
+import anyio
 import httpx2
 import pytest
 from inline_snapshot import Is, snapshot
@@ -3253,3 +3254,108 @@ async def echo_callback() -> AuthorizationCodeResult:
         await auth_flow.asend(httpx2.Response(200, request=final_req))
     except StopAsyncIteration:
         pass
+
+
+@pytest.mark.anyio
+async def test_in_flight_request_does_not_block_a_concurrent_request(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """A request still in flight must not hold up the next one on the same provider.
+
+    The standalone GET SSE stream lives as long as the server keeps it open, so holding
+    ``context.lock`` until its response arrived stalled the first ``tools/call`` for that
+    whole time (#3209).
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() + 1800
+    oauth_provider._initialized = True
+
+    sse_sent = anyio.Event()
+    call_done = anyio.Event()
+
+    async def get_sse_stream() -> None:
+        flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+        request = await flow.__anext__()
+        sse_sent.set()
+        # The server holds the stream open, so the response lands after the call is answered.
+        with anyio.fail_after(5):
+            await call_done.wait()
+        with pytest.raises(StopAsyncIteration):
+            await flow.asend(httpx2.Response(200, request=request))
+
+    async def call_tool() -> None:
+        await sse_sent.wait()
+        flow = oauth_provider.async_auth_flow(httpx2.Request("POST", "https://api.example.com/v1/mcp"))
+        with anyio.fail_after(5):
+            request = await flow.__anext__()
+        assert request.headers["Authorization"] == "Bearer test_access_token"
+        with pytest.raises(StopAsyncIteration):
+            await flow.asend(httpx2.Response(200, request=request))
+        call_done.set()
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(get_sse_stream)
+        tg.start_soon(call_tool)
+
+
+@pytest.mark.anyio
+async def test_step_up_uses_the_protocol_version_of_its_own_request(
+    oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken
+):
+    """Re-authorization must use the protocol version of the request that was challenged.
+
+    ``context.protocol_version`` is shared and the lock is no longer held across the
+    protected request, so a second request can stamp its own version in between and
+    otherwise flip the ``resource`` parameter for the first one.
+    """
+    oauth_provider.context.current_tokens = valid_tokens
+    oauth_provider.context.token_expiry_time = time.time() + 1800
+    oauth_provider.context.client_info = OAuthClientInformationFull(
+        client_id="test_client_id",
+        client_secret="test_client_secret",
+        redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+    )
+    oauth_provider._initialized = True
+
+    captured_state: str | None = None
+
+    async def capture_redirect(url: str) -> None:
+        nonlocal captured_state
+        captured_state = parse_qs(urlparse(url).query).get("state", [None])[0]
+
+    async def mock_callback() -> AuthorizationCodeResult:
+        return AuthorizationCodeResult(code="auth_code", state=captured_state)
+
+    oauth_provider.context.redirect_handler = capture_redirect
+    oauth_provider.context.callback_handler = mock_callback
+
+    flow = oauth_provider.async_auth_flow(
+        httpx2.Request("GET", "https://api.example.com/v1/mcp", headers={"mcp-protocol-version": "2025-06-18"})
+    )
+    request = await flow.__anext__()
+
+    # A request on an older protocol version goes out while the first one is in flight.
+    other = oauth_provider.async_auth_flow(
+        httpx2.Request("POST", "https://api.example.com/v1/mcp", headers={"mcp-protocol-version": "2025-03-26"})
+    )
+    await other.__anext__()
+    await other.aclose()
+
+    response_403 = httpx2.Response(
+        403,
+        headers={"WWW-Authenticate": 'Bearer error="insufficient_scope", scope="admin:write"'},
+        request=request,
+    )
+    token_exchange_request = await flow.asend(response_403)
+
+    assert "resource=" in token_exchange_request.content.decode()
+
+    # Drive the flow to completion so the context lock is released cleanly
+    token_response = httpx2.Response(
+        200,
+        json={"access_token": "new", "token_type": "Bearer", "expires_in": 3600, "scope": "admin:write"},
+        request=token_exchange_request,
+    )
+    final_request = await flow.asend(token_response)
+    with pytest.raises(StopAsyncIteration):
+        await flow.asend(httpx2.Response(200, request=final_request))