Fix OAuth redirect URLs with existing query params

Author: he-yufeng

src/mcp/client/auth/oauth2.py

@@ -12,7 +12,7 @@
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass, field
 from typing import Any, Protocol
-from urllib.parse import quote, urlencode, urljoin, urlparse
+from urllib.parse import parse_qsl, quote, urlencode, urljoin, urlparse, urlunparse
 
 import anyio
 import httpx
@@ -53,6 +53,12 @@
 logger = logging.getLogger(__name__)
 
 
+def _url_with_query_params(url: str, params: dict[str, str]) -> str:
+    parsed = urlparse(url)
+    query = urlencode([*parse_qsl(parsed.query, keep_blank_values=True), *params.items()])
+    return urlunparse(parsed._replace(query=query))
+
+
 class PKCEParameters(BaseModel):
     """PKCE (Proof Key for Code Exchange) parameters."""
 
@@ -325,14 +331,14 @@ async def _perform_authorization_code_grant(self) -> tuple[str, str]:
             auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
             auth_endpoint = urljoin(auth_base_url, "/authorize")
 
-        if not self.context.client_info:
+        if not self.context.client_info or not self.context.client_info.client_id:
             raise OAuthFlowError("No client info available for authorization")  # pragma: no cover
 
         # Generate PKCE parameters
         pkce_params = PKCEParameters.generate()
         state = secrets.token_urlsafe(32)
 
-        auth_params = {
+        auth_params: dict[str, str] = {
             "response_type": "code",
             "client_id": self.context.client_info.client_id,
             "redirect_uri": str(self.context.client_metadata.redirect_uris[0]),
@@ -353,7 +359,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 = _url_with_query_params(auth_endpoint, auth_params)
         await self.context.redirect_handler(authorization_url)
 
         # Wait for callback

tests/client/test_auth.py

@@ -1167,6 +1167,47 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide
         assert oauth_provider.context.current_tokens.access_token == "new_access_token"
         assert oauth_provider.context.token_expiry_time is not None
 
+    @pytest.mark.anyio
+    async def test_authorization_endpoint_preserves_existing_query(self, oauth_provider: OAuthClientProvider):
+        captured_auth_url: str | None = None
+        captured_state: str | None = None
+
+        async def redirect_handler(url: str) -> None:
+            nonlocal captured_auth_url, captured_state
+            captured_auth_url = url
+            captured_state = parse_qs(urlparse(url).query)["state"][0]
+
+        async def callback_handler() -> tuple[str, str | None]:
+            return "test_auth_code", captured_state
+
+        oauth_provider.context.redirect_handler = redirect_handler
+        oauth_provider.context.callback_handler = callback_handler
+        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",
+            client_secret="test_secret",
+            redirect_uris=[AnyUrl("http://localhost:3030/callback")],
+        )
+
+        auth_code, _ = await oauth_provider._perform_authorization_code_grant()
+
+        assert auth_code == "test_auth_code"
+        assert captured_auth_url is not None
+        parsed = urlparse(captured_auth_url)
+        params = parse_qs(parsed.query)
+        assert parsed.path == "/services/oauth2/authorize"
+        assert "?" not in parsed.query
+        assert params["prompt"] == ["select_account"]
+        assert params["response_type"] == ["code"]
+        assert params["client_id"] == ["test_client"]
+        assert params["redirect_uri"] == ["http://localhost:3030/callback"]
+
     @pytest.mark.anyio
     async def test_auth_flow_no_unnecessary_retry_after_oauth(
         self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken