diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52b9..cb165e23e8 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -1,4 +1,4 @@ -"""OAuth2 Authentication implementation for httpx2. +"""OAuth2 Authentication implementation for HTTPX. Implements authorization code flow with PKCE and automatic token refresh. """ @@ -11,15 +11,14 @@ import time from collections.abc import AsyncGenerator, Awaitable, Callable from dataclasses import dataclass, field -from typing import Any, Protocol, get_args +from typing import Any, Protocol from urllib.parse import quote, urlencode, urljoin, urlparse import anyio -import httpx2 -from mcp_types.version import is_version_at_least +import httpx from pydantic import BaseModel, Field, ValidationError -from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError +from mcp.client.auth.exceptions import OAuthFlowError, OAuthTokenError from mcp.client.auth.utils import ( build_oauth_authorization_server_metadata_discovery_urls, build_protected_resource_metadata_discovery_urls, @@ -41,6 +40,7 @@ validate_authorization_response_iss, validate_metadata_issuer, ) +from mcp.client.streamable_http import MCP_PROTOCOL_VERSION from mcp.shared.auth import ( AuthorizationCodeResult, OAuthClientInformationFull, @@ -48,66 +48,16 @@ OAuthMetadata, OAuthToken, ProtectedResourceMetadata, - TokenEndpointAuthMethod, ) from mcp.shared.auth_utils import ( calculate_token_expiry, check_resource_allowed, resource_url_from_server_url, ) -from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER +from mcp.shared.version import is_version_at_least logger = logging.getLogger(__name__) -# Methods a registered client's record may carry without a token request being an error, -# derived from the set the SDK is willing to request so the two cannot drift. `None`/"none" -# send no client secret. `private_key_jwt` sends none from here either: only -# `PrivateKeyJWTOAuthProvider` signs the assertion, and only in its client-credentials -# exchange, so its inherited refresh path must pass through here without raising - a refresh -# the server then rejects falls back to a fresh client-credentials exchange, which signs. -# Anything else is a method no client here can apply. -_KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = (None, *get_args(TokenEndpointAuthMethod)) - -# Methods that authenticate the token request with the minted `client_secret`; a -# registration assigning one is only usable if the server issued that secret. -_SECRET_TOKEN_ENDPOINT_AUTH_METHODS = ("client_secret_post", "client_secret_basic") - -# Methods a registration completed by the authorization-code flow can act on. That flow -# authenticates the token request with the minted client secret (or nothing); it holds no key -# to sign a `private_key_jwt` assertion, so a server assigning that method has registered a -# client this flow cannot use. `PrivateKeyJWTOAuthProvider` never registers dynamically. -_REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = tuple( - method for method in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS if method != "private_key_jwt" -) - - -def check_registration_usable(client_info: OAuthClientInformationFull) -> None: - """Confirm a registration this flow completed is one it can act on. - - RFC 7591 §3.2.1 lets the authorization server replace requested metadata and leaves it to - the client to "check the values in the response to determine if the registration is - sufficient for use". Two substitutions make the minted credentials unusable, and both are - judged here - before the record is persisted or any interactive authorization begins - - rather than surfacing later as an opaque failure at the token endpoint: a token-endpoint - auth method the authorization-code flow cannot apply (one it does not implement, or - `private_key_jwt`, whose assertion this flow has no key to sign), and a secret-based - method the flow could apply but for which the server issued no `client_secret`. - - Raises: - OAuthRegistrationError: The server registered the client with a - `token_endpoint_auth_method` this flow cannot apply, or with a secret-based - method but no `client_secret`. - """ - method = client_info.token_endpoint_auth_method - if method not in _REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: - raise OAuthRegistrationError( - f"Authorization server registered the client with unsupported token_endpoint_auth_method {method!r}" - ) - if method in _SECRET_TOKEN_ENDPOINT_AUTH_METHODS and client_info.client_secret is None: - raise OAuthRegistrationError( - f"Authorization server registered the client for {method!r} but issued no client_secret" - ) - class PKCEParameters(BaseModel): """PKCE (Proof Key for Code Exchange) parameters.""" @@ -153,6 +103,7 @@ class OAuthContext: storage: TokenStorage redirect_handler: Callable[[str], Awaitable[None]] | None callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None + timeout: float = 300.0 client_metadata_url: str | None = None # Discovered metadata @@ -169,7 +120,18 @@ class OAuthContext: token_expiry_time: float | None = None # State + # + # `lock` guards short-lived reads/writes of provider state (initialization + # flag, token cache mutation, protocol_version assignment). It is held only + # while mutating state and is released before any HTTP request is yielded + # so a long-running request (e.g. GET SSE long-poll) does not block + # unrelated concurrent requests. + # + # `refresh_lock` provides single-flight semantics for token refresh: only + # one concurrent refresh fires; other waiters block on this lock, then + # re-check the token cache and proceed without re-refreshing. lock: anyio.Lock = field(default_factory=anyio.Lock) + refresh_lock: anyio.Lock = field(default_factory=anyio.Lock) def get_authorization_base_url(self, server_url: str) -> str: """Extract base URL by removing path component.""" @@ -240,12 +202,6 @@ def prepare_token_auth( Returns: Tuple of (updated_data, updated_headers) - - Raises: - OAuthTokenError: The client record carries a `token_endpoint_auth_method` this - client does not know. A dynamic registration assigning an unusable method is - rejected earlier, by `check_registration_usable`; this fires for a stored or - pre-registered record that reaches a token request with such a method. """ if headers is None: headers = {} # pragma: no cover @@ -255,7 +211,7 @@ def prepare_token_auth( auth_method = self.client_info.token_endpoint_auth_method - if auth_method == "client_secret_basic" and self.client_info.client_secret: + if auth_method == "client_secret_basic" and self.client_info.client_id and self.client_info.client_secret: # URL-encode client ID and secret per RFC 6749 Section 2.3.1 encoded_id = quote(self.client_info.client_id, safe="") encoded_secret = quote(self.client_info.client_secret, safe="") @@ -264,20 +220,17 @@ def prepare_token_auth( headers["Authorization"] = f"Basic {encoded_credentials}" # Don't include client_secret in body for basic auth data = {k: v for k, v in data.items() if k != "client_secret"} - elif auth_method == "client_secret_post" and self.client_info.client_secret: + elif auth_method == "client_secret_post" and self.client_info.client_id and self.client_info.client_secret: # Include client_id and client_secret in request body (RFC 6749 §2.3.1) data["client_id"] = self.client_info.client_id data["client_secret"] = self.client_info.client_secret - elif auth_method not in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: - raise OAuthTokenError(f"Registered client uses unsupported token_endpoint_auth_method {auth_method!r}") - # For "none" (or absent), don't add any client_secret; "private_key_jwt" adds its - # assertion in the provider that implements it, not here. + # For auth_method == "none", don't add any client_secret return data, headers -class OAuthClientProvider(httpx2.Auth): - """OAuth2 authentication for httpx2. +class OAuthClientProvider(httpx.Auth): + """OAuth2 authentication for httpx. Handles OAuth flow with automatic client registration and token storage. """ @@ -291,6 +244,7 @@ def __init__( storage: TokenStorage, redirect_handler: Callable[[str], Awaitable[None]] | None = None, callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None = None, + timeout: float = 300.0, client_metadata_url: str | None = None, validate_resource_url: Callable[[str, str | None], Awaitable[None]] | None = None, ): @@ -302,6 +256,7 @@ def __init__( storage: Token storage implementation. redirect_handler: Handler for authorization redirects. callback_handler: Handler for authorization callbacks. + timeout: Timeout for the OAuth flow. client_metadata_url: URL-based client ID. When provided and the server advertises client_id_metadata_document_supported=True, this URL will be used as the client_id instead of performing dynamic client registration. @@ -327,12 +282,13 @@ def __init__( storage=storage, redirect_handler=redirect_handler, callback_handler=callback_handler, + timeout=timeout, client_metadata_url=client_metadata_url, ) self._validate_resource_url_callback = validate_resource_url self._initialized = False - async def _handle_protected_resource_response(self, response: httpx2.Response) -> bool: + async def _handle_protected_resource_response(self, response: httpx.Response) -> bool: """Handle protected resource metadata discovery response. Per SEP-985, supports fallback when discovery fails at one URL. @@ -363,7 +319,7 @@ async def _handle_protected_resource_response(self, response: httpx2.Response) - f"Protected Resource Metadata request failed: {response.status_code}" ) # pragma: no cover - async def _perform_authorization(self) -> httpx2.Request: + async def _perform_authorization(self) -> httpx.Request: """Perform the authorization flow.""" auth_code, code_verifier = await self._perform_authorization_code_grant() token_request = await self._exchange_token_authorization_code(auth_code, code_verifier) @@ -438,7 +394,9 @@ def _get_token_endpoint(self) -> str: token_url = urljoin(auth_base_url, "/token") return token_url - async def _exchange_token_authorization_code(self, auth_code: str, code_verifier: str) -> httpx2.Request: + async def _exchange_token_authorization_code( + self, auth_code: str, code_verifier: str, *, token_data: dict[str, Any] | None = {} + ) -> httpx.Request: """Build token exchange request for authorization_code flow.""" if self.context.client_metadata.redirect_uris is None: raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover @@ -446,13 +404,16 @@ async def _exchange_token_authorization_code(self, auth_code: str, code_verifier raise OAuthFlowError("Missing client info") # pragma: no cover token_url = self._get_token_endpoint() - token_data: dict[str, Any] = { - "grant_type": "authorization_code", - "code": auth_code, - "redirect_uri": str(self.context.client_metadata.redirect_uris[0]), - "client_id": self.context.client_info.client_id, - "code_verifier": code_verifier, - } + token_data = token_data or {} + token_data.update( + { + "grant_type": "authorization_code", + "code": auth_code, + "redirect_uri": str(self.context.client_metadata.redirect_uris[0]), + "client_id": self.context.client_info.client_id, + "code_verifier": code_verifier, + } + ) # Only include resource param if conditions are met if self.context.should_include_resource_param(self.context.protocol_version): @@ -462,14 +423,14 @@ async def _exchange_token_authorization_code(self, auth_code: str, code_verifier headers = {"Content-Type": "application/x-www-form-urlencoded"} token_data, headers = self.context.prepare_token_auth(token_data, headers) - return httpx2.Request("POST", token_url, data=token_data, headers=headers) + return httpx.Request("POST", token_url, data=token_data, headers=headers) - async def _handle_token_response(self, response: httpx2.Response) -> None: + async def _handle_token_response(self, response: httpx.Response) -> None: """Handle token exchange response.""" if response.status_code not in {200, 201}: - body = await response.aread() - body_text = body.decode("utf-8") - raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body_text}") + body = await response.aread() # pragma: no cover + body_text = body.decode("utf-8") # pragma: no cover + raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body_text}") # pragma: no cover # Parse and validate response with scope validation token_response = await handle_token_response_scopes(response) @@ -486,7 +447,7 @@ async def _handle_token_response(self, response: httpx2.Response) -> None: self.context.update_token_expiry(token_response) await self.context.storage.set_tokens(token_response) - async def _refresh_token(self) -> httpx2.Request: + async def _refresh_token(self) -> httpx.Request: """Build token refresh request.""" if not self.context.current_tokens or not self.context.current_tokens.refresh_token: raise OAuthTokenError("No refresh token available") # pragma: no cover @@ -514,9 +475,9 @@ async def _refresh_token(self) -> httpx2.Request: headers = {"Content-Type": "application/x-www-form-urlencoded"} refresh_data, headers = self.context.prepare_token_auth(refresh_data, headers) - return httpx2.Request("POST", token_url, data=refresh_data, headers=headers) + return httpx.Request("POST", token_url, data=refresh_data, headers=headers) - async def _handle_refresh_response(self, response: httpx2.Response) -> bool: + async def _handle_refresh_response(self, response: httpx.Response) -> bool: """Handle token refresh response. Returns True if successful.""" if response.status_code != 200: logger.warning(f"Token refresh failed: {response.status_code}") @@ -542,7 +503,7 @@ async def _handle_refresh_response(self, response: httpx2.Response) -> bool: await self.context.storage.set_tokens(token_response) return True - except ValidationError: # pragma: no cover + except ValidationError: logger.exception("Invalid refresh response") self.context.clear_tokens() return False @@ -553,12 +514,12 @@ async def _initialize(self) -> None: self.context.client_info = await self.context.storage.get_client_info() self._initialized = True - def _add_auth_header(self, request: httpx2.Request) -> None: + def _add_auth_header(self, request: httpx.Request) -> None: """Add authorization header to request if we have valid tokens.""" if self.context.current_tokens and self.context.current_tokens.access_token: # pragma: no branch request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}" - async def _handle_oauth_metadata_response(self, response: httpx2.Response) -> None: + async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None: content = await response.aread() metadata = OAuthMetadata.model_validate_json(content) self.context.oauth_metadata = metadata @@ -577,30 +538,88 @@ async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None if not check_resource_allowed(requested_resource=default_resource, configured_resource=prm_resource): raise OAuthFlowError(f"Protected resource {prm_resource} does not match expected {default_resource}") - async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - """httpx2 auth flow integration.""" + async def _prepare_and_decide_refresh(self, request: httpx.Request) -> bool: + """Phase 1: initialize + capture protocol version, then decide whether a + proactive token refresh is needed. Holds ``self.context.lock`` only + briefly. Returns ``True`` when the token is invalid but refreshable. + """ 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) - - if not self.context.is_token_valid() and self.context.can_refresh_token(): - # Try to refresh token - refresh_request = await self._refresh_token() - refresh_response = yield refresh_request - - if not await self._handle_refresh_response(refresh_response): - # Refresh failed, need full re-authentication - self._initialized = False - - if self.context.is_token_valid(): - self._add_auth_header(request) + self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION) + + # pragma: no branch — coverage.py on Python 3.10/3.11 (sys.settrace + # backend) cannot reliably track both arms of compound boolean + # predicates inside an ``async with`` block in an async generator. + # Python 3.12+ (sys.monitoring) handles this correctly; the pragmas + # below are workarounds for the legacy backend only. + if not self.context.is_token_valid() and self.context.can_refresh_token(): # pragma: no branch + return True + return False - response = yield request + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + """HTTPX auth flow integration. - if response.status_code == 401: + Lock scope: + ``self.context.lock`` is held only while reading/mutating provider + state. The actual HTTP request yield (which may be a long-poll GET + SSE stream) runs outside any lock so concurrent unrelated requests + are not blocked. ``self.context.refresh_lock`` provides + single-flight semantics for token refresh. + """ + # === Phase 1: state read + refresh decision (brief context.lock) === + needs_refresh = await self._prepare_and_decide_refresh(request) + + # === Phase 2: single-flight token refresh (yield outside context.lock) === + if needs_refresh: + async with self.context.refresh_lock: + # Re-check under context.lock: another coroutine may already have + # refreshed while we were waiting on refresh_lock. + refresh_request: httpx.Request | None = None + async with self.context.lock: + if not self.context.is_token_valid() and self.context.can_refresh_token(): # pragma: no branch + refresh_request = await self._refresh_token() + if refresh_request is not None: # pragma: no branch + # yield runs outside any lock so a long network round trip + # does not block unrelated concurrent requests. + refresh_response = yield refresh_request + async with self.context.lock: + if not await self._handle_refresh_response(refresh_response): # pragma: no branch + # Refresh failed; fall through to 401 handling below. + self._initialized = False + + # === Phase 3: send request (no lock; safe for long-poll GET SSE) === + if self.context.is_token_valid(): + self._add_auth_header(request) + + # Capture the access token actually used to send this request so the + # 401 handler below can detect a token change made by a concurrent + # request while this one was in flight. + sent_access_token = self.context.current_tokens.access_token if self.context.current_tokens else None + + response = yield request + + # === Phase 4: 401 / 403 full OAuth flow === + # NOTE: Phase 4 yields multiple sub-requests (discovery, registration, + # token exchange) under context.lock. This is the existing behavior and + # is acceptable because the 401 path is exceptional and not concurrent + # with steady-state traffic. A future refactor could narrow the lock + # here in the same pattern as Phase 1-2. + if response.status_code == 401: + async with self.context.lock: + # Concurrency guard: while this request was in flight, another + # request holding ``context.lock`` may have already completed a + # token refresh or a full re-authorization. If the stored access + # token changed since we sent this request, the 401 is stale - + # retry once with the new token instead of running a second, + # duplicate ``authorization_code`` exchange. + current_access_token = self.context.current_tokens.access_token if self.context.current_tokens else None + if current_access_token is not None and current_access_token != sent_access_token: + self._add_auth_header(request) + yield request + return # Perform full OAuth flow try: # OAuth flow must be inline due to generator constraints @@ -723,7 +742,6 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx ) registration_response = yield registration_request client_information = await handle_registration_response(registration_response) - check_registration_usable(client_information) # Only record the issuer when the registration above actually targeted # the discovered AS — either via its published registration_endpoint, # or because the resource-origin /register fallback is on the issuer's @@ -752,7 +770,8 @@ 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: + 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") diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8eec..8cd5796b09 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -1,18 +1,20 @@ """Tests for refactored OAuth client authentication implementation.""" import base64 +import contextlib import json import time from unittest import mock from urllib.parse import parse_qs, quote, unquote, urlparse -import httpx2 +import anyio +import httpx import pytest from inline_snapshot import Is, snapshot from pydantic import AnyHttpUrl, AnyUrl from mcp.client.auth import OAuthClientProvider, PKCEParameters -from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError +from mcp.client.auth.exceptions import OAuthFlowError from mcp.client.auth.utils import ( build_oauth_authorization_server_metadata_discovery_urls, build_protected_resource_metadata_discovery_urls, @@ -111,7 +113,7 @@ async def callback_handler() -> AuthorizationCodeResult: @pytest.fixture def prm_metadata_response(): """PRM metadata response with scopes.""" - return httpx2.Response( + return httpx.Response( 200, content=( b'{"resource": "https://api.example.com/v1/mcp", ' @@ -124,7 +126,7 @@ def prm_metadata_response(): @pytest.fixture def prm_metadata_without_scopes_response(): """PRM metadata response without scopes.""" - return httpx2.Response( + return httpx.Response( 200, content=( b'{"resource": "https://api.example.com/v1/mcp", ' @@ -137,20 +139,20 @@ def prm_metadata_without_scopes_response(): @pytest.fixture def init_response_with_www_auth_scope(): """Initial 401 response with WWW-Authenticate header containing scope.""" - return httpx2.Response( + return httpx.Response( 401, headers={"WWW-Authenticate": 'Bearer scope="special:scope from:www-authenticate"'}, - request=httpx2.Request("GET", "https://api.example.com/test"), + request=httpx.Request("GET", "https://api.example.com/test"), ) @pytest.fixture def init_response_without_www_auth_scope(): """Initial 401 response without WWW-Authenticate scope.""" - return httpx2.Response( + return httpx.Response( 401, headers={}, - request=httpx2.Request("GET", "https://api.example.com/test"), + request=httpx.Request("GET", "https://api.example.com/test"), ) @@ -192,6 +194,7 @@ async def test_oauth_provider_initialization( assert oauth_provider.context.server_url == "https://api.example.com/v1/mcp" assert oauth_provider.context.client_metadata == client_metadata assert oauth_provider.context.storage == mock_storage + assert oauth_provider.context.timeout == 300.0 assert oauth_provider.context is not None def test_context_url_parsing(self, oauth_provider: OAuthClientProvider): @@ -289,8 +292,8 @@ async def callback_handler() -> AuthorizationCodeResult: ) # Test without WWW-Authenticate (fallback) - init_response = httpx2.Response( - status_code=401, headers={}, request=httpx2.Request("GET", "https://request-api.example.com") + init_response = httpx.Response( + status_code=401, headers={}, request=httpx.Request("GET", "https://request-api.example.com") ) urls = build_protected_resource_metadata_discovery_urls( @@ -407,7 +410,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl ) # Create a test request - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") # Mock the auth flow auth_flow = oauth_provider.async_auth_flow(test_request) @@ -417,7 +420,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl assert "Authorization" not in request.headers # Send a 401 response to trigger the OAuth flow - response = httpx2.Response( + response = httpx.Response( 401, headers={ "WWW-Authenticate": 'Bearer resource_metadata="https://api.example.com/.well-known/oauth-protected-resource"' @@ -432,7 +435,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl # Send a successful discovery response with minimal protected resource metadata # Note: auth server URL has a path (/v1/mcp), so only path-based URLs will be tried - discovery_response = httpx2.Response( + discovery_response = httpx.Response( 200, content=b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com/v1/mcp"]}', request=discovery_request, @@ -447,7 +450,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl assert oauth_metadata_request_1.method == "GET" # Send a 404 response - oauth_metadata_response_1 = httpx2.Response( + oauth_metadata_response_1 = httpx.Response( 404, content=b"Not Found", request=oauth_metadata_request_1, @@ -459,7 +462,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl assert oauth_metadata_request_2.method == "GET" # Send a 400 response - oauth_metadata_response_2 = httpx2.Response( + oauth_metadata_response_2 = httpx.Response( 400, content=b"Bad Request", request=oauth_metadata_request_2, @@ -471,7 +474,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl assert oauth_metadata_request_3.method == "GET" # Send a 500 response - oauth_metadata_response_3 = httpx2.Response( + oauth_metadata_response_3 = httpx.Response( 500, content=b"Internal Server Error", request=oauth_metadata_request_3, @@ -489,7 +492,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl assert token_request.method == "POST" # Send a successful token response - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=( b'{"access_token": "new_access_token", "token_type": "Bearer", "expires_in": 3600, ' @@ -505,7 +508,7 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl assert str(final_request.url) == "https://api.example.com/v1/mcp" # Send final success response to properly close the generator - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -520,7 +523,7 @@ async def test_handle_metadata_response_success(self, oauth_provider: OAuthClien "authorization_endpoint": "https://auth.example.com/authorize", "token_endpoint": "https://auth.example.com/token" }""" - response = httpx2.Response(200, content=content) + response = httpx.Response(200, content=content) # Should set metadata; the empty path is preserved (no trailing slash added) await oauth_provider._handle_oauth_metadata_response(response) @@ -531,8 +534,8 @@ async def test_handle_metadata_response_success(self, oauth_provider: OAuthClien async def test_prioritize_www_auth_scope_over_prm( self, oauth_provider: OAuthClientProvider, - prm_metadata_response: httpx2.Response, - init_response_with_www_auth_scope: httpx2.Response, + prm_metadata_response: httpx.Response, + init_response_with_www_auth_scope: httpx.Response, ): """Test that WWW-Authenticate scope is prioritized over PRM scopes.""" # First, process PRM metadata to set protected_resource_metadata with scopes @@ -551,8 +554,8 @@ async def test_prioritize_www_auth_scope_over_prm( async def test_prioritize_prm_scopes_when_no_www_auth_scope( self, oauth_provider: OAuthClientProvider, - prm_metadata_response: httpx2.Response, - init_response_without_www_auth_scope: httpx2.Response, + prm_metadata_response: httpx.Response, + init_response_without_www_auth_scope: httpx.Response, ): """Test that PRM scopes are prioritized when WWW-Authenticate header has no scopes.""" # Process the PRM metadata to set protected_resource_metadata with scopes @@ -571,8 +574,8 @@ async def test_prioritize_prm_scopes_when_no_www_auth_scope( async def test_omit_scope_when_no_prm_scopes_or_www_auth( self, oauth_provider: OAuthClientProvider, - prm_metadata_without_scopes_response: httpx2.Response, - init_response_without_www_auth_scope: httpx2.Response, + prm_metadata_without_scopes_response: httpx.Response, + init_response_without_www_auth_scope: httpx.Response, ): """Test that scope is omitted when PRM has no scopes and WWW-Authenticate doesn't specify scope.""" # Process the PRM metadata without scopes @@ -980,7 +983,7 @@ async def test_handle_registration_response_reads_before_accessing_text(self): """Test that response.aread() is called before accessing response.text.""" # Track if aread() was called - class MockResponse(httpx2.Response): + class MockResponse(httpx.Response): def __init__(self): self.status_code = 400 self._aread_called = False @@ -1008,91 +1011,6 @@ def text(self): assert "Registration failed: 400" in str(exc_info.value) -@pytest.mark.anyio -async def test_registration_response_with_substituted_metadata_yields_the_credentials(): - """A 201 whose echoed metadata differs from the request still registers the client. - - The authorization server returned an application_type outside OIDC Registration's set, - a null redirect_uris, and an auth method the SDK does not implement. RFC 7591 §3.2.1 - permits the server to substitute values; the client keeps the credentials it minted. - """ - body = ( - b'{"client_id": "issued-id", "client_secret": "issued-secret", ' - b'"application_type": "confidential", "redirect_uris": null, ' - b'"token_endpoint_auth_method": "client_secret_jwt"}' - ) - response = httpx2.Response(201, content=body) - - client_info = await handle_registration_response(response) - - assert client_info.client_id == "issued-id" - assert client_info.client_secret == "issued-secret" - assert client_info.application_type == "confidential" - - -@pytest.mark.anyio -@pytest.mark.parametrize("echoed_issuer", ["https://not-the-flow.example", 12345], ids=["string", "not-a-string"]) -async def test_registration_response_does_not_seed_the_issuer_binding_from_the_body(echoed_issuer: object): - """The issuer binding (SEP-2352) is the SDK's record of which server it registered with, - stamped by the auth flow; an "issuer" member in the untrusted response body is dropped - before parsing - never populating the binding, and never failing the parse either, so a - mismatched or malformed value cannot discard the credentials on every 401.""" - body = json.dumps({"client_id": "issued-id", "issuer": echoed_issuer}).encode() - - client_info = await handle_registration_response(httpx2.Response(201, content=body)) - - assert client_info.client_id == "issued-id" - assert client_info.issuer is None - - -@pytest.mark.anyio -@pytest.mark.parametrize( - "content", - [b"not json", b'["json", "but", "not", "an", "object"]', '{"client_id": "café"}'.encode("latin-1")], - ids=["not-json", "not-an-object", "not-utf8"], -) -async def test_a_2xx_body_that_is_not_client_information_is_an_oauth_registration_error(content: bytes): - """A success status whose body is not client information - unparseable, not an object, or - not valid UTF-8 - surfaces as OAuthRegistrationError rather than a raw parse failure, so a - single OAuthFlowError handler still covers registration.""" - response = httpx2.Response(201, content=content) - - with pytest.raises(OAuthRegistrationError): - await handle_registration_response(response) - - -@pytest.mark.anyio -async def test_token_exchange_reports_an_unimplemented_registered_auth_method(oauth_provider: OAuthClientProvider): - """A server-assigned auth method the SDK cannot apply (RFC 7591 §3.2.1 lets the server - substitute one) is reported at the token exchange rather than sending the request - unauthenticated for the server to reject as invalid_client.""" - oauth_provider.context.client_info = OAuthClientInformationFull( - client_id="registered-id", - client_secret="registered-secret", - token_endpoint_auth_method="client_secret_jwt", - ) - - with pytest.raises(OAuthTokenError): - await oauth_provider._exchange_token_authorization_code("test_auth_code", "test_verifier") - - -def test_prepare_token_auth_leaves_a_private_key_jwt_client_to_its_provider(oauth_provider: OAuthClientProvider): - """private_key_jwt is recognized, so the base leaves the request untouched rather than - raising - PrivateKeyJWTOAuthProvider's inherited refresh path passes through here, and a - refresh the server then rejects (no assertion is signed on it) falls back to a fresh, - signed client-credentials exchange instead of aborting the flow.""" - oauth_provider.context.client_info = OAuthClientInformationFull( - client_id="registered-id", - client_secret="registered-secret", - token_endpoint_auth_method="private_key_jwt", - ) - - data, headers = oauth_provider.context.prepare_token_auth({"grant_type": "refresh_token"}, {}) - - assert data == {"grant_type": "refresh_token"} - assert headers == {} - - class TestCreateClientRegistrationRequest: """Test client registration request creation.""" @@ -1152,7 +1070,7 @@ def test_registration_request_sends_application_type(): class TestAuthFlow: - """Test the auth flow in httpx2.""" + """Test the auth flow in httpx.""" @pytest.mark.anyio async def test_auth_flow_with_valid_tokens( @@ -1166,7 +1084,7 @@ async def test_auth_flow_with_valid_tokens( oauth_provider._initialized = True # Create a test request - test_request = httpx2.Request("GET", "https://api.example.com/test") + test_request = httpx.Request("GET", "https://api.example.com/test") # Mock the auth flow auth_flow = oauth_provider.async_auth_flow(test_request) @@ -1176,7 +1094,7 @@ async def test_auth_flow_with_valid_tokens( assert request.headers["Authorization"] == "Bearer test_access_token" # Send a successful response - response = httpx2.Response(200) + response = httpx.Response(200) try: await auth_flow.asend(response) except StopAsyncIteration: @@ -1191,7 +1109,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide oauth_provider._initialized = True # Create a test request - test_request = httpx2.Request("GET", "https://api.example.com/mcp") + test_request = httpx.Request("GET", "https://api.example.com/mcp") # Mock the auth flow auth_flow = oauth_provider.async_auth_flow(test_request) @@ -1201,7 +1119,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide assert "Authorization" not in request.headers # Send a 401 response to trigger the OAuth flow - response = httpx2.Response( + response = httpx.Response( 401, headers={ "WWW-Authenticate": 'Bearer resource_metadata="https://api.example.com/.well-known/oauth-protected-resource"' @@ -1215,7 +1133,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide assert str(discovery_request.url) == "https://api.example.com/.well-known/oauth-protected-resource" # Send a successful discovery response with minimal protected resource metadata - discovery_response = httpx2.Response( + discovery_response = httpx.Response( 200, content=b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}', request=discovery_request, @@ -1228,7 +1146,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide assert "mcp-protocol-version" in oauth_metadata_request.headers # Send a successful OAuth metadata response - oauth_metadata_response = httpx2.Response( + oauth_metadata_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com", ' @@ -1245,7 +1163,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide assert str(registration_request.url) == "https://auth.example.com/register" # Send a successful registration response - registration_response = httpx2.Response( + registration_response = httpx.Response( 201, content=b'{"client_id": "test_client_id", "client_secret": "test_client_secret", "redirect_uris": ["http://localhost:3030/callback"]}', request=registration_request, @@ -1263,7 +1181,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide assert "code=test_auth_code" in token_request.content.decode() # Send a successful token response - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=( b'{"access_token": "new_access_token", "token_type": "Bearer", "expires_in": 3600, ' @@ -1279,7 +1197,7 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide assert str(final_request.url) == "https://api.example.com/mcp" # Send final success response to properly close the generator - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -1301,7 +1219,7 @@ async def test_auth_flow_no_unnecessary_retry_after_oauth( oauth_provider.context.token_expiry_time = time.time() + 1800 oauth_provider._initialized = True - test_request = httpx2.Request("GET", "https://api.example.com/mcp") + test_request = httpx.Request("GET", "https://api.example.com/mcp") auth_flow = oauth_provider.async_auth_flow(test_request) # Count how many times the request is yielded @@ -1313,7 +1231,7 @@ async def test_auth_flow_no_unnecessary_retry_after_oauth( assert request.headers["Authorization"] == "Bearer test_access_token" # Send a successful 200 response - response = httpx2.Response(200, request=request) + response = httpx.Response(200, request=request) # In the buggy version, this would yield the request AGAIN unconditionally # In the fixed version, this should end the generator @@ -1344,7 +1262,7 @@ async def test_token_exchange_accepts_201_status( oauth_provider._initialized = True # Create a test request - test_request = httpx2.Request("GET", "https://api.example.com/mcp") + test_request = httpx.Request("GET", "https://api.example.com/mcp") # Mock the auth flow auth_flow = oauth_provider.async_auth_flow(test_request) @@ -1354,7 +1272,7 @@ async def test_token_exchange_accepts_201_status( assert "Authorization" not in request.headers # Send a 401 response to trigger the OAuth flow - response = httpx2.Response( + response = httpx.Response( 401, headers={ "WWW-Authenticate": 'Bearer resource_metadata="https://api.example.com/.well-known/oauth-protected-resource"' @@ -1368,7 +1286,7 @@ async def test_token_exchange_accepts_201_status( assert str(discovery_request.url) == "https://api.example.com/.well-known/oauth-protected-resource" # Send a successful discovery response with minimal protected resource metadata - discovery_response = httpx2.Response( + discovery_response = httpx.Response( 200, content=b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}', request=discovery_request, @@ -1381,7 +1299,7 @@ async def test_token_exchange_accepts_201_status( assert "mcp-protocol-version" in oauth_metadata_request.headers # Send a successful OAuth metadata response - oauth_metadata_response = httpx2.Response( + oauth_metadata_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com", ' @@ -1398,7 +1316,7 @@ async def test_token_exchange_accepts_201_status( assert str(registration_request.url) == "https://auth.example.com/register" # Send a successful registration response with 201 status - registration_response = httpx2.Response( + registration_response = httpx.Response( 201, content=b'{"client_id": "test_client_id", "client_secret": "test_client_secret", "redirect_uris": ["http://localhost:3030/callback"]}', request=registration_request, @@ -1416,7 +1334,7 @@ async def test_token_exchange_accepts_201_status( assert "code=test_auth_code" in token_request.content.decode() # Send a successful token response with 201 status code (test both 200 and 201 are accepted) - token_response = httpx2.Response( + token_response = httpx.Response( 201, content=( b'{"access_token": "new_access_token", "token_type": "Bearer", "expires_in": 3600, ' @@ -1432,7 +1350,7 @@ async def test_token_exchange_accepts_201_status( assert str(final_request.url) == "https://api.example.com/mcp" # Send final success response to properly close the generator - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -1489,14 +1407,14 @@ async def mock_callback() -> AuthorizationCodeResult: oauth_provider.context.callback_handler = mock_callback - test_request = httpx2.Request("GET", "https://api.example.com/mcp") + test_request = httpx.Request("GET", "https://api.example.com/mcp") auth_flow = oauth_provider.async_auth_flow(test_request) # First request request = await auth_flow.__anext__() # Send 403 with new scope requirement - response_403 = httpx2.Response( + response_403 = httpx.Response( 403, headers={"WWW-Authenticate": 'Bearer error="insufficient_scope", scope="admin:write admin:delete"'}, request=request, @@ -1510,7 +1428,7 @@ async def mock_callback() -> AuthorizationCodeResult: assert redirect_captured # Complete the flow with successful token response - token_response = httpx2.Response( + token_response = httpx.Response( 200, json={ "access_token": "new_token_with_new_scope", @@ -1525,7 +1443,7 @@ async def mock_callback() -> AuthorizationCodeResult: final_request = await auth_flow.asend(token_response) # Send success response - flow should complete - success_response = httpx2.Response(200, request=final_request) + success_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(success_response) pytest.fail("Should have stopped after successful response") # pragma: no cover @@ -1569,9 +1487,9 @@ async def mock_callback() -> AuthorizationCodeResult: oauth_provider.context.redirect_handler = capture_redirect oauth_provider.context.callback_handler = mock_callback - auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/mcp")) + auth_flow = oauth_provider.async_auth_flow(httpx.Request("GET", "https://api.example.com/mcp")) request = await auth_flow.__anext__() - response_403 = httpx2.Response( + response_403 = httpx.Response( 403, headers={"WWW-Authenticate": 'Bearer error="insufficient_scope", scope="write"'}, request=request, @@ -1581,14 +1499,14 @@ async def mock_callback() -> AuthorizationCodeResult: assert reauthorize_scope == "read write" # Drive the flow to completion so the context lock is released cleanly - token_response = httpx2.Response( + token_response = httpx.Response( 200, json={"access_token": "new", "token_type": "Bearer", "expires_in": 3600, "scope": "read write"}, request=token_exchange_request, ) final_request = await auth_flow.asend(token_response) try: - await auth_flow.asend(httpx2.Response(200, request=final_request)) + await auth_flow.asend(httpx.Response(200, request=final_request)) except StopAsyncIteration: pass @@ -1703,7 +1621,7 @@ async def callback_handler() -> AuthorizationCodeResult: redirect_uris=[AnyUrl("http://localhost:3030/callback")], ) - test_request = httpx2.Request("GET", "https://mcp.linear.app/sse") + test_request = httpx.Request("GET", "https://mcp.linear.app/sse") auth_flow = provider.async_auth_flow(test_request) # First request @@ -1711,21 +1629,21 @@ async def callback_handler() -> AuthorizationCodeResult: assert "Authorization" not in request.headers # Send 401 without WWW-Authenticate header (typical legacy server) - response = httpx2.Response(401, headers={}, request=test_request) + response = httpx.Response(401, headers={}, request=test_request) # Should try path-based PRM first prm_request_1 = await auth_flow.asend(response) assert str(prm_request_1.url) == "https://mcp.linear.app/.well-known/oauth-protected-resource/sse" # PRM returns 404 - prm_response_1 = httpx2.Response(404, request=prm_request_1) + prm_response_1 = httpx.Response(404, request=prm_request_1) # Should try root-based PRM prm_request_2 = await auth_flow.asend(prm_response_1) assert str(prm_request_2.url) == "https://mcp.linear.app/.well-known/oauth-protected-resource" # PRM returns 404 again - all PRM URLs failed - prm_response_2 = httpx2.Response(404, request=prm_request_2) + prm_response_2 = httpx.Response(404, request=prm_request_2) # Should fall back to root OAuth discovery (March 2025 spec behavior) oauth_metadata_request = await auth_flow.asend(prm_response_2) @@ -1733,7 +1651,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert oauth_metadata_request.method == "GET" # Send successful OAuth metadata response - oauth_metadata_response = httpx2.Response( + oauth_metadata_response = httpx.Response( 200, content=( b'{"issuer": "https://mcp.linear.app", ' @@ -1753,7 +1671,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert str(token_request.url) == "https://mcp.linear.app/token" # Send successful token response - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=b'{"access_token": "linear_token", "token_type": "Bearer", "expires_in": 3600}', request=token_request, @@ -1765,7 +1683,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert str(final_request.url) == "https://mcp.linear.app/sse" # Complete flow - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -1800,13 +1718,13 @@ async def callback_handler() -> AuthorizationCodeResult: redirect_uris=[AnyUrl("http://localhost:3030/callback")], ) - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") auth_flow = provider.async_auth_flow(test_request) await auth_flow.__anext__() # 401 with custom WWW-Authenticate PRM URL - response = httpx2.Response( + response = httpx.Response( 401, headers={ "WWW-Authenticate": 'Bearer resource_metadata="https://custom.prm.com/.well-known/oauth-protected-resource"' @@ -1819,28 +1737,28 @@ async def callback_handler() -> AuthorizationCodeResult: assert str(prm_request_1.url) == "https://custom.prm.com/.well-known/oauth-protected-resource" # Returns 500 - prm_response_1 = httpx2.Response(500, request=prm_request_1) + prm_response_1 = httpx.Response(500, request=prm_request_1) # Try path-based fallback prm_request_2 = await auth_flow.asend(prm_response_1) assert str(prm_request_2.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp" # Returns 404 - prm_response_2 = httpx2.Response(404, request=prm_request_2) + prm_response_2 = httpx.Response(404, request=prm_request_2) # Try root fallback prm_request_3 = await auth_flow.asend(prm_response_2) assert str(prm_request_3.url) == "https://api.example.com/.well-known/oauth-protected-resource" # Also returns 404 - all PRM URLs failed - prm_response_3 = httpx2.Response(404, request=prm_request_3) + prm_response_3 = httpx.Response(404, request=prm_request_3) # Should fall back to root OAuth discovery oauth_metadata_request = await auth_flow.asend(prm_response_3) assert str(oauth_metadata_request.url) == "https://api.example.com/.well-known/oauth-authorization-server" # Complete the flow - oauth_metadata_response = httpx2.Response( + oauth_metadata_response = httpx.Response( 200, content=( b'{"issuer": "https://api.example.com", ' @@ -1857,7 +1775,7 @@ async def callback_handler() -> AuthorizationCodeResult: token_request = await auth_flow.asend(oauth_metadata_response) assert str(token_request.url) == "https://api.example.com/token" - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=b'{"access_token": "test_token", "token_type": "Bearer", "expires_in": 3600}', request=token_request, @@ -1866,7 +1784,7 @@ async def callback_handler() -> AuthorizationCodeResult: final_request = await auth_flow.asend(token_response) assert final_request.headers["Authorization"] == "Bearer test_token" - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -1897,8 +1815,8 @@ async def callback_handler() -> AuthorizationCodeResult: ) # Test with 401 response without WWW-Authenticate header - init_response = httpx2.Response( - status_code=401, headers={}, request=httpx2.Request("GET", "https://api.example.com/v1/mcp") + init_response = httpx.Response( + status_code=401, headers={}, request=httpx.Request("GET", "https://api.example.com/v1/mcp") ) # Build discovery URLs @@ -1943,7 +1861,7 @@ async def callback_handler() -> AuthorizationCodeResult: ) # Create a test request - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") # Mock the auth flow auth_flow = provider.async_auth_flow(test_request) @@ -1953,7 +1871,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert "Authorization" not in request.headers # Send a 401 response without WWW-Authenticate header - response = httpx2.Response(401, headers={}, request=test_request) + response = httpx.Response(401, headers={}, request=test_request) # Next request should be to discover protected resource metadata (path-based) discovery_request_1 = await auth_flow.asend(response) @@ -1961,7 +1879,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert discovery_request_1.method == "GET" # Send 404 response for path-based discovery - discovery_response_1 = httpx2.Response(404, request=discovery_request_1) + discovery_response_1 = httpx.Response(404, request=discovery_request_1) # Next request should be to root-based well-known URI discovery_request_2 = await auth_flow.asend(discovery_response_1) @@ -1969,7 +1887,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert discovery_request_2.method == "GET" # Send successful discovery response - discovery_response_2 = httpx2.Response( + discovery_response_2 = httpx.Response( 200, content=( b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}' @@ -1985,7 +1903,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert oauth_metadata_request.method == "GET" # Complete the flow - oauth_metadata_response = httpx2.Response( + oauth_metadata_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com", ' @@ -1996,7 +1914,7 @@ async def callback_handler() -> AuthorizationCodeResult: ) token_request = await auth_flow.asend(oauth_metadata_response) - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=( b'{"access_token": "new_access_token", "token_type": "Bearer", "expires_in": 3600, ' @@ -2006,7 +1924,7 @@ async def callback_handler() -> AuthorizationCodeResult: ) final_request = await auth_flow.asend(token_response) - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -2033,12 +1951,12 @@ async def callback_handler() -> AuthorizationCodeResult: ) # Test with 401 response with WWW-Authenticate header - init_response = httpx2.Response( + init_response = httpx.Response( status_code=401, headers={ "WWW-Authenticate": 'Bearer resource_metadata="https://custom.example.com/.well-known/oauth-protected-resource"' }, - request=httpx2.Request("GET", "https://api.example.com/v1/mcp"), + request=httpx.Request("GET", "https://api.example.com/v1/mcp"), ) # Build discovery URLs @@ -2112,10 +2030,10 @@ def test_extract_field_from_www_auth_valid_cases( ): """Test extraction of various fields from valid WWW-Authenticate headers.""" - init_response = httpx2.Response( + init_response = httpx.Response( status_code=401, headers={"WWW-Authenticate": www_auth_header}, - request=httpx2.Request("GET", "https://api.example.com/test"), + request=httpx.Request("GET", "https://api.example.com/test"), ) result = extract_field_from_www_auth(init_response, field_name) @@ -2147,8 +2065,8 @@ def test_extract_field_from_www_auth_invalid_cases( """Test extraction returns None for invalid cases.""" headers = {"WWW-Authenticate": www_auth_header} if www_auth_header is not None else {} - init_response = httpx2.Response( - status_code=401, headers=headers, request=httpx2.Request("GET", "https://api.example.com/test") + init_response = httpx.Response( + status_code=401, headers=headers, request=httpx.Request("GET", "https://api.example.com/test") ) result = extract_field_from_www_auth(init_response, field_name) @@ -2295,7 +2213,7 @@ async def callback_handler() -> AuthorizationCodeResult: provider.context.token_expiry_time = None provider._initialized = True - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") auth_flow = provider.async_auth_flow(test_request) # First request @@ -2303,11 +2221,11 @@ async def callback_handler() -> AuthorizationCodeResult: assert "Authorization" not in request.headers # Send 401 response - response = httpx2.Response(401, headers={}, request=test_request) + response = httpx.Response(401, headers={}, request=test_request) # PRM discovery prm_request = await auth_flow.asend(response) - prm_response = httpx2.Response( + prm_response = httpx.Response( 200, content=b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}', request=prm_request, @@ -2315,7 +2233,7 @@ async def callback_handler() -> AuthorizationCodeResult: # OAuth metadata discovery oauth_request = await auth_flow.asend(prm_response) - oauth_response = httpx2.Response( + oauth_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com", ' @@ -2346,7 +2264,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert provider.context.client_info.token_endpoint_auth_method == "none" # Complete the flow - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=b'{"access_token": "test_token", "token_type": "Bearer", "expires_in": 3600}', request=token_request, @@ -2355,7 +2273,7 @@ async def callback_handler() -> AuthorizationCodeResult: final_request = await auth_flow.asend(token_response) assert final_request.headers["Authorization"] == "Bearer test_token" - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -2386,18 +2304,18 @@ async def callback_handler() -> AuthorizationCodeResult: provider.context.token_expiry_time = None provider._initialized = True - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") auth_flow = provider.async_auth_flow(test_request) # First request await auth_flow.__anext__() # Send 401 response - response = httpx2.Response(401, headers={}, request=test_request) + response = httpx.Response(401, headers={}, request=test_request) # PRM discovery prm_request = await auth_flow.asend(response) - prm_response = httpx2.Response( + prm_response = httpx.Response( 200, content=b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}', request=prm_request, @@ -2405,7 +2323,7 @@ async def callback_handler() -> AuthorizationCodeResult: # OAuth metadata discovery - server does NOT support CIMD oauth_request = await auth_flow.asend(prm_response) - oauth_response = httpx2.Response( + oauth_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com", ' @@ -2422,7 +2340,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert str(registration_request.url) == "https://auth.example.com/register" # Complete the flow to avoid generator cleanup issues - registration_response = httpx2.Response( + registration_response = httpx.Response( 201, content=b'{"client_id": "dcr_client_id", "redirect_uris": ["http://localhost:3030/callback"]}', request=registration_request, @@ -2434,14 +2352,14 @@ async def callback_handler() -> AuthorizationCodeResult: ) token_request = await auth_flow.asend(registration_response) - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=b'{"access_token": "test_token", "token_type": "Bearer", "expires_in": 3600}', request=token_request, ) final_request = await auth_flow.asend(token_response) - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -2625,7 +2543,7 @@ async def callback_handler() -> AuthorizationCodeResult: redirect_uris=[AnyUrl("http://localhost:3030/callback")], ) - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") auth_flow = provider.async_auth_flow(test_request) # First request @@ -2633,11 +2551,11 @@ async def callback_handler() -> AuthorizationCodeResult: assert "Authorization" not in request.headers # Send 401 - response = httpx2.Response(401, headers={}, request=test_request) + response = httpx.Response(401, headers={}, request=test_request) # PRM discovery prm_request = await auth_flow.asend(response) - prm_response = httpx2.Response( + prm_response = httpx.Response( 200, content=( b'{"resource": "https://api.example.com/v1/mcp",' @@ -2649,7 +2567,7 @@ async def callback_handler() -> AuthorizationCodeResult: # OAuth metadata discovery - AS advertises offline_access oauth_request = await auth_flow.asend(prm_response) - oauth_response = httpx2.Response( + oauth_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com",' @@ -2677,7 +2595,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert params["prompt"][0] == "consent" # Complete the token exchange - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=( b'{"access_token": "new_access_token", "token_type": "Bearer",' @@ -2690,7 +2608,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert final_request.headers["Authorization"] == "Bearer new_access_token" # Close the generator - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -2734,18 +2652,18 @@ async def callback_handler() -> AuthorizationCodeResult: redirect_uris=[AnyUrl("http://localhost:3030/callback")], ) - test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + test_request = httpx.Request("GET", "https://api.example.com/v1/mcp") auth_flow = provider.async_auth_flow(test_request) # First request await auth_flow.__anext__() # Send 401 - response = httpx2.Response(401, headers={}, request=test_request) + response = httpx.Response(401, headers={}, request=test_request) # PRM discovery prm_request = await auth_flow.asend(response) - prm_response = httpx2.Response( + prm_response = httpx.Response( 200, content=( b'{"resource": "https://api.example.com/v1/mcp",' @@ -2757,7 +2675,7 @@ async def callback_handler() -> AuthorizationCodeResult: # OAuth metadata discovery - AS does NOT advertise offline_access oauth_request = await auth_flow.asend(prm_response) - oauth_response = httpx2.Response( + oauth_response = httpx.Response( 200, content=( b'{"issuer": "https://auth.example.com",' @@ -2785,7 +2703,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert "prompt" not in params # Complete the token exchange - token_response = httpx2.Response( + token_response = httpx.Response( 200, content=b'{"access_token": "new_access_token", "token_type": "Bearer", "expires_in": 3600}', request=token_request, @@ -2795,7 +2713,7 @@ async def callback_handler() -> AuthorizationCodeResult: assert final_request.headers["Authorization"] == "Bearer new_access_token" # Close the generator - final_response = httpx2.Response(200, request=final_request) + final_response = httpx.Response(200, request=final_request) try: await auth_flow.asend(final_response) except StopAsyncIteration: @@ -2930,10 +2848,10 @@ async def test_handle_token_response_backfills_omitted_scope_from_request( has reverted to its constructor value. """ oauth_provider.context.client_metadata.scope = "read admin" - response = httpx2.Response( + response = httpx.Response( 200, json={"access_token": "t", "token_type": "Bearer", "expires_in": 3600}, - request=httpx2.Request("POST", "https://auth.example.com/token"), + request=httpx.Request("POST", "https://auth.example.com/token"), ) await oauth_provider._handle_token_response(response) @@ -2944,17 +2862,6 @@ async def test_handle_token_response_backfills_omitted_scope_from_request( assert stored.scope == "read admin" -@pytest.mark.anyio -async def test_handle_token_response_raises_on_non_2xx_with_body(oauth_provider: OAuthClientProvider): - response = httpx2.Response( - 400, - json={"error": "invalid_grant"}, - request=httpx2.Request("POST", "https://auth.example.com/token"), - ) - with pytest.raises(OAuthTokenError, match=r"Token exchange failed \(400\).*invalid_grant"): - await oauth_provider._handle_token_response(response) - - @pytest.mark.anyio async def test_handle_refresh_response_carries_prior_scope_and_refresh_token_when_omitted( oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage @@ -2968,10 +2875,10 @@ async def test_handle_refresh_response_carries_prior_scope_and_refresh_token_whe oauth_provider.context.current_tokens = OAuthToken( access_token="old", scope="read write", refresh_token="prior-refresh" ) - response = httpx2.Response( + response = httpx.Response( 200, json={"access_token": "new", "token_type": "Bearer", "expires_in": 3600}, - request=httpx2.Request("POST", "https://auth.example.com/token"), + request=httpx.Request("POST", "https://auth.example.com/token"), ) ok = await oauth_provider._handle_refresh_response(response) @@ -2994,10 +2901,10 @@ async def test_handle_refresh_response_adopts_rotated_refresh_token_when_returne oauth_provider.context.current_tokens = OAuthToken( access_token="old", scope="read write", refresh_token="prior-refresh" ) - response = httpx2.Response( + response = httpx.Response( 200, json={"access_token": "new", "token_type": "Bearer", "expires_in": 3600, "refresh_token": "rotated"}, - request=httpx2.Request("POST", "https://auth.example.com/token"), + request=httpx.Request("POST", "https://auth.example.com/token"), ) ok = await oauth_provider._handle_refresh_response(response) @@ -3027,20 +2934,20 @@ async def test_issuer_binding_re_evaluated_after_asm_when_prm_discovery_failed( issuer="https://old-as.example.com", ) - auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + auth_flow = oauth_provider.async_auth_flow(httpx.Request("GET", "https://api.example.com/v1/mcp")) request = await auth_flow.__anext__() - response_401 = httpx2.Response(401, request=request) + response_401 = httpx.Response(401, request=request) # PRM discovery: path-based then root, both 404. prm_req = await auth_flow.asend(response_401) assert str(prm_req.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp" - prm_req = await auth_flow.asend(httpx2.Response(404, request=prm_req)) + prm_req = await auth_flow.asend(httpx.Response(404, request=prm_req)) assert str(prm_req.url) == "https://api.example.com/.well-known/oauth-protected-resource" # ASM discovery via root fallback (no auth_server_url) succeeds with a different issuer. - asm_req = await auth_flow.asend(httpx2.Response(404, request=prm_req)) + asm_req = await auth_flow.asend(httpx.Response(404, request=prm_req)) assert str(asm_req.url) == "https://api.example.com/.well-known/oauth-authorization-server" - asm_response = httpx2.Response( + asm_response = httpx.Response( 200, content=( b'{"issuer": "https://api.example.com", ' @@ -3065,12 +2972,12 @@ async def test_issuer_binding_re_evaluated_after_asm_when_prm_discovery_failed( "asm_responses", [ pytest.param( - [httpx2.Response(404), httpx2.Response(404)], + [httpx.Response(404), httpx.Response(404)], id="asm-discovery-failed", ), pytest.param( [ - httpx2.Response( + httpx.Response( 200, content=( b'{"issuer": "https://new-as.example.com", ' @@ -3084,7 +2991,7 @@ async def test_issuer_binding_re_evaluated_after_asm_when_prm_discovery_failed( ], ) async def test_issuer_is_not_stamped_when_registration_falls_back_to_the_resource_origin( - oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, asm_responses: list[httpx2.Response] + oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, asm_responses: list[httpx.Response] ): """SEP-2352: a fallback registration is not recorded as bound to the PRM-advertised AS. @@ -3116,9 +3023,9 @@ async def echo_callback() -> AuthorizationCodeResult: oauth_provider.context.redirect_handler = capture_redirect oauth_provider.context.callback_handler = echo_callback - auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + auth_flow = oauth_provider.async_auth_flow(httpx.Request("GET", "https://api.example.com/v1/mcp")) request = await auth_flow.__anext__() - response_401 = httpx2.Response( + response_401 = httpx.Response( 401, headers={ "WWW-Authenticate": ( @@ -3131,7 +3038,7 @@ async def echo_callback() -> AuthorizationCodeResult: # PRM succeeds and advertises a new AS — the discard block fires. prm_req = await auth_flow.asend(response_401) assert str(prm_req.url) == "https://api.example.com/.well-known/oauth-protected-resource" - prm_response = httpx2.Response( + prm_response = httpx.Response( 200, content=( b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://new-as.example.com"]}' @@ -3153,7 +3060,7 @@ async def echo_callback() -> AuthorizationCodeResult: dcr_req = next_req assert dcr_req.method == "POST" assert str(dcr_req.url) == "https://api.example.com/register" - dcr_response = httpx2.Response( + dcr_response = httpx.Response( 201, json={"client_id": "fallback-client", "redirect_uris": ["http://localhost:3030/callback"]}, request=dcr_req, @@ -3167,12 +3074,12 @@ async def echo_callback() -> AuthorizationCodeResult: assert stored.issuer is None # Drive the flow to completion so the context lock is released cleanly. - token_response = httpx2.Response( + token_response = httpx.Response( 200, json={"access_token": "t", "token_type": "Bearer", "expires_in": 3600}, request=token_req ) final_req = await auth_flow.asend(token_response) try: - await auth_flow.asend(httpx2.Response(200, request=final_req)) + await auth_flow.asend(httpx.Response(200, request=final_req)) except StopAsyncIteration: pass @@ -3205,19 +3112,19 @@ async def echo_callback() -> AuthorizationCodeResult: oauth_provider.context.redirect_handler = capture_redirect oauth_provider.context.callback_handler = echo_callback - auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + auth_flow = oauth_provider.async_auth_flow(httpx.Request("GET", "https://api.example.com/v1/mcp")) request = await auth_flow.__anext__() # PRM discovery 404s on both well-known URLs. - prm_req = await auth_flow.asend(httpx2.Response(401, request=request)) + prm_req = await auth_flow.asend(httpx.Response(401, request=request)) assert str(prm_req.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp" - prm_req = await auth_flow.asend(httpx2.Response(404, request=prm_req)) + prm_req = await auth_flow.asend(httpx.Response(404, request=prm_req)) assert str(prm_req.url) == "https://api.example.com/.well-known/oauth-protected-resource" # Root ASM discovery succeeds with the resource origin as issuer and no registration_endpoint. - asm_req = await auth_flow.asend(httpx2.Response(404, request=prm_req)) + asm_req = await auth_flow.asend(httpx.Response(404, request=prm_req)) assert str(asm_req.url) == "https://api.example.com/.well-known/oauth-authorization-server" - asm_response = httpx2.Response( + asm_response = httpx.Response( 200, content=( b'{"issuer": "https://api.example.com", ' @@ -3231,7 +3138,7 @@ async def echo_callback() -> AuthorizationCodeResult: dcr_req = await auth_flow.asend(asm_response) assert dcr_req.method == "POST" assert str(dcr_req.url) == "https://api.example.com/register" - dcr_response = httpx2.Response( + dcr_response = httpx.Response( 201, json={"client_id": "embedded-client", "redirect_uris": ["http://localhost:3030/callback"]}, request=dcr_req, @@ -3245,11 +3152,276 @@ async def echo_callback() -> AuthorizationCodeResult: assert stored.issuer == str(oauth_provider.context.oauth_metadata.issuer) assert urlparse(stored.issuer).netloc == "api.example.com" - token_response = httpx2.Response( + token_response = httpx.Response( 200, json={"access_token": "t", "token_type": "Bearer", "expires_in": 3600}, request=token_req ) final_req = await auth_flow.asend(token_response) try: - await auth_flow.asend(httpx2.Response(200, request=final_req)) + await auth_flow.asend(httpx.Response(200, request=final_req)) except StopAsyncIteration: pass + + +@pytest.mark.anyio +async def test_concurrent_request_not_blocked_by_pending_long_running_request( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """Regression for #1326: a second request reaches its yield while the + first is still suspended (= simulating a server-side long-poll). + + Before the lock-scope fix, ``async_auth_flow`` held ``context.lock`` + across ``yield request``. A GET SSE long-poll would therefore hold the + lock for the entire SSE lifetime, blocking any concurrent request + waiting on the same provider's lock. + """ + # Set up valid tokens so neither refresh (Phase 2) nor full OAuth + # flow (Phase 4) is triggered — we exercise the steady-state Phase 3 + # yield path that previously held the lock. + 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 + + # Flow 1: drive to yield, then leave suspended (simulating long-poll). + slow_request = httpx.Request("GET", "https://api.example.com/v1/mcp") + slow_flow = oauth_provider.async_auth_flow(slow_request) + yielded_slow = await slow_flow.__anext__() + assert yielded_slow.headers.get("Authorization") == "Bearer test_access_token" + + # Flow 2: concurrent request. With the fix this reaches its yield + # immediately; without the fix it would block on context.lock. + fast_request = httpx.Request("POST", "https://api.example.com/v1/mcp") + fast_flow = oauth_provider.async_auth_flow(fast_request) + with anyio.fail_after(5): + yielded_fast = await fast_flow.__anext__() + assert yielded_fast.headers.get("Authorization") == "Bearer test_access_token" + + with contextlib.suppress(StopAsyncIteration): + await fast_flow.asend(httpx.Response(200, request=yielded_fast)) + with contextlib.suppress(StopAsyncIteration): + await slow_flow.asend(httpx.Response(200, request=yielded_slow)) + + +@pytest.mark.anyio +async def test_refresh_lock_double_check_skips_redundant_refresh( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """Two flows enter Phase 2 with an expired token. After the first + completes a refresh, the second observes the fresh token via the + Phase 2 double-check inside ``refresh_lock`` (or directly in Phase 1 + if it arrives late) and skips its own refresh. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + 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 + + # Flow A: drive to refresh yield, then complete refresh. + request_a = httpx.Request("GET", "https://api.example.com/v1/mcp") + flow_a = oauth_provider.async_auth_flow(request_a) + refresh_a = await flow_a.__anext__() + assert "grant_type=refresh_token" in refresh_a.read().decode() + + refresh_response = httpx.Response( + 200, + content=( + b'{"access_token": "new_access_token", "token_type": "Bearer", ' + b'"expires_in": 3600, "refresh_token": "new_refresh_token"}' + ), + request=refresh_a, + ) + request_a_post = await flow_a.asend(refresh_response) + assert request_a_post.headers.get("Authorization") == "Bearer new_access_token" + + # Flow B: state already refreshed; Phase 1 sees valid token, skips Phase 2. + request_b = httpx.Request("POST", "https://api.example.com/v1/mcp") + flow_b = oauth_provider.async_auth_flow(request_b) + with anyio.fail_after(5): + request_b_yielded = await flow_b.__anext__() + assert request_b_yielded.method == "POST" + assert request_b_yielded.headers.get("Authorization") == "Bearer new_access_token" + + with contextlib.suppress(StopAsyncIteration): + await flow_b.asend(httpx.Response(200, request=request_b_yielded)) + with contextlib.suppress(StopAsyncIteration): + await flow_a.asend(httpx.Response(200, request=request_a_post)) + + +@pytest.mark.anyio +async def test_refresh_with_failed_status_clears_tokens(oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken): + """A non-2xx refresh response clears stored tokens and marks the provider + uninitialized so the next request triggers a full OAuth flow. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 + 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 + + request = httpx.Request("POST", "https://api.example.com/v1/mcp") + flow = oauth_provider.async_auth_flow(request) + refresh_request = await flow.__anext__() + assert "grant_type=refresh_token" in refresh_request.read().decode() + + # Refresh server returns 401. + refresh_response = httpx.Response(401, content=b'{"error": "invalid_grant"}', request=refresh_request) + with contextlib.suppress(StopAsyncIteration): + # After failed refresh, the flow proceeds to Phase 3 yielding the + # original request without a fresh Authorization header. We don't + # exercise the subsequent 401/full OAuth path here. + await flow.asend(refresh_response) + + assert oauth_provider.context.current_tokens is None + + +@pytest.mark.anyio +async def test_refresh_with_invalid_json_clears_tokens(oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken): + """A 200 refresh response with a malformed body clears stored tokens — + the pydantic ValidationError branch is taken. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 + 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 + + request = httpx.Request("POST", "https://api.example.com/v1/mcp") + flow = oauth_provider.async_auth_flow(request) + refresh_request = await flow.__anext__() + + # Body does not parse as OAuthToken. + refresh_response = httpx.Response(200, content=b"not json", request=refresh_request) + with contextlib.suppress(StopAsyncIteration): + await flow.asend(refresh_response) + + assert oauth_provider.context.current_tokens is None + + +@pytest.mark.anyio +async def test_double_check_inside_refresh_lock_skips_second_refresh( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken, monkeypatch: pytest.MonkeyPatch +): + """Exercise the double-check branch inside ``refresh_lock``: ``is_token_valid`` + returns False in Phase 1 (= the flow decides to refresh) but True inside + the inner ``context.lock`` block (= another coroutine refreshed while we + were waiting on ``refresh_lock``). The flow must skip ``_refresh_token`` + and proceed straight to Phase 3. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + 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 + + # Toggle is_token_valid: False on the first call (Phase 1 decision), + # True on the second (double-check inside refresh_lock). + call_count = {"n": 0} + original_is_valid = oauth_provider.context.__class__.is_token_valid + + def fake_is_token_valid(self: object) -> bool: + call_count["n"] += 1 + if call_count["n"] == 1: + return False + # By the second call, "another coroutine" refreshed; reset token expiry + # so callers downstream see a valid token. + oauth_provider.context.token_expiry_time = time.time() + 1800 + return True + + monkeypatch.setattr(oauth_provider.context.__class__, "is_token_valid", fake_is_token_valid) + try: + request = httpx.Request("POST", "https://api.example.com/v1/mcp") + flow = oauth_provider.async_auth_flow(request) + # No refresh yield is expected — the flow goes directly to its own + # request yield with the (now-valid) token header attached. + with anyio.fail_after(5): + yielded = await flow.__anext__() + assert yielded.method == "POST" + assert yielded.headers.get("Authorization") == "Bearer test_access_token" + with contextlib.suppress(StopAsyncIteration): + await flow.asend(httpx.Response(200, request=yielded)) + finally: + monkeypatch.setattr(oauth_provider.context.__class__, "is_token_valid", original_is_valid) + + +@pytest.mark.anyio +async def test_phase1_skips_refresh_when_token_valid(oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken): + """Phase 1 branch where ``is_token_valid()`` is True so ``needs_refresh`` stays + False and Phase 2 is skipped entirely (covers oauth2.py 536->540). + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() + 3600 # valid + 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 + + request = httpx.Request("POST", "https://api.example.com/v1/mcp") + flow = oauth_provider.async_auth_flow(request) + # No refresh yield: Phase 1 sees valid token, skips Phase 2 (536->540 False branch). + yielded = await flow.__anext__() + assert yielded.method == "POST" + assert yielded.headers.get("Authorization") == "Bearer test_access_token" + with contextlib.suppress(StopAsyncIteration): + await flow.asend(httpx.Response(200, request=yielded)) + + +@pytest.mark.anyio +async def test_refresh_success_proceeds_to_phase3_without_resetting_initialized( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """After a successful refresh, ``_handle_refresh_response`` returns True so the + ``_initialized = False`` reset is skipped and Phase 3 proceeds with the fresh + token (covers oauth2.py 553->558 branch where the False arm is taken). + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + 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 + + request = httpx.Request("POST", "https://api.example.com/v1/mcp") + flow = oauth_provider.async_auth_flow(request) + # Phase 2 yields the refresh request. + refresh_request = await flow.__anext__() + assert "grant_type=refresh_token" in refresh_request.read().decode() + + # 200 OK with a parseable token body → _handle_refresh_response returns True. + refresh_response = httpx.Response( + 200, + content=( + b'{"access_token": "fresh_access_token", "token_type": "Bearer", ' + b'"expires_in": 3600, "refresh_token": "fresh_refresh_token"}' + ), + request=refresh_request, + ) + actual_request = await flow.asend(refresh_response) + # Phase 3 yields the *original* request with the fresh Authorization header. + assert actual_request.method == "POST" + assert actual_request.headers.get("Authorization") == "Bearer fresh_access_token" + # _initialized must NOT have been reset (the True branch of _handle_refresh_response). + assert oauth_provider._initialized is True + + with contextlib.suppress(StopAsyncIteration): + await flow.asend(httpx.Response(200, request=actual_request))