diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index fbef33c32ed..871f6d60e70 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -126,6 +126,12 @@ def encrypt_credentials( new_encryption_key=encryption_key, ) # aws_region_name and aws_service_name are NOT secrets — stored as-is + refresh_token = credentials.get("refresh_token") + if refresh_token is not None: + credentials["refresh_token"] = encrypt_value_helper( + value=refresh_token, + new_encryption_key=encryption_key, + ) return credentials @@ -137,6 +143,7 @@ def decrypt_credentials( "auth_value", "client_id", "client_secret", + "refresh_token", "aws_access_key_id", "aws_secret_access_key", "aws_session_token", diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 07309eb57f2..2736a71401b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -5,6 +5,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -277,9 +278,148 @@ async def exchange_token_with_server( if "scope" in token_response and token_response["scope"]: result["scope"] = token_response["scope"] + # Persist OAuth credentials and discovery URLs to the database. + # When the upstream server does not rotate refresh tokens (common pattern), + # the token response omits refresh_token — fall back to the incoming + # refresh_token parameter so the previously stored token is preserved. + persisted_refresh_token = token_response.get("refresh_token") or refresh_token + await _persist_oauth_credentials( + mcp_server=mcp_server, + client_id=token_data["client_id"], + client_secret=token_data.get("client_secret"), + access_token=access_token, + refresh_token=persisted_refresh_token, + expires_in=token_response.get("expires_in"), + ) + return JSONResponse(result) +async def _persist_oauth_credentials( + mcp_server: MCPServer, + client_id: Optional[str], + client_secret: Optional[str], + access_token: str, + refresh_token: Optional[str], + expires_in: Optional[int], +) -> None: + """ + Persist OAuth credentials and discovery URLs to the MCP server record + in the database after a successful token exchange. + + Merges new credential fields on top of any existing credentials in the DB, + following the same pattern as ``update_mcp_server`` to avoid wiping fields + that may have been set independently (e.g. by an admin). + """ + try: + from datetime import datetime, timedelta, timezone + + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + verbose_logger.warning( + "Cannot persist OAuth credentials: no database connected" + ) + return + + import json + + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + from litellm.proxy.common_utils.encrypt_decrypt_utils import _get_salt_key + + # Build new credential fields from the token exchange + new_credentials: dict = {} + if client_id: + new_credentials["client_id"] = client_id + if client_secret: + new_credentials["client_secret"] = client_secret + new_credentials["auth_value"] = access_token + if refresh_token: + new_credentials["refresh_token"] = refresh_token + if expires_in is not None: + new_credentials["expires_at"] = ( + datetime.now(timezone.utc) + timedelta(seconds=expires_in) + ).isoformat() + new_credentials["type"] = "oauth2" + + # Resolve the persistent DB record for this server. + # During an OAuth PKCE session flow the MCPServer object carries an + # ephemeral session UUID as server_id (not the real DB row id). Try + # the exact server_id first; if not found, fall back to server_name / + # alias so the credentials land on the correct persistent row. + existing_record = await prisma_client.db.litellm_mcpservertable.find_unique( + where={"server_id": mcp_server.server_id} + ) + if existing_record is None and mcp_server.name: + existing_record = await prisma_client.db.litellm_mcpservertable.find_first( + where={ + "OR": [ + {"server_name": mcp_server.name}, + {"alias": mcp_server.name}, + ] + } + ) + + if existing_record is None: + verbose_logger.warning( + "Cannot persist OAuth credentials: no DB record found for MCP server %s (%s)", + mcp_server.name, + mcp_server.server_id, + ) + return + + # Merge new credentials on top of existing to preserve unrelated fields + if existing_record.credentials: + existing_creds = ( + json.loads(existing_record.credentials) + if isinstance(existing_record.credentials, str) + else dict(existing_record.credentials) + ) + merged = {**existing_creds, **new_credentials} + else: + merged = new_credentials + + encrypted_credentials = encrypt_credentials( + credentials=merged, + encryption_key=_get_salt_key(), + ) + + # Build the update data — persist credentials and discovery URLs + update_data: dict = { + "auth_type": "oauth2", + "credentials": safe_dumps(encrypted_credentials), + "updated_by": "oauth_token_exchange", + } + + # Persist discovery URLs if they were found during the flow + if mcp_server.token_url: + update_data["token_url"] = mcp_server.token_url + if mcp_server.authorization_url: + update_data["authorization_url"] = mcp_server.authorization_url + if mcp_server.registration_url: + update_data["registration_url"] = mcp_server.registration_url + + await prisma_client.db.litellm_mcpservertable.update( + where={"server_id": existing_record.server_id}, + data=update_data, + ) + + verbose_logger.info( + "Persisted OAuth credentials for MCP server %s (db_id=%s, session_id=%s)", + mcp_server.name, + existing_record.server_id, + mcp_server.server_id, + ) + except Exception as e: + verbose_logger.error( + "Failed to persist OAuth credentials for MCP server %s: %s", + mcp_server.server_id, + e, + ) + + async def register_client_with_server( request: Request, mcp_server: MCPServer, 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 954f2703e3b..4ff57b0c3f6 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 @@ -1804,3 +1804,284 @@ async def test_token_endpoint_authorization_code_missing_code(): ) assert exc_info.value.status_code == 400 assert "code is required" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_token_exchange_persists_credentials_to_db(): + """Test that token exchange persists OAuth credentials and discovery URLs to the database.""" + try: + import json + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + 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") + + oauth2_server = MCPServer( + server_id="test-persist-server", + name="test_persist", + server_name="test_persist", + alias="test_persist", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + authorization_url="https://auth.example.com/authorize", + registration_url="https://auth.example.com/register", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm-proxy.example.com/" + mock_request.headers = {} + + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "test_access_token_123", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "test_refresh_token_456", + "scope": "tools", + } + mock_response.raise_for_status = MagicMock() + + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None) + # No existing record — first token exchange + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client, patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), patch( + "litellm.proxy.common_utils.encrypt_decrypt_utils._get_salt_key", + return_value="test-salt-key", + ): + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_async_client + + response = await exchange_token_with_server( + request=mock_request, + mcp_server=oauth2_server, + grant_type="authorization_code", + code="test_auth_code", + redirect_uri="https://litellm-proxy.example.com/callback", + client_id="test-client-id", + client_secret=None, + code_verifier="test_verifier", + ) + + response_data = json.loads(response.body) + assert response_data["access_token"] == "test_access_token_123" + assert response_data["refresh_token"] == "test_refresh_token_456" + + mock_prisma.db.litellm_mcpservertable.update.assert_called_once() + update_call = mock_prisma.db.litellm_mcpservertable.update.call_args + + assert update_call[1]["where"]["server_id"] == "test-persist-server" + + update_data = update_call[1]["data"] + assert update_data["auth_type"] == "oauth2" + assert update_data["token_url"] == "https://auth.example.com/token" + assert update_data["authorization_url"] == "https://auth.example.com/authorize" + assert update_data["registration_url"] == "https://auth.example.com/register" + assert update_data["updated_by"] == "oauth_token_exchange" + + credentials_json = json.loads(update_data["credentials"]) + assert "auth_value" in credentials_json + assert "client_id" in credentials_json + assert "refresh_token" in credentials_json + assert "type" in credentials_json + + # Verify expires_at is ISO-8601 string (not Unix epoch int) + assert "expires_at" in credentials_json + from datetime import datetime + + datetime.fromisoformat(credentials_json["expires_at"]) # should not raise + + # Verify sensitive fields are encrypted (not plaintext) + assert credentials_json["auth_value"] != "test_access_token_123" + assert credentials_json["client_id"] != "test-client-id" + assert credentials_json["refresh_token"] != "test_refresh_token_456" + + +@pytest.mark.asyncio +async def test_token_exchange_persists_without_optional_fields(): + """Test that token exchange persists correctly when refresh_token and expires_in are absent.""" + try: + import json + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + 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") + + oauth2_server = MCPServer( + server_id="test-persist-minimal", + name="test_minimal", + server_name="test_minimal", + alias="test_minimal", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.example.com/" + mock_request.headers = {} + + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "minimal_token", + "token_type": "Bearer", + } + mock_response.raise_for_status = MagicMock() + + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client, patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), patch( + "litellm.proxy.common_utils.encrypt_decrypt_utils._get_salt_key", + return_value="test-salt-key", + ): + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_async_client + + response = await exchange_token_with_server( + request=mock_request, + mcp_server=oauth2_server, + grant_type="authorization_code", + code="test_code", + redirect_uri="https://proxy.example.com/callback", + client_id="test-client", + client_secret=None, + code_verifier=None, + ) + + response_data = json.loads(response.body) + assert response_data["access_token"] == "minimal_token" + assert "refresh_token" not in response_data + + mock_prisma.db.litellm_mcpservertable.update.assert_called_once() + update_data = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + assert update_data["auth_type"] == "oauth2" + assert update_data["token_url"] == "https://auth.example.com/token" + assert "authorization_url" not in update_data + assert "registration_url" not in update_data + + +@pytest.mark.asyncio +async def test_token_refresh_preserves_existing_refresh_token(): + """ + When a non-rotating OAuth server returns a new access_token without a + refresh_token, the previously stored refresh_token must be preserved. + """ + try: + import json + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + 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") + + oauth2_server = MCPServer( + server_id="test-refresh-preserve", + name="test_refresh", + server_name="test_refresh", + alias="test_refresh", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + token_url="https://auth.example.com/token", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://proxy.example.com/" + mock_request.headers = {} + + # Upstream returns new access_token but NO refresh_token (non-rotating) + mock_response = MagicMock() + mock_response.json.return_value = { + "access_token": "new_access_token", + "token_type": "Bearer", + "expires_in": 3600, + } + mock_response.raise_for_status = MagicMock() + + # Existing DB record has the original refresh_token + mock_existing = MagicMock() + mock_existing.credentials = json.dumps({ + "auth_value": "old_encrypted_access", + "refresh_token": "original_refresh_token_encrypted", + "client_id": "existing_client_id", + "type": "oauth2", + }) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=mock_existing + ) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" + ) as mock_get_client, patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma, + ), patch( + "litellm.proxy.common_utils.encrypt_decrypt_utils._get_salt_key", + return_value="test-salt-key", + ): + mock_async_client = MagicMock() + mock_async_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_async_client + + # Simulate refresh_token grant — caller passes the stored refresh token + response = await exchange_token_with_server( + request=mock_request, + mcp_server=oauth2_server, + grant_type="refresh_token", + code=None, + redirect_uri=None, + client_id="test-client", + client_secret=None, + code_verifier=None, + refresh_token="original_refresh_token", + ) + + response_data = json.loads(response.body) + assert response_data["access_token"] == "new_access_token" + + mock_prisma.db.litellm_mcpservertable.update.assert_called_once() + update_data = mock_prisma.db.litellm_mcpservertable.update.call_args[1]["data"] + credentials_json = json.loads(update_data["credentials"]) + + # The refresh_token should be preserved (merged from existing + fallback) + assert "refresh_token" in credentials_json + # auth_value should be the new access token (encrypted, so not plaintext) + assert credentials_json["auth_value"] != "new_access_token" diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index f97d8ffab04..7028879dd86 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -99,11 +99,26 @@ const UserDashboard: React.FC = ({ const [selectedTeam, setSelectedTeam] = useState(null); // Clear session storage on page unload so next load fetches fresh data. - // Note: MCP auth tokens are persistent and should not be cleared on page refresh - // They are only cleared on logout + // Preserve MCP OAuth flow state keys across the clear — they are needed + // to complete the OAuth PKCE redirect flow (code_verifier, state, etc.). + // Uses prefix matching so future per-server or per-flow keys are + // automatically covered without updating a hardcoded list. useEffect(() => { const handleBeforeUnload = () => { + const saved: Record = {}; + const keysToPreserve = Object.keys(sessionStorage).filter( + (key) => + key.startsWith("litellm-mcp-oauth") || + key.startsWith("litellm-user-mcp-oauth") + ); + keysToPreserve.forEach((key) => { + const val = sessionStorage.getItem(key); + if (val) saved[key] = val; + }); sessionStorage.clear(); + Object.entries(saved).forEach(([key, val]) => { + sessionStorage.setItem(key, val); + }); }; window.addEventListener("beforeunload", handleBeforeUnload); return () => window.removeEventListener("beforeunload", handleBeforeUnload);