fix(auth): preserve user-agent through oauth flow

Author: go165

src/mcp/client/auth/oauth2.py

@@ -9,7 +9,7 @@
 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
 from urllib.parse import quote, urlencode, urljoin, urlparse
@@ -303,10 +303,10 @@ async def _handle_protected_resource_response(self, response: httpx.Response) ->
                 f"Protected Resource Metadata request failed: {response.status_code}"
             )  # pragma: no cover
 
-    async def _perform_authorization(self) -> httpx.Request:
+    async def _perform_authorization(self, headers: Mapping[str, str] | None = None) -> 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)
+        token_request = await self._exchange_token_authorization_code(auth_code, code_verifier, headers=headers)
         return token_request
 
     async def _perform_authorization_code_grant(self) -> tuple[str, str]:
@@ -376,7 +376,12 @@ def _get_token_endpoint(self) -> str:
         return token_url
 
     async def _exchange_token_authorization_code(
-        self, auth_code: str, code_verifier: str, *, token_data: dict[str, Any] | None = {}
+        self,
+        auth_code: str,
+        code_verifier: str,
+        *,
+        token_data: dict[str, Any] | None = {},
+        headers: Mapping[str, str] | None = None,
     ) -> httpx.Request:
         """Build token exchange request for authorization_code flow."""
         if self.context.client_metadata.redirect_uris is None:
@@ -401,10 +406,12 @@ async def _exchange_token_authorization_code(
             token_data["resource"] = self.context.get_resource_url()  # RFC 8707
 
         # Prepare authentication based on preferred method
-        headers = {"Content-Type": "application/x-www-form-urlencoded"}
-        token_data, headers = self.context.prepare_token_auth(token_data, headers)
+        request_headers = {"Content-Type": "application/x-www-form-urlencoded"}
+        if headers:
+            request_headers.update(headers)
+        token_data, request_headers = self.context.prepare_token_auth(token_data, request_headers)
 
-        return httpx.Request("POST", token_url, data=token_data, headers=headers)
+        return httpx.Request("POST", token_url, data=token_data, headers=request_headers)
 
     async def _handle_token_response(self, response: httpx.Response) -> None:
         """Handle token exchange response."""
@@ -421,7 +428,7 @@ async def _handle_token_response(self, response: httpx.Response) -> None:
         self.context.update_token_expiry(token_response)
         await self.context.storage.set_tokens(token_response)
 
-    async def _refresh_token(self) -> httpx.Request:
+    async def _refresh_token(self, headers: Mapping[str, str] | None = None) -> 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
@@ -446,10 +453,12 @@ async def _refresh_token(self) -> httpx.Request:
             refresh_data["resource"] = self.context.get_resource_url()  # RFC 8707
 
         # Prepare authentication based on preferred method
-        headers = {"Content-Type": "application/x-www-form-urlencoded"}
-        refresh_data, headers = self.context.prepare_token_auth(refresh_data, headers)
+        request_headers = {"Content-Type": "application/x-www-form-urlencoded"}
+        if headers:
+            request_headers.update(headers)
+        refresh_data, request_headers = self.context.prepare_token_auth(refresh_data, request_headers)
 
-        return httpx.Request("POST", token_url, data=refresh_data, headers=headers)
+        return httpx.Request("POST", token_url, data=refresh_data, headers=request_headers)
 
     async def _handle_refresh_response(self, response: httpx.Response) -> bool:
         """Handle token refresh response. Returns True if successful."""
@@ -483,6 +492,11 @@ def _add_auth_header(self, request: httpx.Request) -> None:
         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}"
 
+    def _auth_flow_request_headers(self, request: httpx.Request) -> dict[str, str]:
+        """Headers that are safe to carry from the MCP request into OAuth requests."""
+        user_agent = request.headers.get("User-Agent")
+        return {"User-Agent": user_agent} if user_agent else {}
+
     async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None:
         content = await response.aread()
         metadata = OAuthMetadata.model_validate_json(content)
@@ -510,10 +524,11 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
 
             # Capture protocol version from request headers
             self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION)
+            auth_flow_headers = self._auth_flow_request_headers(request)
 
             if not self.context.is_token_valid() and self.context.can_refresh_token():
                 # Try to refresh token
-                refresh_request = await self._refresh_token()
+                refresh_request = await self._refresh_token(headers=auth_flow_headers)
                 refresh_response = yield refresh_request
 
                 if not await self._handle_refresh_response(refresh_response):
@@ -537,7 +552,7 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                     )
 
                     for url in prm_discovery_urls:  # pragma: no branch
-                        discovery_request = create_oauth_metadata_request(url)
+                        discovery_request = create_oauth_metadata_request(url, headers=auth_flow_headers)
 
                         discovery_response = yield discovery_request  # sending request
 
@@ -563,7 +578,7 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
 
                     # Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers)
                     for url in asm_discovery_urls:  # pragma: no branch
-                        oauth_metadata_request = create_oauth_metadata_request(url)
+                        oauth_metadata_request = create_oauth_metadata_request(url, headers=auth_flow_headers)
                         oauth_metadata_response = yield oauth_metadata_request
 
                         ok, asm = await handle_auth_metadata_response(oauth_metadata_response)
@@ -602,14 +617,15 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                                 self.context.oauth_metadata,
                                 self.context.client_metadata,
                                 self.context.get_authorization_base_url(self.context.server_url),
+                                headers=auth_flow_headers,
                             )
                             registration_response = yield registration_request
                             client_information = await handle_registration_response(registration_response)
                             self.context.client_info = client_information
                             await self.context.storage.set_client_info(client_information)
 
                     # Step 5: Perform authorization and complete token exchange
-                    token_response = yield await self._perform_authorization()
+                    token_response = yield await self._perform_authorization(headers=auth_flow_headers)
                     await self._handle_token_response(token_response)
                 except Exception:
                     logger.exception("OAuth flow error")
@@ -634,7 +650,7 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                         )
 
                         # Step 2b: Perform (re-)authorization and token exchange
-                        token_response = yield await self._perform_authorization()
+                        token_response = yield await self._perform_authorization(headers=auth_flow_headers)
                         await self._handle_token_response(token_response)
                     except Exception:  # pragma: no cover
                         logger.exception("OAuth flow error")

src/mcp/client/auth/utils.py

@@ -1,4 +1,5 @@
 import re
+from collections.abc import Mapping
 from urllib.parse import urljoin, urlparse
 
 from httpx import Request, Response
@@ -211,12 +212,18 @@ async def handle_auth_metadata_response(response: Response) -> tuple[bool, OAuth
     return True, None
 
 
-def create_oauth_metadata_request(url: str) -> Request:
-    return Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
+def create_oauth_metadata_request(url: str, headers: Mapping[str, str] | None = None) -> Request:
+    request_headers = {MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION}
+    if headers:
+        request_headers.update(headers)
+    return Request("GET", url, headers=request_headers)
 
 
 def create_client_registration_request(
-    auth_server_metadata: OAuthMetadata | None, client_metadata: OAuthClientMetadata, auth_base_url: str
+    auth_server_metadata: OAuthMetadata | None,
+    client_metadata: OAuthClientMetadata,
+    auth_base_url: str,
+    headers: Mapping[str, str] | None = None,
 ) -> Request:
     """Build a client registration request."""
 
@@ -227,7 +234,11 @@ def create_client_registration_request(
 
     registration_data = client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True)
 
-    return Request("POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"})
+    request_headers = {"Content-Type": "application/json"}
+    if headers:
+        request_headers.update(headers)
+
+    return Request("POST", registration_url, json=registration_data, headers=request_headers)
 
 
 async def handle_registration_response(response: Response) -> OAuthClientInformationFull:

tests/interaction/auth/_harness.py

@@ -394,6 +394,7 @@ async def connect_with_oauth(
     client_metadata_url: str | None = None,
     headless: HeadlessOAuth | None = None,
     auth: httpx.Auth | None = None,
+    headers: Mapping[str, str] | None = None,
     verify_tokens: bool = True,
     app_shim: Callable[[ASGIApp], ASGIApp] | None = None,
     on_request: Callable[[httpx.Request], None] | None = None,
@@ -455,7 +456,11 @@ async def hook(request: httpx.Request) -> None:
         await stack.enter_async_context(server.session_manager.run())
         http_client = await stack.enter_async_context(
             httpx.AsyncClient(
-                transport=StreamingASGITransport(app), base_url=BASE_URL, auth=oauth, event_hooks=event_hooks
+                transport=StreamingASGITransport(app),
+                base_url=BASE_URL,
+                auth=oauth,
+                headers=headers,
+                event_hooks=event_hooks,
             )
         )
         headless.bind(http_client)

tests/interaction/auth/test_flow.py

@@ -118,6 +118,34 @@ async def test_an_unauthenticated_request_is_challenged_then_the_full_oauth_flow
     assert storage.tokens.access_token in provider.access_tokens
 
 
+@requirement("client-transport:http:custom-headers")
+async def test_oauth_flow_preserves_custom_user_agent_on_auth_requests() -> None:
+    """OAuth requests keep the caller's User-Agent from the Streamable HTTP client."""
+    requests: list[httpx.Request] = []
+    server = Server("guarded", on_list_tools=list_tools)
+    user_agent = "mcp-python-sdk/issue-1664"
+
+    with anyio.fail_after(5):
+        async with connect_with_oauth(
+            server,
+            provider=InMemoryAuthorizationServerProvider(),
+            headers={"User-Agent": user_agent},
+            on_request=requests.append,
+        ) as (client, _):
+            await client.list_tools()
+
+    auth_paths = {
+        "/.well-known/oauth-protected-resource/mcp",
+        "/.well-known/oauth-authorization-server",
+        "/register",
+        "/token",
+    }
+    auth_requests = [request for request in requests if request.url.path in auth_paths]
+
+    assert {request.url.path for request in auth_requests} == auth_paths
+    assert all(request.headers["user-agent"] == user_agent for request in auth_requests)
+
+
 @requirement("hosting:auth:authinfo-propagates")
 async def test_the_access_token_reaches_the_tool_handler_via_get_access_token() -> None:
     """A tool handler reads the request's access token through `get_access_token()`."""