Keep re-authorization on its own request's protocol version

Releasing context.lock before the protected request is yielded let a second
request restamp the shared context.protocol_version while the first was in
flight. The first request's 401 or 403 re-authorization then ran on the other
request's MCP-Protocol-Version, which decides whether the resource parameter is
sent. Holding the lock across the yield used to hide this.

Read the header into a local once, and restamp context.protocol_version from it
inside the 401 and 403 blocks after the lock is re-acquired. Every read of
context.protocol_version happens under the lock in one of those regions, so this
needs no new parameter on the token-request helpers.

Also wrap the concurrency test's call_done.wait() in anyio.fail_after(5) so a
failure in the sibling task fails fast instead of hanging.

Author: guptaishaan

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
@@ -605,6 +607,10 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
 
         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
+
                 # Perform full OAuth flow
                 try:
                     # OAuth flow must be inline due to generator constraints
@@ -759,6 +765,8 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
             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")
 

tests/client/test_auth.py

@@ -3278,7 +3278,8 @@ async def get_sse_stream() -> None:
         request = await flow.__anext__()
         sse_sent.set()
         # The server holds the stream open, so the response lands after the call is answered.
-        await call_done.wait()
+        with anyio.fail_after(5):
+            await call_done.wait()
         with pytest.raises(StopAsyncIteration):
             await flow.asend(httpx2.Response(200, request=request))
 
@@ -3295,3 +3296,66 @@ async def call_tool() -> None:
     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))