From 0e4d0a103515d2c91894e60e0807d952c4f0585c Mon Sep 17 00:00:00 2001 From: Jack Venberg Date: Wed, 18 Mar 2026 16:12:09 -0700 Subject: [PATCH] fix(mcp): support refresh_token grant type in OAuth token endpoint (#23701) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): support refresh_token grant type in OAuth token endpoint (#23700) The .well-known/oauth-authorization-server metadata advertises refresh_token as a supported grant type, but the token endpoint rejected it with HTTP 400. This adds refresh_token grant support so MCP clients can refresh expired tokens without re-authenticating. * test(mcp): add tests for refresh_token grant type in OAuth token endpoint * fix(mcp): move code_verifier guard into authorization_code branch code_verifier is only relevant for authorization_code grants (PKCE). Move it inside the else branch so it doesn't apply to refresh_token. * fix(mcp): guard None client_secret and forward scope in token exchange - Conditionally include client_secret in form data to prevent httpx from sending the literal string "None" (applies to both authorization_code and refresh_token branches) - Forward optional scope parameter per RFC 6749 ยง6, allowing clients to request a subset of originally-granted scopes on refresh * fix(mcp): validate code param in authorization_code grant Guard against None code being form-encoded as literal string "None" by httpx, symmetric with the existing refresh_token guard. --- .../mcp_server/discoverable_endpoints.py | 56 +++++-- .../mcp_management_endpoints.py | 4 + .../mcp_server/test_discoverable_endpoints.py | 138 ++++++++++++++++++ .../test_mcp_management_endpoints.py | 52 +++++++ 4 files changed, 237 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3385e7feef6..4dec41a4251 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -208,26 +208,52 @@ async def exchange_token_with_server( client_id: str, client_secret: Optional[str], code_verifier: Optional[str], + refresh_token: Optional[str] = None, + scope: Optional[str] = None, ): - if grant_type != "authorization_code": + if grant_type not in ("authorization_code", "refresh_token"): raise HTTPException(status_code=400, detail="Unsupported grant_type") if mcp_server.token_url is None: raise HTTPException(status_code=400, detail="MCP server token url is not set") - proxy_base_url = get_request_base_url(request) - token_data = { - "grant_type": "authorization_code", - "client_id": mcp_server.client_id if mcp_server.client_id else client_id, - "client_secret": mcp_server.client_secret - if mcp_server.client_secret - else client_secret, - "code": code, - "redirect_uri": f"{proxy_base_url}/callback", - } + resolved_client_id = mcp_server.client_id if mcp_server.client_id else client_id + resolved_client_secret = ( + mcp_server.client_secret if mcp_server.client_secret else client_secret + ) - if code_verifier: - token_data["code_verifier"] = code_verifier + if grant_type == "refresh_token": + if not refresh_token: + raise HTTPException( + status_code=400, + detail="refresh_token is required for refresh_token grant", + ) + token_data: dict = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": resolved_client_id, + } + if resolved_client_secret is not None: + token_data["client_secret"] = resolved_client_secret + if scope: + token_data["scope"] = scope + else: + if not code: + raise HTTPException( + status_code=400, + detail="code is required for authorization_code grant", + ) + proxy_base_url = get_request_base_url(request) + token_data: dict = { + "grant_type": "authorization_code", + "client_id": resolved_client_id, + "code": code, + "redirect_uri": f"{proxy_base_url}/callback", + } + if resolved_client_secret is not None: + token_data["client_secret"] = resolved_client_secret + if code_verifier: + token_data["code_verifier"] = code_verifier async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) response = await async_client.post( @@ -375,6 +401,8 @@ async def token_endpoint( client_id: str = Form(...), client_secret: Optional[str] = Form(None), code_verifier: str = Form(None), + refresh_token: Optional[str] = Form(None), + scope: Optional[str] = Form(None), mcp_server_name: Optional[str] = None, ): """ @@ -408,6 +436,8 @@ async def token_endpoint( client_id=client_id, client_secret=client_secret, code_verifier=code_verifier, + refresh_token=refresh_token, + scope=scope, ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3e5b729cea6..f29a721ede8 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1399,6 +1399,8 @@ if MCP_AVAILABLE: client_id: Optional[str] = Form(None), client_secret: Optional[str] = Form(None), code_verifier: Optional[str] = Form(None), + refresh_token: Optional[str] = Form(None), + scope: Optional[str] = Form(None), ): mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) resolved_client_id = mcp_server.client_id or client_id or "" @@ -1422,6 +1424,8 @@ if MCP_AVAILABLE: client_id=resolved_client_id, client_secret=client_secret, code_verifier=code_verifier, + refresh_token=refresh_token, + scope=scope, ) @router.post( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 700ba86b108..954f2703e3b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -1666,3 +1666,141 @@ async def test_oauth_authorize_prefers_request_scope_over_server_config(): redirect_url = response.headers["location"] assert "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url assert "default_scope" not in redirect_url + + +@pytest.mark.asyncio +async def test_token_endpoint_refresh_token_grant(): + """Test that token endpoint supports refresh_token grant type.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + token_endpoint, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="google_mcp", + name="google_mcp", + server_name="google_mcp", + alias="google_mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_secret", + authorization_url="https://accounts.google.com/o/oauth2/v2/auth", + token_url="https://oauth2.googleapis.com/token", + scopes=["openid", "email"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.litellm.example/" + mock_request.headers = {} + + # Mock httpx client response with new tokens + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "new_access_token", + "token_type": "Bearer", + "expires_in": 3599, + "refresh_token": "new_refresh_token", + } + mock_response.raise_for_status = MagicMock() + + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client: + mock_get_client.return_value = mock_async_client + + response = await token_endpoint( + request=mock_request, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="test_client_id", + mcp_server_name="google_mcp", + client_secret="test_secret", + refresh_token="rt-test", + scope="openid email", + ) + + # Verify the POST was called with refresh_token grant data + mock_async_client.post.assert_called_once() + call_args = mock_async_client.post.call_args + + assert call_args[1]["data"]["grant_type"] == "refresh_token" + assert call_args[1]["data"]["refresh_token"] == "rt-test" + assert call_args[1]["data"]["client_id"] == "test_client_id" + assert call_args[1]["data"]["client_secret"] == "test_secret" + assert call_args[1]["data"]["scope"] == "openid email" + + # Verify response contains the new tokens + import json + + token_data = json.loads(response.body) + assert token_data["access_token"] == "new_access_token" + assert token_data["refresh_token"] == "new_refresh_token" + + +@pytest.mark.asyncio +async def test_token_endpoint_authorization_code_missing_code(): + """Test that authorization_code grant rejects missing code param.""" + try: + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + + server = MCPServer( + server_id="test_server", + name="test_server", + server_name="test_server", + alias="test_server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + token_url="https://example.com/token", + ) + global_mcp_server_manager.registry[server.server_id] = server + + mock_request = MagicMock() + mock_request.base_url = "https://proxy.example/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code=None, + redirect_uri="https://example.com/cb", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + assert exc_info.value.status_code == 400 + assert "code is required" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index eeaeb498328..77ac3a040ac 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1519,6 +1519,8 @@ class TestTemporaryMCPSessionEndpoints: client_id="client", client_secret="secret", code_verifier="verifier", + refresh_token=None, + scope=None, ) assert result is exchange_response @@ -1532,6 +1534,56 @@ class TestTemporaryMCPSessionEndpoints: client_id="client", client_secret="secret", code_verifier="verifier", + refresh_token=None, + scope=None, + ) + + @pytest.mark.asyncio + async def test_mcp_token_proxies_refresh_token_grant(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + mcp_token, + ) + + request = MagicMock() + server = generate_mock_mcp_server_config_record(server_id="server-1") + exchange_response = {"access_token": "new-token", "refresh_token": "new-rt"} + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", + return_value=server, + ) as get_server, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.exchange_token_with_server", + AsyncMock(return_value=exchange_response), + ) as exchange_mock, + ): + result = await mcp_token( + request=request, + server_id="server-1", + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="client", + client_secret="secret", + code_verifier=None, + refresh_token="rt-123", + scope=None, + ) + + assert result is exchange_response + get_server.assert_called_once_with("server-1") + exchange_mock.assert_awaited_once_with( + request=request, + mcp_server=server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="client", + client_secret="secret", + code_verifier=None, + refresh_token="rt-123", + scope=None, ) @pytest.mark.asyncio