diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52b9..411223e0b1 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. """ @@ -9,17 +9,17 @@ import secrets import string import time -from collections.abc import AsyncGenerator, Awaitable, Callable +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping from dataclasses import dataclass, field -from typing import Any, Protocol, get_args -from urllib.parse import quote, urlencode, urljoin, urlparse +from typing import Any, Protocol +from urllib.parse import parse_qsl, quote, urlencode, urljoin, urlparse, urlunparse import anyio -import httpx2 +import httpx from mcp_types.version import is_version_at_least 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, @@ -48,7 +48,6 @@ OAuthMetadata, OAuthToken, ProtectedResourceMetadata, - TokenEndpointAuthMethod, ) from mcp.shared.auth_utils import ( calculate_token_expiry, @@ -59,54 +58,28 @@ 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`. +def _build_authorization_url(auth_endpoint: str, auth_params: Mapping[str, str | None]) -> str: + """Build an authorization URL, preserving any query params already on the endpoint. - 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`. + Servers may advertise an ``authorization_endpoint`` that already carries query + parameters (e.g. ``https://example.com/authorize?prompt=select_account``). + Naively appending ``?`` would produce an invalid URL with two ``?`` + separators, so the existing query is parsed and merged with ``auth_params``. + Flow-generated params take precedence on key conflicts; ``None`` values are + dropped rather than serialized as the literal string ``"None"``. Existing + multi-value query params (e.g. ``?scope=a&scope=b``) are preserved rather + than collapsed, except for keys that the flow overrides. """ - 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" - ) + parsed = urlparse(auth_endpoint) + flow_params = {key: value for key, value in auth_params.items() if value is not None} + # Keep existing endpoint params (including duplicate keys) except those the + # flow overrides, then append the authoritative flow params. + existing = [ + (key, value) for key, value in parse_qsl(parsed.query, keep_blank_values=True) if key not in flow_params + ] + merged_params = existing + list(flow_params.items()) + return urlunparse(parsed._replace(query=urlencode(merged_params))) class PKCEParameters(BaseModel): @@ -153,6 +126,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 @@ -240,12 +214,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 +223,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 +232,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 +256,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 +268,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 +294,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 +331,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) @@ -412,7 +380,7 @@ async def _perform_authorization_code_grant(self) -> tuple[str, str]: if "offline_access" in self.context.client_metadata.scope.split(): auth_params["prompt"] = "consent" - authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}" + authorization_url = _build_authorization_url(auth_endpoint, auth_params) await self.context.redirect_handler(authorization_url) # Wait for callback @@ -438,7 +406,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 +416,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,9 +435,9 @@ 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() @@ -486,7 +459,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 +487,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}") @@ -553,12 +526,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,8 +550,8 @@ 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 async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + """HTTPX auth flow integration.""" async with self.context.lock: if not self._initialized: await self._initialize() @@ -723,7 +696,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 diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8eec..3194e05a9f 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -6,13 +6,14 @@ from unittest import mock from urllib.parse import parse_qs, quote, unquote, urlparse -import httpx2 +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, OAuthTokenError +from mcp.client.auth.oauth2 import _build_authorization_url from mcp.client.auth.utils import ( build_oauth_authorization_server_metadata_discovery_urls, build_protected_resource_metadata_discovery_urls, @@ -111,7 +112,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 +125,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 +138,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 +193,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 +291,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 +409,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 +419,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 +434,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 +449,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 +461,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 +473,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 +491,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 +507,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 +522,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 +533,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 +553,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 +573,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 +982,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 +1010,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 +1069,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 +1083,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 +1093,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 +1108,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 +1118,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 +1132,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 +1145,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 +1162,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 +1180,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 +1196,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 +1218,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 +1230,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 +1261,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 +1271,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 +1285,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 +1298,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 +1315,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 +1333,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 +1349,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 +1406,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 +1427,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 +1442,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 +1486,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 +1498,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 +1620,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 +1628,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 +1650,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 +1670,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 +1682,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 +1717,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 +1736,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 +1774,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 +1783,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 +1814,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 +1860,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 +1870,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 +1878,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 +1886,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 +1902,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 +1913,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 +1923,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 +1950,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 +2029,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 +2064,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 +2212,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 +2220,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 +2232,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 +2263,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 +2272,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 +2303,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 +2322,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 +2339,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 +2351,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 +2542,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 +2550,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 +2566,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 +2594,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 +2607,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 +2651,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 +2674,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 +2702,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 +2712,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 +2847,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) @@ -2946,10 +2863,10 @@ async def test_handle_token_response_backfills_omitted_scope_from_request( @pytest.mark.anyio async def test_handle_token_response_raises_on_non_2xx_with_body(oauth_provider: OAuthClientProvider): - response = httpx2.Response( + response = httpx.Response( 400, json={"error": "invalid_grant"}, - request=httpx2.Request("POST", "https://auth.example.com/token"), + request=httpx.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) @@ -2968,10 +2885,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 +2911,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 +2944,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 +2982,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 +3001,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 +3033,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 +3048,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 +3070,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 +3084,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 +3122,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 +3148,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 +3162,105 @@ 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 + + +class TestAuthorizationEndpointWithQuery: + """Regression tests for #2776 - authorization_endpoint carrying query params.""" + + def test_build_authorization_url_no_existing_query(self): + url = _build_authorization_url( + "https://auth.example.com/authorize", + {"response_type": "code", "client_id": "abc"}, + ) + parsed = urlparse(url) + params = parse_qs(parsed.query) + assert parsed.path == "/authorize" + assert params["response_type"] == ["code"] + assert params["client_id"] == ["abc"] + # No malformed double "?" separator. + assert url.count("?") == 1 + + def test_build_authorization_url_preserves_existing_query(self): + # e.g. Salesforce advertises .../authorize?prompt=select_account + url = _build_authorization_url( + "https://test.salesforce.com/services/oauth2/authorize?prompt=select_account", + {"response_type": "code", "client_id": "abc"}, + ) + parsed = urlparse(url) + params = parse_qs(parsed.query) + assert parsed.path == "/services/oauth2/authorize" + # The server-provided param survives... + assert params["prompt"] == ["select_account"] + # ...alongside the flow-generated params. + assert params["response_type"] == ["code"] + assert params["client_id"] == ["abc"] + # Exactly one "?" - the old f-string produced "...?prompt=...?response_type=...". + assert url.count("?") == 1 + + def test_build_authorization_url_flow_params_win_on_conflict(self): + url = _build_authorization_url( + "https://auth.example.com/authorize?response_type=token", + {"response_type": "code"}, + ) + params = parse_qs(urlparse(url).query) + assert params["response_type"] == ["code"] + + def test_build_authorization_url_preserves_multi_value_existing_query(self): + """Duplicate keys on the endpoint query are preserved, not collapsed.""" + url = _build_authorization_url( + "https://auth.example.com/authorize?scope=a&scope=b", + {"response_type": "code"}, + ) + params = parse_qs(urlparse(url).query) + assert params["scope"] == ["a", "b"] + assert params["response_type"] == ["code"] + assert url.count("?") == 1 + + @pytest.mark.anyio + async def test_perform_authorization_preserves_endpoint_query(self, oauth_provider: OAuthClientProvider): + """End-to-end: redirect URL stays valid when the endpoint has a query string.""" + oauth_provider.context.oauth_metadata = OAuthMetadata( + issuer=AnyHttpUrl("https://test.salesforce.com"), + authorization_endpoint=AnyHttpUrl( + "https://test.salesforce.com/services/oauth2/authorize?prompt=select_account" + ), + token_endpoint=AnyHttpUrl("https://test.salesforce.com/services/oauth2/token"), + ) + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="test_client_id", + client_secret="test_client_secret", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + ) + + captured_url: str | None = None + captured_state: str | None = None + + async def capture_redirect(url: str) -> None: + nonlocal captured_url, captured_state + captured_url = url + captured_state = parse_qs(urlparse(url).query).get("state", [None])[0] + + async def mock_callback() -> AuthorizationCodeResult: + return AuthorizationCodeResult(code="test_auth_code", state=captured_state) + + oauth_provider.context.redirect_handler = capture_redirect + oauth_provider.context.callback_handler = mock_callback + + await oauth_provider._perform_authorization_code_grant() + + assert captured_url is not None + parsed = urlparse(captured_url) + params = parse_qs(parsed.query) + assert parsed.path == "/services/oauth2/authorize" + assert params["prompt"] == ["select_account"] + assert params["response_type"] == ["code"] + assert params["client_id"] == ["test_client_id"] + assert captured_url.count("?") == 1