diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52b9..454fb2e498 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -598,9 +598,13 @@ 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: + if response.status_code == 401: + async with self.context.lock: # Perform full OAuth flow try: # OAuth flow must be inline due to generator constraints @@ -751,8 +755,10 @@ 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: # Step 1: Extract error field from WWW-Authenticate header error = extract_field_from_www_auth(response, "error") @@ -782,4 +788,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 diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8eec..cdb1f4d39d 100644 --- a/tests/client/test_auth.py +++ b/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,44 @@ 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. + 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)