mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix: MCP OAuth PKCE — persist credentials, encrypt refresh_token, preserve sessionStorage (#24938)
* fix: MCP OAuth PKCE flow — persist credentials and preserve sessionStorage Three bugs prevented the MCP OAuth PKCE flow from working end-to-end: 1. sessionStorage.clear() in beforeunload wipes OAuth flow state during redirect to the OAuth provider (#19821). Fix: preserve keys matching litellm-mcp-oauth* and litellm-user-mcp-oauth* prefixes. 2. exchange_token_with_server() never persisted credentials to the DB (#20493). Fix: add _persist_oauth_credentials() that encrypts and saves access_token, refresh_token, client_id, and discovery URLs. Merges with existing credentials to avoid wiping unrelated fields. 3. refresh_token was not encrypted by encrypt_credentials(). Fix: add refresh_token to both encrypt_credentials() and decrypt_credentials(). Additional fixes from review: - expires_at stored as ISO-8601 string (matching store_user_oauth_credential) - Refresh token preserved on non-rotating OAuth servers (fallback to incoming refresh_token when upstream omits it) - Credentials merged with existing DB record instead of overwriting - sessionStorage keys matched by prefix instead of hardcoded list Fixes #19821, #20493 * fix: resolve real DB server_id in _persist_oauth_credentials When the OAuth PKCE session uses an ephemeral UUID as server_id, the Prisma update was silently missing the real DB row. Now falls back to find_first by server_name/alias to locate the persistent record before updating. Logs both db_id and session_id for traceability.
This commit is contained in:
madhu19991
GitHub
parent
51c9553920
commit
e82f59472d
4 changed files with 445 additions and 2 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -99,11 +99,26 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
|
|||
const [selectedTeam, setSelectedTeam] = useState<any | null>(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<string, string> = {};
|
||||
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);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue