From 8007773ec59a1099d5bd681bd643674a594dc03d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:08:36 -0700 Subject: [PATCH] fix(mcp): request offline access from Google upstream OAuth providers (#41197) * fix(mcp): request offline access from Google upstream OAuth providers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(mcp): format authorize_with_server redirect Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): type empty OAuth authorization defaults --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 15 +++-- .../_experimental/mcp_server/oauth_utils.py | 18 +++++- .../mcp_server/test_discoverable_endpoints.py | 60 +++++++++++++++++++ 3 files changed, 84 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cca722f45c2..6e5b5486696 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -83,6 +83,7 @@ from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( CIMD_METADATA_PATH, TOKEN_NO_CACHE_HEADERS, + build_upstream_authorize_url, build_upstream_oauth2_token_request, configured_proxy_base_url, get_cimd_client_id, @@ -918,9 +919,9 @@ def _redirect_to_upstream_authorize( **({"scope": scope_value} if scope_value else {}), **({"resource": upstream_resource} if upstream_resource else {}), } - parsed_auth_url: Final = urlparse(mcp_server.effective_authorization_url or "") - merged_params: Final = {**dict(parse_qsl(parsed_auth_url.query)), **passthrough_params} - return RedirectResponse(urlunparse(parsed_auth_url._replace(query=urlencode(merged_params)))) + return RedirectResponse( + build_upstream_authorize_url(mcp_server.effective_authorization_url or "", passthrough_params) + ) def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MCPServer) -> RedirectResponse: @@ -1151,11 +1152,9 @@ async def authorize_with_server( if upstream_resource: params["resource"] = upstream_resource - parsed_auth_url: Final = urlparse(resolved_server.effective_authorization_url) - existing_params: Final = dict(parse_qsl(parsed_auth_url.query)) - existing_params.update(params) - final_url: Final = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params))) - response: Final = RedirectResponse(final_url) + response: Final = RedirectResponse( + build_upstream_authorize_url(resolved_server.effective_authorization_url, params) + ) _set_oauth_state_cookie(response, request, relay_state, encoded_state) return response diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index b551e099aa5..b0448a62c9e 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -3,10 +3,12 @@ import os import re +from collections.abc import Mapping from dataclasses import dataclass from ipaddress import ip_address +from types import MappingProxyType from typing import TYPE_CHECKING, Final, NoReturn -from urllib.parse import ParseResult, urlparse, urlsplit, urlunparse, urlunsplit +from urllib.parse import ParseResult, parse_qsl, urlencode, urlparse, urlsplit, urlunparse, urlunsplit from fastapi import HTTPException, Request from starlette.types import Scope @@ -56,6 +58,9 @@ _DEFAULT_NATIVE_REDIRECT_URIS: Final[list[str]] = [ "cursor://anysphere.cursor-mcp/oauth/callback", ] +_GOOGLE_AUTHORIZATION_HOSTS: Final = frozenset({"accounts.google.com"}) +_GOOGLE_OFFLINE_ACCESS_PARAMS: Final = MappingProxyType({"access_type": "offline", "prompt": "consent"}) + _warned_invalid_proxy_base_url: str | None = None @@ -112,6 +117,17 @@ def redact_mcp_resource_url(url: str | None) -> str | None: _redact_mcp_resource_url: Final = redact_mcp_resource_url +def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, str]) -> str: + parsed: Final = urlparse(authorization_url) + provider_defaults: Final = ( + _GOOGLE_OFFLINE_ACCESS_PARAMS + if parsed.hostname in _GOOGLE_AUTHORIZATION_HOSTS + else MappingProxyType[str, str]({}) + ) + merged: Final = {**provider_defaults, **dict(parse_qsl(parsed.query)), **params} + return urlunparse(parsed._replace(query=urlencode(merged))) + + def _resolve_proxy_base_url_env() -> str | None: global _warned_invalid_proxy_base_url configured: Final = os.environ.get("PROXY_BASE_URL", "").strip() diff --git a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index aa0166864ce..f09409a023d 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -581,6 +581,66 @@ async def test_authorize_endpoint_preserves_existing_query_params(): assert "scope=read+write" in location +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_mode", ["direct", "true_passthrough", "oauth_delegate"]) +@pytest.mark.parametrize( + ("authorization_url", "expected_access_type", "expected_prompt"), + [ + ("https://accounts.google.com/o/oauth2/v2/auth", ["offline"], ["consent"]), + ("https://accounts.google.com/o/oauth2/v2/auth?prompt=select_account", ["offline"], ["select_account"]), + ("https://accounts.google.com/o/oauth2/v2/auth?access_type=online&prompt=none", ["online"], ["none"]), + ("https://provider.example/oauth/authorize", None, None), + ("https://accounts.google.com.example/oauth/authorize", None, None), + ], +) +async def test_authorize_preserves_contract_with_provider_defaults( + monkeypatch: pytest.MonkeyPatch, + auth_mode: str, + authorization_url: str, + expected_access_type: list[str] | None, + expected_prompt: list[str] | None, +) -> None: + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import authorize_with_server + + monkeypatch.setenv("LITELLM_SALT_KEY", "isolated-oauth-test-salt") + is_bridge: Final = auth_mode != "direct" + server: Final = _bridge_server( + auth_type=MCPAuth(auth_mode) if is_bridge else MCPAuth.oauth2, + dcr_bridge=is_bridge, + client_id=None if is_bridge else "configured-client", + authorization_url=authorization_url + ("&" if "?" in authorization_url else "?") + "scope=configured-scope", + ) + response: Final = await authorize_with_server( + request=_bridge_mock_request(), + mcp_server=server, + client_id="caller-client", + redirect_uri=_BRIDGE_CLIENT_REDIRECT if is_bridge else "http://127.0.0.1:60108/callback", + state="caller-state", + scope="requested-scope", + code_challenge="caller-challenge", + code_challenge_method="S256", + ) + query: Final = parse_qs(urlparse(response.headers["location"]).query) + + assert response.status_code == 307 + assert query.get("access_type") == expected_access_type + assert query.get("prompt") == expected_prompt + assert query["scope"] == ["requested-scope"] + assert query["code_challenge"] == ["caller-challenge"] + assert query["code_challenge_method"] == ["S256"] + assert query["client_id"] == ["caller-client" if is_bridge else "configured-client"] + if is_bridge: + assert query["redirect_uri"] == [_BRIDGE_CLIENT_REDIRECT] + assert query["state"] == ["caller-state"] + assert "set-cookie" not in response.headers + else: + assert query["redirect_uri"] == ["https://litellm.example.com/callback"] + assert query["state"] != ["caller-state"] + assert "set-cookie" in response.headers + + @pytest.mark.asyncio async def test_authorize_endpoint_forwards_pkce_parameters(): """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)"""