From 05f39bf9427290e077a6ef07c6d964cc90ef44df Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 10:43:13 -0700 Subject: [PATCH 01/12] fix(mcp): invalidate a browser-authorized upstream token when a mint-relevant field changes An admin who ran Authorize & Fetch and then changed a field that determines which upstream OAuth token gets minted kept using the stale token for tool preview, sessionStorage, and (on the backend) the stored per-user credential and its cache. Grounded in RFC 8707/8693 and the MCP auth spec, a token is bound to one tuple: resource/audience (url), OAuth mode/grant (auth_type, oauth_flow_type), the authorization-server endpoints, and the OAuth client + scopes. A shared getOAuthAuthorizationIdentity captures exactly those fields; transport (http/sse on the same url is the same audience) and delegate_auth_to_upstream (a downstream-usage toggle never sent to the authorize request) are excluded. UI: both the create and edit forms now discard the held token (React state / sessionStorage / hook, plus the fetched token + DCR client in form.credentials) whenever the identity diverges from the one it was authorized against, re-applying the admin's in-flight edit so it is never wiped. The check lives in one shared helper so the two forms cannot drift. Backend: editing an MCP server now compares the pre/post identity and, on a mint-relevant change, purges every stored per-user OAuth credential for the server (DB row + per-user token cache) so no user forwards a token minted for a resource/AS/client that no longer matches. Best-effort; a purge failure never fails the update. --- litellm/proxy/_experimental/mcp_server/db.py | 44 +++++++++ .../mcp_management_endpoints.py | 29 ++++++ .../mcp_server/test_db_credentials.py | 89 +++++++++++++++++++ .../mcp_tools/create_mcp_server.test.tsx | 38 ++++++++ .../mcp_tools/create_mcp_server.tsx | 67 +++++++------- .../components/mcp_tools/mcp_server_edit.tsx | 53 ++++++++++- .../src/components/mcp_tools/types.tsx | 27 ++++++ 7 files changed, 312 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 10081ce19de..8c5e728d86b 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1070,6 +1070,50 @@ async def list_user_oauth_credentials( return results +def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]: + """The upstream-OAuth-token-determining fields of an MCP server: the resource/audience (url), the + OAuth mode/grant (auth_type, oauth2_flow), the authorization-server endpoints, and the OAuth client + + scopes. Mirrors the dashboard's getOAuthAuthorizationIdentity. When any of these change on a server + update, previously stored per-user tokens were minted for the old identity and are stale. Excludes + transport and delegate_auth_to_upstream, which do not affect what token is minted (RFC 8707/8693).""" + creds = getattr(server, "credentials", None) + creds_dict: Dict[str, Any] = creds if isinstance(creds, dict) else {} + return ( + getattr(server, "url", None), + getattr(server, "auth_type", None), + getattr(server, "oauth2_flow", None), + getattr(server, "authorization_url", None), + getattr(server, "token_url", None), + getattr(server, "registration_url", None), + creds_dict.get("client_id"), + creds_dict.get("client_secret"), + creds_dict.get("scopes"), + ) + + +async def purge_user_oauth_credentials_for_server(prisma_client: PrismaClient, server_id: str) -> int: + """Delete every stored per-user OAuth credential for a server and drop each from the per-user token + cache, so no user keeps a token minted for a superseded configuration. Called when a server update + changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows removed.""" + repo = MCPUserCredentialsRepository(prisma_client) + rows = await repo.table.find_many(where={"server_id": server_id}) + if not rows: + return 0 + await repo.table.delete_many(where={"server_id": server_id}) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + mcp_per_user_token_cache, + ) + + for row in rows: + try: + await mcp_per_user_token_cache.delete(row.user_id, server_id) + except Exception as exc: # noqa: BLE001 - cache drop is best-effort; the DB delete is authoritative + verbose_proxy_logger.warning( + "Failed to drop cached MCP OAuth token for user=%s server=%s: %s", row.user_id, server_id, exc + ) + return len(rows) + + async def refresh_user_oauth_token( prisma_client: PrismaClient, user_id: str, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c9952b245c7..8f6779b17b9 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -125,7 +125,9 @@ if MCP_AVAILABLE: get_user_env_vars_bulk, get_user_oauth_credential, list_user_oauth_credentials, + mcp_oauth_token_identity, merge_user_env_vars, + purge_user_oauth_credentials_for_server, reject_mcp_server, store_user_credential, store_user_oauth_credential, @@ -2318,6 +2320,9 @@ if MCP_AVAILABLE: }, ) + # Snapshot the pre-update identity so we can detect a mint-relevant change below. + old_server_record = await get_mcp_server(prisma_client, payload.server_id) + # try to update the mcp server mcp_server_record_updated = await update_mcp_server( prisma_client, @@ -2336,6 +2341,30 @@ if MCP_AVAILABLE: # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() + # If a field that determines which upstream OAuth token gets minted changed (url/audience, OAuth + # mode/grant, authorization-server endpoints, or the OAuth client + scopes), every stored per-user + # token was minted for the old configuration and is stale. Purge them (DB + cache) so the next + # tool call re-authorizes instead of forwarding a token for a resource/AS/client that no longer + # matches. Best-effort: a purge failure must not fail the update, whose primary job already + # succeeded. + if old_server_record is not None and mcp_oauth_token_identity(old_server_record) != mcp_oauth_token_identity( + mcp_server_record_updated + ): + try: + purged = await purge_user_oauth_credentials_for_server(prisma_client, payload.server_id) + if purged: + verbose_logger.info( + "MCP server %s: purged %d stale per-user OAuth token(s) after a mint-relevant config change", + payload.server_id, + purged, + ) + except Exception as exc: # noqa: BLE001 - purge is best-effort; the server update already succeeded + verbose_logger.warning( + "MCP server %s: failed to purge stale per-user OAuth tokens after config change: %s", + payload.server_id, + exc, + ) + # TODO: Enterprise: Finish audit log trail if litellm.store_audit_logs: pass diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 8b8b8a363d5..628f422fbf1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -11,6 +11,7 @@ keeps a plain-base64 fallback on read so existing rows continue to work. import base64 import json from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest @@ -63,6 +64,94 @@ def _legacy_row(payload: str): return row +def _identity_server(**overrides): + base = dict( + url="https://up.example.com/mcp", + auth_type="oauth2", + oauth2_flow="authorization_code", + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + registration_url="https://idp.example.com/register", + credentials={"client_id": "cid", "client_secret": "csec", "scopes": ["a"]}, + server_name="srv", + description="d", + ) + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.mark.parametrize( + "overrides", + [ + {"url": "https://other.example.com/mcp"}, + {"auth_type": "oauth_delegate"}, + {"oauth2_flow": "client_credentials"}, + {"authorization_url": "https://other.example.com/authorize"}, + {"token_url": "https://other.example.com/token"}, + {"registration_url": "https://other.example.com/register"}, + {"credentials": {"client_id": "new", "client_secret": "csec", "scopes": ["a"]}}, + {"credentials": {"client_id": "cid", "client_secret": "rotated", "scopes": ["a"]}}, + {"credentials": {"client_id": "cid", "client_secret": "csec", "scopes": ["b"]}}, + ], +) +def test_mcp_oauth_token_identity_changes_on_mint_relevant_fields(overrides): + from litellm.proxy._experimental.mcp_server.db import mcp_oauth_token_identity + + assert mcp_oauth_token_identity(_identity_server()) != mcp_oauth_token_identity(_identity_server(**overrides)) + + +@pytest.mark.parametrize( + "overrides", + [ + {"server_name": "renamed"}, + {"description": "changed"}, + ], +) +def test_mcp_oauth_token_identity_stable_on_non_mint_fields(overrides): + from litellm.proxy._experimental.mcp_server.db import mcp_oauth_token_identity + + assert mcp_oauth_token_identity(_identity_server()) == mcp_oauth_token_identity(_identity_server(**overrides)) + + +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_deletes_rows_and_cache(monkeypatch): + from litellm.proxy._experimental.mcp_server import oauth2_token_cache + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + r1 = MagicMock(user_id="alice", server_id="srv-1") + r2 = MagicMock(user_id="bob", server_id="srv-1") + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[r1, r2]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() + + cache_deletes = [] + monkeypatch.setattr( + oauth2_token_cache.mcp_per_user_token_cache, + "delete", + AsyncMock(side_effect=lambda uid, sid: cache_deletes.append((uid, sid))), + ) + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + + assert purged == 2 + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + assert set(cache_deletes) == {("alice", "srv-1"), ("bob", "srv-1")} + + +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_noop_when_empty(): + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + + assert purged == 0 + prisma.db.litellm_mcpusercredentials.delete_many.assert_not_awaited() + + def _stored_value(prisma) -> str: """Pull the credential_b64 value passed to the most recent upsert call.""" call = prisma.db.litellm_mcpusercredentials.upsert.call_args diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index eecbc253b6b..7339420ed68 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -681,6 +681,44 @@ describe("CreateMCPServer", () => { // Asserted in setupOAuthInteractive }); + it("invalidates the held token when the auth mode changes after Authorize & Fetch", async () => { + await setupOAuthInteractive(); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://a.example.com/mcp" } }); + }); + act(() => { + oauthHook.onTokenReceived?.({ access_token: "tok-a" }, { clientId: "client-a", clientSecret: "secret-a" }); + }); + oauthHook.reset.mockClear(); + + // Switching the Authentication mode changes the OAuth identity, so the held token is discarded. + await selectAntOption("Authentication", "True Passthrough (no LiteLLM auth)"); + + await waitFor(() => expect(oauthHook.reset).toHaveBeenCalled()); + }); + + it("does NOT invalidate the held token when a non-mint field (server name) changes", async () => { + await setupOAuthInteractive(); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://a.example.com/mcp" } }); + }); + act(() => { + oauthHook.onTokenReceived?.({ access_token: "tok-a" }, { clientId: "client-a", clientSecret: "secret-a" }); + }); + oauthHook.reset.mockClear(); + + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "Renamed_Server" } }); + }); + + // server_name is not part of the OAuth identity, so the held token must survive the edit. + await waitFor(() => expect(screen.getAllByRole("button", { name: "Add MCP Server" }).length).toBeGreaterThan(0)); + expect(oauthHook.reset).not.toHaveBeenCalled(); + }); + it("includes token_validation in payload when token_validation_json is filled with valid JSON", async () => { vi.mocked(networking.createMCPServer).mockResolvedValue({ server_id: "new-server-oauth", diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 0b39add234c..17c7e7c0b9a 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -15,6 +15,7 @@ import { MCP_OAUTH2_FLOW_M2M, MCP_OAUTH2_FLOW_INTERACTIVE, isClientForwardedTokenMode, + getOAuthAuthorizationIdentity, } from "./types"; import OAuthFormFields from "./OAuthFormFields"; import TruePassthroughWarning from "./TruePassthroughWarning"; @@ -99,7 +100,10 @@ const CreateMCPServer: React.FC = ({ const [oauthAccessToken, setOauthAccessToken] = useState(null); const [logoUrl, setLogoUrl] = useState(undefined); const [oauthDocsUrl, setOauthDocsUrl] = useState(null); - const [authorizedUrl, setAuthorizedUrl] = useState(undefined); + // The OAuth authorization identity (see getOAuthAuthorizationIdentity) captured at the moment a token + // was fetched; undefined when no valid token is held. If any mint-relevant field diverges from this, + // the held token is stale and is discarded so the admin must re-authorize. + const [authorizedIdentity, setAuthorizedIdentity] = useState(undefined); // Single hook call shared by MCPConnectionStatus and MCPToolConfiguration to avoid duplicate requests. const { @@ -125,12 +129,6 @@ const CreateMCPServer: React.FC = ({ const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4; const isM2MFlow = isOAuthAuthType && formValues.oauth_flow_type === OAUTH_FLOW.M2M; - const getOAuthAuthorizationTarget = (values: Record): string | undefined => { - const transport = values.transport || transportType; - const target = transport === TRANSPORT.OPENAPI ? values.spec_path : values.url; - return typeof target === "string" ? target : undefined; - }; - const persistCreateUiState = () => { if (typeof window === "undefined") { return; @@ -207,6 +205,7 @@ const CreateMCPServer: React.FC = ({ // and committed to sessionStorage on submit; it must never be written into form.credentials, // which would persist it as server-level credentials on the created server row. Mirrors the // edit form's onTokenReceived early return. + setAuthorizedIdentity(getOAuthAuthorizationIdentity(form.getFieldsValue(true))); NotificationsManager.success( "Token held for this browser session. Tools can now be previewed and configured; nothing will be saved to LiteLLM.", ); @@ -223,7 +222,9 @@ const CreateMCPServer: React.FC = ({ }; form.setFieldsValue({ credentials }); - setAuthorizedUrl(getOAuthAuthorizationTarget(form.getFieldsValue(true))); + // Capture the identity AFTER writing the DCR'd credentials so the held token is not spuriously + // invalidated by its own credential write. + setAuthorizedIdentity(getOAuthAuthorizationIdentity(form.getFieldsValue(true))); NotificationsManager.success( "OAuth authorization successful! Please click 'Create MCP Server' to save the configuration.", @@ -233,13 +234,24 @@ const CreateMCPServer: React.FC = ({ flowSource: "create", }); - const clearAuthorizedOAuthState = (values: Record) => { - form.resetFields(["credentials", "authorization_url", "token_url", "registration_url"]); - form.setFieldsValue(values); + // Discard the held browser-authorized token and its tool preview when the authorization identity + // changes (or the modal closes). For oauth2 the fetched token + DCR client also live in + // form.credentials, and the discovered endpoints in authorization_url/token_url/registration_url, so + // those form fields are reset too; whatever the admin just changed (passed via changedValues) is + // re-applied so the invalidation never wipes their in-flight edit. + const CLEARED_ON_INVALIDATION = ["credentials", "authorization_url", "token_url", "registration_url"] as const; + const clearHeldOAuthToken = (changedValues: Record = {}) => { setOauthAccessToken(null); clearTools(); resetOAuthFlow(); - setAuthorizedUrl(undefined); + setAuthorizedIdentity(undefined); + form.resetFields([...CLEARED_ON_INVALIDATION]); + const preserved = Object.fromEntries( + CLEARED_ON_INVALIDATION.filter((key) => key in changedValues).map((key) => [key, changedValues[key]]), + ); + if (Object.keys(preserved).length > 0) { + form.setFieldsValue(preserved); + } }; React.useEffect(() => { @@ -577,7 +589,7 @@ const CreateMCPServer: React.FC = ({ : { spec_path: undefined, command: undefined, args: undefined, env: undefined }; const nextValues = - authorizedUrl === undefined + authorizedIdentity === undefined ? transportValues : { ...transportValues, @@ -587,10 +599,9 @@ const CreateMCPServer: React.FC = ({ registration_url: undefined, }; - if (authorizedUrl !== undefined) { - clearAuthorizedOAuthState(nextValues); - } else { - form.setFieldsValue(nextValues); + form.setFieldsValue(nextValues); + if (authorizedIdentity !== undefined) { + clearHeldOAuthToken(); } }; @@ -652,28 +663,18 @@ const CreateMCPServer: React.FC = ({ setOauthAccessToken(null); clearTools(); resetOAuthFlow(); - setAuthorizedUrl(undefined); + setAuthorizedIdentity(undefined); } }, [isModalVisible, form, clearTools, resetOAuthFlow]); const isAdmin = isAdminRole(userRole); const handleFormValuesChange = (changedValues: Record, allValues: Record) => { - const changedAuthorizationTarget = "url" in changedValues || "spec_path" in changedValues; - if ( - changedAuthorizationTarget && - authorizedUrl !== undefined && - getOAuthAuthorizationTarget(allValues) !== authorizedUrl - ) { - const invalidated = { - credentials: undefined, - authorization_url: changedValues.authorization_url, - token_url: changedValues.token_url, - registration_url: changedValues.registration_url, - }; - clearAuthorizedOAuthState(invalidated); - setFormValues({ ...allValues, ...invalidated }); - return; + // Any change to a mint-relevant field (url, auth_type, oauth_flow_type, client creds/scopes, or the + // authorization/token/registration endpoints — see getOAuthAuthorizationIdentity) makes a held token + // stale, so discard it and force a fresh authorize. + if (authorizedIdentity !== undefined && getOAuthAuthorizationIdentity(allValues) !== authorizedIdentity) { + clearHeldOAuthToken(changedValues); } setFormValues(allValues); }; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index ee7938d2904..55adcf2bb59 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -5,6 +5,7 @@ import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/rea import { AUTH_TYPE, isClientForwardedTokenMode, + getOAuthAuthorizationIdentity, OAUTH_FLOW, MCP_OAUTH2_FLOW_M2M, MCP_OAUTH2_FLOW_INTERACTIVE, @@ -15,7 +16,7 @@ import { oauth2FlowToFormValue, } from "./types"; import { updateMCPServer, listMCPTools, storeMCPOAuthUserCredential } from "../networking"; -import { getToken, isTokenValid, setToken } from "@/utils/mcpTokenStore"; +import { getToken, isTokenValid, removeToken, setToken } from "@/utils/mcpTokenStore"; import { buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils"; import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPPermissionManagement from "./MCPPermissionManagement"; @@ -136,11 +137,17 @@ const MCPServerEdit: React.FC = ({ // that read only mcpServer.auth_type go stale the moment the admin switches modes in the form. const getEffectiveAuthType = () => form.getFieldValue("auth_type") ?? mcpServer.auth_type; + // The OAuth authorization identity (see getOAuthAuthorizationIdentity) captured when a token is fetched + // in this edit session; undefined when none is held. If a mint-relevant field later diverges from it, + // the held token (hook response + sessionStorage) is discarded so the admin must re-authorize. + const authorizedIdentityRef = React.useRef(undefined); + const { startOAuthFlow, status: oauthStatus, error: oauthError, tokenResponse: oauthTokenResponse, + reset: resetOAuthFlow, } = useMcpOAuthFlow({ accessToken, getCredentials: () => form.getFieldValue("credentials"), @@ -183,6 +190,7 @@ const MCPServerEdit: React.FC = ({ return; } + authorizedIdentityRef.current = getOAuthAuthorizationIdentity(form.getFieldsValue(true)); if (isClientForwardedTokenMode(getEffectiveAuthType())) { const browserHeldToken = { access_token: token.access_token, @@ -205,6 +213,8 @@ const MCPServerEdit: React.FC = ({ }; form.setFieldsValue({ credentials }); + // Re-capture after writing credentials so the token is not invalidated by its own credential write. + authorizedIdentityRef.current = getOAuthAuthorizationIdentity(form.getFieldsValue(true)); NotificationsManager.success( "OAuth authorization successful! Please click 'Update MCP Server' to save the credentials.", @@ -378,6 +388,39 @@ const MCPServerEdit: React.FC = ({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token]); + // Invalidate a token authorized in this edit session once any mint-relevant field diverges from the + // identity it was minted against (url, auth_type, oauth_flow_type, client creds/scopes, or the + // authorization/token/registration endpoints — see getOAuthAuthorizationIdentity). Discards the hook + // token (resetOAuthFlow, which re-runs fetchTools to prompt a fresh authorize), the sessionStorage + // token (removeToken, browser-held modes), and the fetched token/DCR client in form.credentials + the + // discovered endpoint fields; the admin's in-flight edit is re-applied so it is never wiped. Only fires + // when a token was actually authorized here (ref set), so a token already valid for the saved server on + // mount is left untouched. Driven from onValuesChange (user input only), never programmatic resets. + const CLEARED_ON_INVALIDATION = ["credentials", "authorization_url", "token_url", "registration_url"] as const; + const clearHeldOAuthToken = (changedValues: Record = {}) => { + authorizedIdentityRef.current = undefined; + if (mcpServer.server_id) { + removeToken(mcpServer.server_id, userID); + } + resetOAuthFlow(); + form.resetFields([...CLEARED_ON_INVALIDATION]); + const preserved = Object.fromEntries( + CLEARED_ON_INVALIDATION.filter((key) => key in changedValues).map((key) => [key, changedValues[key]]), + ); + if (Object.keys(preserved).length > 0) { + form.setFieldsValue(preserved); + } + }; + + const handleFormValuesChange = (changedValues: Record) => { + if ( + authorizedIdentityRef.current !== undefined && + getOAuthAuthorizationIdentity(form.getFieldsValue(true)) !== authorizedIdentityRef.current + ) { + clearHeldOAuthToken(changedValues); + } + }; + const fetchTools = async () => { if (!accessToken || !mcpServer.server_id) return; @@ -805,7 +848,13 @@ const MCPServerEdit: React.FC = ({ -
+ sse on the same url is the same audience; a transport switch +// only matters when it changes the target url, which `target` already captures), delegate_auth_to_upstream +// (a downstream-usage toggle that is never sent to the authorize request), and all metadata/RBAC/routing +// fields. Shared by the create and edit forms so their invalidation logic cannot drift. +export const getOAuthAuthorizationIdentity = (values: Record): string => { + const credentials = (values.credentials ?? {}) as Record; + const target = values.transport === TRANSPORT.OPENAPI ? values.spec_path : values.url; + const identity = { + target: typeof target === "string" ? target : null, + auth_type: values.auth_type ?? null, + oauth_flow_type: values.oauth_flow_type ?? null, + client_id: credentials.client_id ?? null, + client_secret: credentials.client_secret ?? null, + scopes: credentials.scopes ?? null, + authorization_url: values.authorization_url ?? null, + token_url: values.token_url ?? null, + registration_url: values.registration_url ?? null, + }; + return JSON.stringify(identity); +}; + // Backend value of `oauth2_flow` that marks a machine-to-machine server. Distinct // from the UI-local OAUTH_FLOW.M2M ("m2m"); this is what the API actually returns. export const MCP_OAUTH2_FLOW_M2M = "client_credentials"; From 48124734a08d40ec1b17c3de86c1af0cafc6fa97 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 12:41:15 -0700 Subject: [PATCH 02/12] fix(mcp): compare the token identity decrypted and invalidate every per-user token store Review follow-ups on the stale-token invalidation. The backend identity now decrypts client_id and client_secret before comparing: the stored values are NaCl-encrypted with a fresh nonce on every write, so comparing ciphertext flagged every routine save as a mint-relevant change and purged per-user tokens that were still valid. The identity also gains spec_path, the audience for OpenAPI servers, and parses credentials stored as a JSON string The purge now routes each (user, server) through the manager's invalidate_user_oauth_token_cache, which becomes the single invalidation point covering both the legacy per-user token cache and the v2 per-user OAuth token store; previously the purge evicted only the legacy cache while the revoke path evicted only the v2 store, so each path left the other cache serving a replaced token until its TTL. A credential row racing in between the find and the delete is now detected via the delete_many count and logged; its cache entry expires by TTL On the dashboard, CLEARED_ON_INVALIDATION and the staleness check move to types.tsx as the single shared implementation for both forms. The edit form's transport handler now rechecks the identity after its programmatic setFieldsValue calls, which antd does not report through onValuesChange, so a token no longer survives a transport switch that clears the mint target. The create form rebuilds formValues from the post-reset form state after an invalidation instead of publishing the pre-reset snapshot, so the tool preview can no longer refetch with the discarded DCR client. Both transport handlers now share the recheck, which also stops the create form from over-invalidating on an http to sse swap that keeps the same url and therefore the same audience --- litellm/proxy/_experimental/mcp_server/db.py | 76 ++++++-- .../mcp_server/mcp_server_manager.py | 17 +- .../mcp_server/test_db_credentials.py | 182 +++++++++--------- .../mcp_server/test_mcp_server_manager.py | 37 +++- .../mcp_tools/create_mcp_server.test.tsx | 50 +++++ .../mcp_tools/create_mcp_server.tsx | 32 ++- .../mcp_tools/mcp_server_edit.test.tsx | 81 +++++++- .../components/mcp_tools/mcp_server_edit.tsx | 21 +- .../src/components/mcp_tools/types.tsx | 15 ++ 9 files changed, 367 insertions(+), 144 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 8c5e728d86b..135a9e055d6 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1070,48 +1070,82 @@ async def list_user_oauth_credentials( return results +def _decrypted_credential_field(creds: Dict[str, Any], field: str) -> Any: + """Return one credential field decrypted with the global salt key; non-string and legacy + plaintext values come back unchanged (decrypt_value_helper returns the original on failure).""" + value = creds.get(field) + if not isinstance(value, str): + return value + return decrypt_value_helper( + value=value, + key=field, + exception_type="debug", + return_original_value=True, + ) + + def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]: - """The upstream-OAuth-token-determining fields of an MCP server: the resource/audience (url), the - OAuth mode/grant (auth_type, oauth2_flow), the authorization-server endpoints, and the OAuth client + - scopes. Mirrors the dashboard's getOAuthAuthorizationIdentity. When any of these change on a server - update, previously stored per-user tokens were minted for the old identity and are stale. Excludes - transport and delegate_auth_to_upstream, which do not affect what token is minted (RFC 8707/8693).""" + """The upstream-OAuth-token-determining fields of an MCP server: the resource/audience (url, or + spec_path for OpenAPI servers), the OAuth mode/grant (auth_type, oauth2_flow), the + authorization-server endpoints, and the OAuth client + scopes. Mirrors the dashboard's + getOAuthAuthorizationIdentity. When any of these change on a server update, previously stored + per-user tokens were minted for the old identity and are stale. Excludes transport and + delegate_auth_to_upstream, which do not affect what token is minted (RFC 8707/8693). + + client_id/client_secret are compared decrypted: stored values are NaCl-encrypted with a fresh + nonce on every write, so comparing ciphertext would flag every routine save as an identity + change and purge tokens that are still valid.""" creds = getattr(server, "credentials", None) - creds_dict: Dict[str, Any] = creds if isinstance(creds, dict) else {} + if isinstance(creds, str): + try: + parsed: Any = json.loads(creds) + except ValueError: + parsed = None + else: + parsed = creds + creds_dict: Dict[str, Any] = parsed if isinstance(parsed, dict) else {} return ( getattr(server, "url", None), + getattr(server, "spec_path", None), getattr(server, "auth_type", None), getattr(server, "oauth2_flow", None), getattr(server, "authorization_url", None), getattr(server, "token_url", None), getattr(server, "registration_url", None), - creds_dict.get("client_id"), - creds_dict.get("client_secret"), + _decrypted_credential_field(creds_dict, "client_id"), + _decrypted_credential_field(creds_dict, "client_secret"), creds_dict.get("scopes"), ) async def purge_user_oauth_credentials_for_server(prisma_client: PrismaClient, server_id: str) -> int: - """Delete every stored per-user OAuth credential for a server and drop each from the per-user token - cache, so no user keeps a token minted for a superseded configuration. Called when a server update - changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows removed.""" + """Delete every stored per-user OAuth credential for a server and invalidate each user's cached + token everywhere it can be served from (the legacy per-user token cache and the v2 per-user OAuth + token store), so no user keeps a token minted for a superseded configuration. Called when a server + update changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows + removed. A row inserted between the find and the delete is removed from the DB but cannot be + evicted from the caches (its user_id was never seen); that case is detected, logged, and bounded + by the cache TTL.""" repo = MCPUserCredentialsRepository(prisma_client) rows = await repo.table.find_many(where={"server_id": server_id}) if not rows: return 0 - await repo.table.delete_many(where={"server_id": server_id}) - from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( - mcp_per_user_token_cache, + deleted_count = await repo.table.delete_many(where={"server_id": server_id}) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, ) for row in rows: - try: - await mcp_per_user_token_cache.delete(row.user_id, server_id) - except Exception as exc: # noqa: BLE001 - cache drop is best-effort; the DB delete is authoritative - verbose_proxy_logger.warning( - "Failed to drop cached MCP OAuth token for user=%s server=%s: %s", row.user_id, server_id, exc - ) - return len(rows) + await global_mcp_server_manager.invalidate_user_oauth_token_cache(row.user_id, server_id) + if deleted_count != len(rows): + verbose_proxy_logger.warning( + "MCP server %s: purge removed %d credential row(s) but %d were enumerated; " + "row(s) raced in during the purge and their cached tokens will expire by TTL", + server_id, + deleted_count, + len(rows), + ) + return deleted_count async def refresh_user_oauth_token( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 356ed7a2729..cc4a9e63105 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -57,7 +57,10 @@ from litellm.proxy._experimental.mcp_server.elicitation_handler import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) -from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth +from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + mcp_per_user_token_cache, + resolve_mcp_auth, +) from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, Ok, @@ -4053,10 +4056,13 @@ class MCPServerManager: return await self._cred_provider.has_user_token(to_subject(user_api_key_auth, None), spec) async def invalidate_user_oauth_token_cache(self, user_id: str, server_id: str) -> None: - """Drop the v2 chain's cached token for ``(user_id, server_id)`` after the credential row - changes (re-auth, revoke), so the next resolve reads the new row instead of serving the - replaced token until its cache TTL. Best-effort: a cache-drop failure is logged, never - raised, because the DB write already succeeded and the TTL remains the backstop. + """Drop every cached token for ``(user_id, server_id)`` after the credential row changes + (re-auth, revoke, config-change purge): the v2 chain's cache and the legacy per-user token + cache, so the next resolve reads the new row instead of serving the replaced token until its + cache TTL, whichever path resolves it. This is the single invalidation point for per-user + OAuth tokens; callers must not evict individual caches directly. Best-effort: a cache-drop + failure is logged, never raised, because the DB write already succeeded and the TTL remains + the backstop. """ try: await self._per_user_oauth_token_store.invalidate(user_id, server_id) @@ -4064,6 +4070,7 @@ class MCPServerManager: verbose_logger.warning( "Failed to invalidate cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc ) + await mcp_per_user_token_cache.delete(user_id, server_id) async def _resolve_oauth2_headers_for_tool_call( self, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 628f422fbf1..51641991ef9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -84,6 +84,7 @@ def _identity_server(**overrides): "overrides", [ {"url": "https://other.example.com/mcp"}, + {"spec_path": "https://up.example.com/openapi.json"}, {"auth_type": "oauth_delegate"}, {"oauth2_flow": "client_credentials"}, {"authorization_url": "https://other.example.com/authorize"}, @@ -113,29 +114,91 @@ def test_mcp_oauth_token_identity_stable_on_non_mint_fields(overrides): assert mcp_oauth_token_identity(_identity_server()) == mcp_oauth_token_identity(_identity_server(**overrides)) +def _encrypted_creds_json(client_id: str = "cid", client_secret: str = "csec") -> str: + from litellm.proxy._experimental.mcp_server.db import encrypt_credentials + + encrypted = encrypt_credentials( + credentials={"client_id": client_id, "client_secret": client_secret, "scopes": ["a"]}, + encryption_key=None, + ) + return json.dumps(encrypted) + + +def test_mcp_oauth_token_identity_stable_across_reencryption(): + """Stored client_id/client_secret are NaCl-encrypted with a fresh nonce on every write, so two + saves of the SAME plaintext produce different ciphertext. The identity must compare decrypted + values; comparing ciphertext would flag every routine save as a mint-relevant change and purge + per-user tokens that are still valid.""" + from litellm.proxy._experimental.mcp_server.db import mcp_oauth_token_identity + + first = _encrypted_creds_json() + second = _encrypted_creds_json() + assert first != second + + assert mcp_oauth_token_identity(_identity_server(credentials=first)) == mcp_oauth_token_identity( + _identity_server(credentials=second) + ) + + +def test_mcp_oauth_token_identity_detects_change_under_encryption(): + from litellm.proxy._experimental.mcp_server.db import mcp_oauth_token_identity + + unchanged = _identity_server(credentials=_encrypted_creds_json()) + changed = _identity_server(credentials=_encrypted_creds_json(client_id="other")) + assert mcp_oauth_token_identity(unchanged) != mcp_oauth_token_identity(changed) + + @pytest.mark.asyncio -async def test_purge_user_oauth_credentials_for_server_deletes_rows_and_cache(monkeypatch): - from litellm.proxy._experimental.mcp_server import oauth2_token_cache +async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(monkeypatch): + """The purge must route each (user, server) through the manager's shared invalidation, which is + the single point covering both the legacy per-user token cache and the v2 per-user OAuth token + store; evicting only one cache lets the other keep serving a token minted for the old config.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server r1 = MagicMock(user_id="alice", server_id="srv-1") r2 = MagicMock(user_id="bob", server_id="srv-1") prisma = MagicMock() prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[r1, r2]) - prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) - cache_deletes = [] + invalidations = [] monkeypatch.setattr( - oauth2_token_cache.mcp_per_user_token_cache, - "delete", - AsyncMock(side_effect=lambda uid, sid: cache_deletes.append((uid, sid))), + mcp_server_manager.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + AsyncMock(side_effect=lambda uid, sid: invalidations.append((uid, sid))), ) purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") assert purged == 2 prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() - assert set(cache_deletes) == {("alice", "srv-1"), ("bob", "srv-1")} + assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} + + +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypatch): + from litellm.proxy._experimental.mcp_server import db as db_module + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[MagicMock(user_id="alice", server_id="srv-1")] + ) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) + monkeypatch.setattr( + mcp_server_manager.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + AsyncMock(), + ) + warning = MagicMock() + monkeypatch.setattr(db_module.verbose_proxy_logger, "warning", warning) + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + + assert purged == 2 + warning.assert_called_once() @pytest.mark.asyncio @@ -225,9 +288,7 @@ async def test_store_user_oauth_credential_does_not_persist_plaintext(): access_token = "ya29.a0AfH6SMBverysecretaccesstoken" prisma = _make_prisma_with_existing(row=None) - await store_user_oauth_credential( - prisma, "alice", "srv-1", access_token, refresh_token="rfr-xyz" - ) + await store_user_oauth_credential(prisma, "alice", "srv-1", access_token, refresh_token="rfr-xyz") stored = _stored_value(prisma) try: @@ -310,9 +371,7 @@ async def test_byok_guard_rejects_overwriting_encrypted_byok(): encrypted_row = MagicMock() encrypted_row.credential_b64 = _stored_value(prisma) - prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock( - return_value=encrypted_row - ) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=encrypted_row) with pytest.raises(ValueError, match="could not be verified as an OAuth2"): await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") @@ -354,18 +413,14 @@ async def test_list_oauth_credentials_filters_byok_and_returns_payloads(): "connected_at": "2024-01-01T00:00:00Z", } legacy_row = MagicMock() - legacy_row.credential_b64 = base64.urlsafe_b64encode( - json.dumps(legacy_payload).encode() - ).decode() + legacy_row.credential_b64 = base64.urlsafe_b64encode(json.dumps(legacy_payload).encode()).decode() legacy_row.server_id = "srv-legacy" byok_row = MagicMock() byok_row.credential_b64 = base64.urlsafe_b64encode(b"plain-byok-key").decode() byok_row.server_id = "srv-byok" - prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( - return_value=[encrypted_row, legacy_row, byok_row] - ) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[encrypted_row, legacy_row, byok_row]) results = await list_user_oauth_credentials(prisma, "alice") @@ -415,9 +470,7 @@ async def test_rotate_re_encrypts_byok_with_new_key(monkeypatch): prisma.db.litellm_mcpusercredentials.update = AsyncMock() new_master_key = "rotated-salt-key-9999-9999-9999-9999" - await rotate_mcp_user_credentials_master_key( - prisma_client=prisma, new_master_key=new_master_key - ) + await rotate_mcp_user_credentials_master_key(prisma_client=prisma, new_master_key=new_master_key) update_call = prisma.db.litellm_mcpusercredentials.update.call_args new_stored = update_call.kwargs["data"]["credential_b64"] @@ -445,19 +498,13 @@ async def test_rotate_migrates_legacy_plaintext_rows(monkeypatch): legacy_row.user_id = "alice" legacy_row.server_id = "srv-legacy" legacy_row.credential_b64 = base64.urlsafe_b64encode(b"legacy-plain").decode() - prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( - return_value=[legacy_row] - ) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[legacy_row]) prisma.db.litellm_mcpusercredentials.update = AsyncMock() new_key = "another-rotation-key-aaaa-bbbb-cccc-dddd" - await rotate_mcp_user_credentials_master_key( - prisma_client=prisma, new_master_key=new_key - ) + await rotate_mcp_user_credentials_master_key(prisma_client=prisma, new_master_key=new_key) - new_stored = prisma.db.litellm_mcpusercredentials.update.call_args.kwargs["data"][ - "credential_b64" - ] + new_stored = prisma.db.litellm_mcpusercredentials.update.call_args.kwargs["data"]["credential_b64"] monkeypatch.setenv("LITELLM_SALT_KEY", new_key) assert ( decrypt_value_helper( @@ -485,14 +532,10 @@ async def test_rotate_skips_undecodable_rows(): good_row.server_id = "srv-ok" good_row.credential_b64 = base64.urlsafe_b64encode(b"good-byok").decode() - prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( - return_value=[bad_row, good_row] - ) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[bad_row, good_row]) prisma.db.litellm_mcpusercredentials.update = AsyncMock() - await rotate_mcp_user_credentials_master_key( - prisma_client=prisma, new_master_key="new-key-xxxx" - ) + await rotate_mcp_user_credentials_master_key(prisma_client=prisma, new_master_key="new-key-xxxx") # Only one update call — the good row. assert prisma.db.litellm_mcpusercredentials.update.call_count == 1 @@ -508,9 +551,7 @@ def _oauth_cred(access_token="at-live", refresh_token=None, expires_in_seconds=N if refresh_token is not None: cred["refresh_token"] = refresh_token if expires_in_seconds is not None: - cred["expires_at"] = ( - datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds) - ).isoformat() + cred["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat() return cred @@ -527,12 +568,7 @@ def test_expiry_buffer_treats_soon_to_expire_as_expired(): cred = _oauth_cred(expires_in_seconds=30) assert is_oauth_credential_expired(cred, buffer_seconds=60) is True # A token comfortably beyond the buffer stays valid. - assert ( - is_oauth_credential_expired( - _oauth_cred(expires_in_seconds=600), buffer_seconds=60 - ) - is False - ) + assert is_oauth_credential_expired(_oauth_cred(expires_in_seconds=600), buffer_seconds=60) is False def test_expiry_past_is_expired_regardless_of_buffer(): @@ -554,9 +590,7 @@ async def test_resolve_returns_valid_token_without_refreshing(monkeypatch): refresh = AsyncMock() monkeypatch.setattr(db_mod, "refresh_user_oauth_token", refresh) - cred = _oauth_cred( - access_token="at-live", refresh_token="rt-1", expires_in_seconds=600 - ) + cred = _oauth_cred(access_token="at-live", refresh_token="rt-1", expires_in_seconds=600) result = await resolve_valid_user_oauth_token( user_id="alice", server=MagicMock(), cred=cred, prisma_client=MagicMock() ) @@ -572,15 +606,11 @@ async def test_resolve_refreshes_expired_token_with_refresh_token(monkeypatch): # new token rather than returning None (which left the UI tool list empty). import litellm.proxy._experimental.mcp_server.db as db_mod - refreshed = _oauth_cred( - access_token="at-fresh", refresh_token="rt-2", expires_in_seconds=3600 - ) + refreshed = _oauth_cred(access_token="at-fresh", refresh_token="rt-2", expires_in_seconds=3600) refresh = AsyncMock(return_value=refreshed) monkeypatch.setattr(db_mod, "refresh_user_oauth_token", refresh) - expired = _oauth_cred( - access_token="at-dead", refresh_token="rt-1", expires_in_seconds=-5 - ) + expired = _oauth_cred(access_token="at-dead", refresh_token="rt-1", expires_in_seconds=-5) result = await resolve_valid_user_oauth_token( user_id="alice", server=MagicMock(), cred=expired, prisma_client=MagicMock() ) @@ -599,9 +629,7 @@ async def test_resolve_refreshes_token_expiring_within_buffer(monkeypatch): refresh = AsyncMock(return_value=refreshed) monkeypatch.setattr(db_mod, "refresh_user_oauth_token", refresh) - soon = _oauth_cred( - access_token="at-soon", refresh_token="rt-1", expires_in_seconds=30 - ) + soon = _oauth_cred(access_token="at-soon", refresh_token="rt-1", expires_in_seconds=30) result = await resolve_valid_user_oauth_token( user_id="alice", server=MagicMock(), cred=soon, prisma_client=MagicMock() ) @@ -635,9 +663,7 @@ async def test_resolve_returns_none_when_refresh_fails(monkeypatch): refresh = AsyncMock(return_value=None) monkeypatch.setattr(db_mod, "refresh_user_oauth_token", refresh) - expired = _oauth_cred( - access_token="at-dead", refresh_token="rt-1", expires_in_seconds=-5 - ) + expired = _oauth_cred(access_token="at-dead", refresh_token="rt-1", expires_in_seconds=-5) result = await resolve_valid_user_oauth_token( user_id="alice", server=MagicMock(), cred=expired, prisma_client=MagicMock() ) @@ -654,9 +680,7 @@ async def test_resolve_returns_none_for_missing_credential(monkeypatch): monkeypatch.setattr(db_mod, "refresh_user_oauth_token", refresh) assert ( - await resolve_valid_user_oauth_token( - user_id="alice", server=MagicMock(), cred=None, prisma_client=MagicMock() - ) + await resolve_valid_user_oauth_token(user_id="alice", server=MagicMock(), cred=None, prisma_client=MagicMock()) is None ) assert ( @@ -690,19 +714,13 @@ async def test_rotate_user_env_vars_re_encrypts_with_new_key(monkeypatch): encrypted_old = encrypt_value_helper(json.dumps(values)) prisma = MagicMock() - prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock( - return_value=[_env_var_row(encrypted_old)] - ) + prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[_env_var_row(encrypted_old)]) prisma.db.litellm_mcpuserenvvars.update = AsyncMock() new_master_key = "rotated-env-key-1111-2222-3333-4444" - await rotate_mcp_user_env_vars_master_key( - prisma_client=prisma, new_master_key=new_master_key - ) + await rotate_mcp_user_env_vars_master_key(prisma_client=prisma, new_master_key=new_master_key) - new_stored = prisma.db.litellm_mcpuserenvvars.update.call_args.kwargs["data"][ - "values_b64" - ] + new_stored = prisma.db.litellm_mcpuserenvvars.update.call_args.kwargs["data"]["values_b64"] assert new_stored != encrypted_old, "rotation must produce different ciphertext" monkeypatch.setenv("LITELLM_SALT_KEY", new_master_key) @@ -719,18 +737,14 @@ async def test_rotate_user_env_vars_re_encrypts_with_new_key(monkeypatch): async def test_rotate_user_env_vars_skips_undecryptable_rows(): # A corrupt row must be skipped (not overwritten) so recoverable data is # preserved and one bad row does not abort the rest of the rotation. - good = _env_var_row( - encrypt_value_helper(json.dumps({"A": "1"})), server_id="srv-ok" - ) + good = _env_var_row(encrypt_value_helper(json.dumps({"A": "1"})), server_id="srv-ok") bad = _env_var_row("!!! not encrypted !!!", server_id="srv-corrupt") prisma = MagicMock() prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[bad, good]) prisma.db.litellm_mcpuserenvvars.update = AsyncMock() - await rotate_mcp_user_env_vars_master_key( - prisma_client=prisma, new_master_key="new-key-xxxx" - ) + await rotate_mcp_user_env_vars_master_key(prisma_client=prisma, new_master_key="new-key-xxxx") assert prisma.db.litellm_mcpuserenvvars.update.call_count == 1 where = prisma.db.litellm_mcpuserenvvars.update.call_args.kwargs["where"] @@ -758,9 +772,7 @@ async def test_refresh_user_oauth_token_uses_client_secret_basic(monkeypatch): monkeypatch.setattr(db_mod, "get_async_httpx_client", lambda **kwargs: mock_client) monkeypatch.setattr(db_mod, "store_user_oauth_credential", AsyncMock()) - monkeypatch.setattr( - db_mod, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "new-at"}) - ) + monkeypatch.setattr(db_mod, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "new-at"})) result = await db_mod.refresh_user_oauth_token( prisma_client=MagicMock(), @@ -799,9 +811,7 @@ async def test_refresh_user_oauth_token_defaults_to_client_secret_post(monkeypat monkeypatch.setattr(db_mod, "get_async_httpx_client", lambda **kwargs: mock_client) monkeypatch.setattr(db_mod, "store_user_oauth_credential", AsyncMock()) - monkeypatch.setattr( - db_mod, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "new-at"}) - ) + monkeypatch.setattr(db_mod, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "new-at"})) await db_mod.refresh_user_oauth_token( prisma_client=MagicMock(), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e8fca9ac6ab..e11d78d07ae 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -3328,8 +3328,34 @@ class TestMCPServerManager: assert store.invalidations == [("alice", "srv-1")] @pytest.mark.asyncio - async def test_invalidate_user_oauth_token_cache_swallows_store_errors(self): - """A cache-drop failure must not fail the credential write that triggered it.""" + async def test_invalidate_user_oauth_token_cache_drops_legacy_cache_too(self, monkeypatch): + """A per-user token can be served from the legacy per-user token cache as well as the v2 + store; the shared invalidation must evict both, or the path not evicted keeps serving a + token minted for a replaced credential row until its TTL.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + + class _Store: + async def fetch(self, user_id: str, server_id: str): + return None + + async def invalidate(self, user_id: str, server_id: str) -> None: + return None + + legacy_deletes: list[tuple[str, str]] = [] + monkeypatch.setattr( + manager_module.mcp_per_user_token_cache, + "delete", + AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))), + ) + manager = MCPServerManager(per_user_oauth_token_store=_Store()) + await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + assert legacy_deletes == [("alice", "srv-1")] + + @pytest.mark.asyncio + async def test_invalidate_user_oauth_token_cache_swallows_store_errors(self, monkeypatch): + """A cache-drop failure must not fail the credential write that triggered it, and the + legacy cache must still be evicted after the v2 store drop fails.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module class _Store: async def fetch(self, user_id: str, server_id: str): @@ -3338,8 +3364,15 @@ class TestMCPServerManager: async def invalidate(self, user_id: str, server_id: str) -> None: raise RuntimeError("redis down") + legacy_deletes: list[tuple[str, str]] = [] + monkeypatch.setattr( + manager_module.mcp_per_user_token_cache, + "delete", + AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))), + ) manager = MCPServerManager(per_user_oauth_token_store=_Store()) await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + assert legacy_deletes == [("alice", "srv-1")] @pytest.mark.asyncio async def test_resolve_oauth2_headers_no_user_id(self): diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index 7339420ed68..7093b7c650e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -719,6 +719,56 @@ describe("CreateMCPServer", () => { expect(oauthHook.reset).not.toHaveBeenCalled(); }); + it("does not refetch the tool preview with a discarded token after invalidation", async () => { + // Regression: handleFormValuesChange used to publish the pre-reset antd snapshot into + // formValues after clearHeldOAuthToken, so useTestMCPConnection kept the discarded OAuth + // material (the DCR client minted for the old identity) and sent it on the next tool-preview + // request. + await setupOAuthInteractive(); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://a.example.com/mcp" } }); + }); + act(() => { + oauthHook.onTokenReceived?.({ access_token: "stale-tok" }, { clientId: "client-a", clientSecret: "secret-a" }); + }); + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "Sync_FormValues" } }); + }); + vi.mocked(networking.testMCPToolsListRequest).mockClear(); + + await selectAntOption("Authentication", "API Key"); + + await waitFor(() => expect(vi.mocked(networking.testMCPToolsListRequest)).toHaveBeenCalled()); + for (const call of vi.mocked(networking.testMCPToolsListRequest).mock.calls) { + expect(call[1]?.credentials?.client_id).not.toBe("client-a"); + expect(call[1]?.credentials?.client_secret).not.toBe("secret-a"); + expect(call[1]?.credentials?.access_token).not.toBe("stale-tok"); + } + }); + + it("keeps the held token on an http to sse switch with the same url", async () => { + // Same url means the same resource/audience (RFC 8707): the minted token is still valid, so a + // pure transport swap between the two MCP wire protocols must not force a re-authorize. + await setupOAuthInteractive(); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://a.example.com/mcp" } }); + }); + act(() => { + oauthHook.onTokenReceived?.({ access_token: "tok-a" }, { clientId: "client-a", clientSecret: "secret-a" }); + }); + oauthHook.reset.mockClear(); + + await selectAntOption("Transport Type", "Server-Sent Events (SSE)"); + + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + expect(oauthHook.reset).not.toHaveBeenCalled(); + }); + it("includes token_validation in payload when token_validation_json is filled with valid JSON", async () => { vi.mocked(networking.createMCPServer).mockResolvedValue({ server_id: "new-server-oauth", diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 17c7e7c0b9a..24b8e9eaafb 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -16,6 +16,8 @@ import { MCP_OAUTH2_FLOW_INTERACTIVE, isClientForwardedTokenMode, getOAuthAuthorizationIdentity, + CLEARED_ON_INVALIDATION, + isHeldOAuthTokenStale, } from "./types"; import OAuthFormFields from "./OAuthFormFields"; import TruePassthroughWarning from "./TruePassthroughWarning"; @@ -235,11 +237,9 @@ const CreateMCPServer: React.FC = ({ }); // Discard the held browser-authorized token and its tool preview when the authorization identity - // changes (or the modal closes). For oauth2 the fetched token + DCR client also live in - // form.credentials, and the discovered endpoints in authorization_url/token_url/registration_url, so - // those form fields are reset too; whatever the admin just changed (passed via changedValues) is + // changes (or the modal closes). The CLEARED_ON_INVALIDATION form fields (shared with the edit form + // via types.tsx) are reset too; whatever the admin just changed (passed via changedValues) is // re-applied so the invalidation never wipes their in-flight edit. - const CLEARED_ON_INVALIDATION = ["credentials", "authorization_url", "token_url", "registration_url"] as const; const clearHeldOAuthToken = (changedValues: Record = {}) => { setOauthAccessToken(null); clearTools(); @@ -588,21 +588,11 @@ const CreateMCPServer: React.FC = ({ ? { url: undefined, command: undefined, args: undefined, env: undefined } : { spec_path: undefined, command: undefined, args: undefined, env: undefined }; - const nextValues = - authorizedIdentity === undefined - ? transportValues - : { - ...transportValues, - credentials: undefined, - authorization_url: undefined, - token_url: undefined, - registration_url: undefined, - }; - - form.setFieldsValue(nextValues); - if (authorizedIdentity !== undefined) { + form.setFieldsValue(transportValues); + if (isHeldOAuthTokenStale(form.getFieldsValue(true), authorizedIdentity)) { clearHeldOAuthToken(); } + setFormValues(form.getFieldsValue(true)); }; // Generate options with existing groups and potential new group @@ -672,9 +662,13 @@ const CreateMCPServer: React.FC = ({ const handleFormValuesChange = (changedValues: Record, allValues: Record) => { // Any change to a mint-relevant field (url, auth_type, oauth_flow_type, client creds/scopes, or the // authorization/token/registration endpoints — see getOAuthAuthorizationIdentity) makes a held token - // stale, so discard it and force a fresh authorize. - if (authorizedIdentity !== undefined && getOAuthAuthorizationIdentity(allValues) !== authorizedIdentity) { + // stale, so discard it and force a fresh authorize. When that happens, formValues must be rebuilt + // from the form's post-reset state, not the pre-reset allValues snapshot: the snapshot still holds + // the discarded token in credentials, and useTestMCPConnection reads formValues for tool preview. + if (isHeldOAuthTokenStale(allValues, authorizedIdentity)) { clearHeldOAuthToken(changedValues); + setFormValues({ ...form.getFieldsValue(true), ...changedValues }); + return; } setFormValues(allValues); }; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index d55f993b926..4f3d7b69b01 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -22,15 +22,22 @@ vi.mock("../molecules/notifications_manager", () => ({ const mockOauth: { tokenResponse: any; getTemporaryPayload: (() => Record | null) | null; -} = { tokenResponse: null, getTemporaryPayload: null }; + onTokenReceived: ((token: Record | null) => void) | null; + reset: ReturnType; +} = { tokenResponse: null, getTemporaryPayload: null, onTokenReceived: null, reset: vi.fn() }; vi.mock("@/hooks/useMcpOAuthFlow", () => ({ - useMcpOAuthFlow: (opts: { getTemporaryPayload?: () => Record | null }) => { + useMcpOAuthFlow: (opts: { + getTemporaryPayload?: () => Record | null; + onTokenReceived?: (token: Record | null) => void; + }) => { mockOauth.getTemporaryPayload = opts?.getTemporaryPayload ?? null; + mockOauth.onTokenReceived = opts?.onTokenReceived ?? null; return { startOAuthFlow: vi.fn(), status: "idle", error: null, tokenResponse: mockOauth.tokenResponse, + reset: mockOauth.reset, }; }, })); @@ -92,10 +99,12 @@ vi.mock("./mcp_tool_configuration", () => ({ const mockGetToken = vi.fn(); const mockIsTokenValid = vi.fn(); const mockSetToken = vi.fn(); +const mockRemoveToken = vi.fn(); vi.mock("@/utils/mcpTokenStore", () => ({ getToken: (...args: any[]) => mockGetToken(...args), isTokenValid: (...args: any[]) => mockIsTokenValid(...args), setToken: (...args: any[]) => mockSetToken(...args), + removeToken: (...args: unknown[]) => mockRemoveToken(...args), })); // ── fixtures ────────────────────────────────────────────────────────────────── @@ -451,6 +460,74 @@ describe("MCPServerEdit (auth type switch)", () => { }); }); +describe("MCPServerEdit OAuth token invalidation", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + const renderOAuthEdit = () => + render( + , + ); + + it("invalidates a session-authorized token when the transport switches to stdio", async () => { + // Switching to stdio clears url/auth_type via programmatic form.setFieldsValue, which antd does + // not report through onValuesChange; the explicit recheck in handleTransportChange must catch it. + // Regression: the token used to survive this switch (sessionStorage + hook state kept the old + // token minted for the http url). + renderOAuthEdit(); + + act(() => { + mockOauth.onTokenReceived?.({ access_token: "tok-1" }); + }); + mockOauth.reset.mockClear(); + + await selectAntOption("Transport Type", "Standard Input/Output (stdio)"); + + await waitFor(() => expect(mockOauth.reset).toHaveBeenCalled()); + expect(mockRemoveToken).toHaveBeenCalledWith("oauth_server_1", undefined); + }); + + it("invalidates a session-authorized token when the server URL changes", async () => { + renderOAuthEdit(); + + act(() => { + mockOauth.onTokenReceived?.({ access_token: "tok-1" }); + }); + mockOauth.reset.mockClear(); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://other.example.com/mcp" } }); + }); + + await waitFor(() => expect(mockOauth.reset).toHaveBeenCalled()); + expect(mockRemoveToken).toHaveBeenCalledWith("oauth_server_1", undefined); + }); + + it("keeps a session-authorized token on an http to sse switch with the same url", async () => { + // Same url means the same resource/audience (RFC 8707): the minted token is still valid, so a + // pure transport swap between the two MCP wire protocols must not force a re-authorize. + renderOAuthEdit(); + + act(() => { + mockOauth.onTokenReceived?.({ access_token: "tok-1" }); + }); + mockOauth.reset.mockClear(); + + await selectAntOption("Transport Type", "Server-Sent Events (SSE)"); + + expect(mockOauth.reset).not.toHaveBeenCalled(); + expect(mockRemoveToken).not.toHaveBeenCalled(); + }); +}); + describe("MCPServerEdit (tool allowlist)", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index 55adcf2bb59..de3528ea6a9 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -6,6 +6,8 @@ import { AUTH_TYPE, isClientForwardedTokenMode, getOAuthAuthorizationIdentity, + CLEARED_ON_INVALIDATION, + isHeldOAuthTokenStale, OAUTH_FLOW, MCP_OAUTH2_FLOW_M2M, MCP_OAUTH2_FLOW_INTERACTIVE, @@ -392,11 +394,12 @@ const MCPServerEdit: React.FC = ({ // identity it was minted against (url, auth_type, oauth_flow_type, client creds/scopes, or the // authorization/token/registration endpoints — see getOAuthAuthorizationIdentity). Discards the hook // token (resetOAuthFlow, which re-runs fetchTools to prompt a fresh authorize), the sessionStorage - // token (removeToken, browser-held modes), and the fetched token/DCR client in form.credentials + the - // discovered endpoint fields; the admin's in-flight edit is re-applied so it is never wiped. Only fires - // when a token was actually authorized here (ref set), so a token already valid for the saved server on - // mount is left untouched. Driven from onValuesChange (user input only), never programmatic resets. - const CLEARED_ON_INVALIDATION = ["credentials", "authorization_url", "token_url", "registration_url"] as const; + // token (removeToken, browser-held modes), and the fetched token/DCR client in the shared + // CLEARED_ON_INVALIDATION form fields; the admin's in-flight edit is re-applied so it is never wiped. + // Only fires when a token was actually authorized here (ref set), so a token already valid for the + // saved server on mount is left untouched. Driven from onValuesChange for user input, plus an explicit + // recheck after programmatic setFieldsValue paths (handleTransportChange), which antd does not report + // through onValuesChange. const clearHeldOAuthToken = (changedValues: Record = {}) => { authorizedIdentityRef.current = undefined; if (mcpServer.server_id) { @@ -413,10 +416,7 @@ const MCPServerEdit: React.FC = ({ }; const handleFormValuesChange = (changedValues: Record) => { - if ( - authorizedIdentityRef.current !== undefined && - getOAuthAuthorizationIdentity(form.getFieldsValue(true)) !== authorizedIdentityRef.current - ) { + if (isHeldOAuthTokenStale(form.getFieldsValue(true), authorizedIdentityRef.current)) { clearHeldOAuthToken(changedValues); } }; @@ -539,6 +539,9 @@ const MCPServerEdit: React.FC = ({ stdio_config: undefined, }); } + if (isHeldOAuthTokenStale(form.getFieldsValue(true), authorizedIdentityRef.current)) { + clearHeldOAuthToken(); + } }; const handleSave = async (values: Record) => { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 7386fc98bc2..e894cf4b8dc 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -84,6 +84,21 @@ export const getOAuthAuthorizationIdentity = (values: Record): return JSON.stringify(identity); }; +// The form fields wiped when a held OAuth token is invalidated: the fetched token + DCR client live in +// `credentials`, and the three endpoint fields were discovered by the authorize flow, so all of them are +// stale together with the token. Shared by the create and edit forms so what gets wiped cannot drift. +export const CLEARED_ON_INVALIDATION = ["credentials", "authorization_url", "token_url", "registration_url"] as const; + +// True when a token was authorized in this session (authorizedIdentity recorded at mint time) and the +// form's current identity no longer matches it. Every invalidation decision in both forms goes through +// this single check: onValuesChange for user edits, and an explicit recheck after any programmatic +// form.setFieldsValue (antd does not fire onValuesChange for those), so a missed event path cannot let a +// stale token survive. +export const isHeldOAuthTokenStale = ( + values: Record, + authorizedIdentity: string | undefined, +): boolean => authorizedIdentity !== undefined && getOAuthAuthorizationIdentity(values) !== authorizedIdentity; + // Backend value of `oauth2_flow` that marks a machine-to-machine server. Distinct // from the UI-local OAUTH_FLOW.M2M ("m2m"); this is what the API actually returns. export const MCP_OAUTH2_FLOW_M2M = "client_credentials"; From 42388c3d689807f5e94de9311c40a09448bb488f Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 12:52:18 -0700 Subject: [PATCH 03/12] refactor(mcp): align the invalidation code with the v2 DI and typing discipline The purge takes an injectable invalidate_token_cache callable defaulting to the manager's shared invalidation, and MCPServerManager takes an injectable per_user_token_cache alongside the existing per_user_oauth_token_store, so tests inject fakes instead of monkeypatching the global manager and the module-level cache. The new identity helpers drop Any for object throughout --- litellm/proxy/_experimental/mcp_server/db.py | 32 +++++++++----- .../mcp_server/mcp_server_manager.py | 5 ++- .../mcp_server/test_db_credentials.py | 28 +++++-------- .../mcp_server/test_mcp_server_manager.py | 42 ++++++++++--------- 4 files changed, 57 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 135a9e055d6..96b28afc093 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -3,7 +3,7 @@ import binascii import hashlib import json from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Union, cast +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterable, List, Optional, Set, Union, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -1070,7 +1070,7 @@ async def list_user_oauth_credentials( return results -def _decrypted_credential_field(creds: Dict[str, Any], field: str) -> Any: +def _decrypted_credential_field(creds: Dict[str, object], field: str) -> object: """Return one credential field decrypted with the global salt key; non-string and legacy plaintext values come back unchanged (decrypt_value_helper returns the original on failure).""" value = creds.get(field) @@ -1084,7 +1084,7 @@ def _decrypted_credential_field(creds: Dict[str, Any], field: str) -> Any: ) -def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]: +def mcp_oauth_token_identity(server: object) -> tuple[object, ...]: """The upstream-OAuth-token-determining fields of an MCP server: the resource/audience (url, or spec_path for OpenAPI servers), the OAuth mode/grant (auth_type, oauth2_flow), the authorization-server endpoints, and the OAuth client + scopes. Mirrors the dashboard's @@ -1098,12 +1098,12 @@ def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]: creds = getattr(server, "credentials", None) if isinstance(creds, str): try: - parsed: Any = json.loads(creds) + parsed: object = json.loads(creds) except ValueError: parsed = None else: parsed = creds - creds_dict: Dict[str, Any] = parsed if isinstance(parsed, dict) else {} + creds_dict: Dict[str, object] = parsed if isinstance(parsed, dict) else {} return ( getattr(server, "url", None), getattr(server, "spec_path", None), @@ -1118,25 +1118,35 @@ def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]: ) -async def purge_user_oauth_credentials_for_server(prisma_client: PrismaClient, server_id: str) -> int: +async def purge_user_oauth_credentials_for_server( + prisma_client: PrismaClient, + server_id: str, + invalidate_token_cache: Optional[Callable[[str, str], Awaitable[None]]] = None, +) -> int: """Delete every stored per-user OAuth credential for a server and invalidate each user's cached token everywhere it can be served from (the legacy per-user token cache and the v2 per-user OAuth token store), so no user keeps a token minted for a superseded configuration. Called when a server update changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows removed. A row inserted between the find and the delete is removed from the DB but cannot be evicted from the caches (its user_id was never seen); that case is detected, logged, and bounded - by the cache TTL.""" + by the cache TTL. + + invalidate_token_cache is injectable for tests; it defaults to the manager's shared + invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens.""" repo = MCPUserCredentialsRepository(prisma_client) rows = await repo.table.find_many(where={"server_id": server_id}) if not rows: return 0 deleted_count = await repo.table.delete_many(where={"server_id": server_id}) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + if invalidate_token_cache is None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache for row in rows: - await global_mcp_server_manager.invalidate_user_oauth_token_cache(row.user_id, server_id) + await invalidate_token_cache(row.user_id, server_id) if deleted_count != len(rows): verbose_proxy_logger.warning( "MCP server %s: purge removed %d credential row(s) but %d were enumerated; " diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cc4a9e63105..41a87a17d58 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -58,6 +58,7 @@ from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + MCPPerUserTokenCache, mcp_per_user_token_cache, resolve_mcp_auth, ) @@ -802,10 +803,12 @@ class MCPServerManager: self, cred_provider: Optional[UpstreamCredentialProvider] = None, per_user_oauth_token_store: Optional[InvalidatableOAuthTokenStore] = None, + per_user_token_cache: Optional[MCPPerUserTokenCache] = None, ): self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( self.get_mcp_server_by_id ) + self._per_user_token_cache = per_user_token_cache or mcp_per_user_token_cache self._cred_provider = cred_provider or UpstreamCredentialProvider( oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), @@ -4070,7 +4073,7 @@ class MCPServerManager: verbose_logger.warning( "Failed to invalidate cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc ) - await mcp_per_user_token_cache.delete(user_id, server_id) + await self._per_user_token_cache.delete(user_id, server_id) async def _resolve_oauth2_headers_for_tool_call( self, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 51641991ef9..1615d81fae9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -149,11 +149,11 @@ def test_mcp_oauth_token_identity_detects_change_under_encryption(): @pytest.mark.asyncio -async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(monkeypatch): - """The purge must route each (user, server) through the manager's shared invalidation, which is - the single point covering both the legacy per-user token cache and the v2 per-user OAuth token - store; evicting only one cache lets the other keep serving a token minted for the old config.""" - from litellm.proxy._experimental.mcp_server import mcp_server_manager +async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(): + """The purge must route each (user, server) through the injected invalidator (defaulting to the + manager's shared invalidation, the single point covering both the legacy per-user token cache and + the v2 per-user OAuth token store); evicting only one cache lets the other keep serving a token + minted for the old config.""" from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server r1 = MagicMock(user_id="alice", server_id="srv-1") @@ -163,13 +163,11 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(m prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) invalidations = [] - monkeypatch.setattr( - mcp_server_manager.global_mcp_server_manager, - "invalidate_user_oauth_token_cache", - AsyncMock(side_effect=lambda uid, sid: invalidations.append((uid, sid))), - ) - purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + async def record_invalidation(user_id: str, server_id: str) -> None: + invalidations.append((user_id, server_id)) + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) assert purged == 2 prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() @@ -179,7 +177,6 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(m @pytest.mark.asyncio async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypatch): from litellm.proxy._experimental.mcp_server import db as db_module - from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server prisma = MagicMock() @@ -187,15 +184,10 @@ async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypat return_value=[MagicMock(user_id="alice", server_id="srv-1")] ) prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) - monkeypatch.setattr( - mcp_server_manager.global_mcp_server_manager, - "invalidate_user_oauth_token_cache", - AsyncMock(), - ) warning = MagicMock() monkeypatch.setattr(db_module.verbose_proxy_logger, "warning", warning) - purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=AsyncMock()) assert purged == 2 warning.assert_called_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e11d78d07ae..97fdd186dda 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -3328,11 +3328,10 @@ class TestMCPServerManager: assert store.invalidations == [("alice", "srv-1")] @pytest.mark.asyncio - async def test_invalidate_user_oauth_token_cache_drops_legacy_cache_too(self, monkeypatch): + async def test_invalidate_user_oauth_token_cache_drops_legacy_cache_too(self): """A per-user token can be served from the legacy per-user token cache as well as the v2 store; the shared invalidation must evict both, or the path not evicted keeps serving a token minted for a replaced credential row until its TTL.""" - from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module class _Store: async def fetch(self, user_id: str, server_id: str): @@ -3341,21 +3340,22 @@ class TestMCPServerManager: async def invalidate(self, user_id: str, server_id: str) -> None: return None - legacy_deletes: list[tuple[str, str]] = [] - monkeypatch.setattr( - manager_module.mcp_per_user_token_cache, - "delete", - AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))), - ) - manager = MCPServerManager(per_user_oauth_token_store=_Store()) + class _LegacyCache: + def __init__(self) -> None: + self.deletes: list[tuple[str, str]] = [] + + async def delete(self, user_id: str, server_id: str) -> None: + self.deletes.append((user_id, server_id)) + + legacy_cache = _LegacyCache() + manager = MCPServerManager(per_user_oauth_token_store=_Store(), per_user_token_cache=legacy_cache) await manager.invalidate_user_oauth_token_cache("alice", "srv-1") - assert legacy_deletes == [("alice", "srv-1")] + assert legacy_cache.deletes == [("alice", "srv-1")] @pytest.mark.asyncio - async def test_invalidate_user_oauth_token_cache_swallows_store_errors(self, monkeypatch): + async def test_invalidate_user_oauth_token_cache_swallows_store_errors(self): """A cache-drop failure must not fail the credential write that triggered it, and the legacy cache must still be evicted after the v2 store drop fails.""" - from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module class _Store: async def fetch(self, user_id: str, server_id: str): @@ -3364,15 +3364,17 @@ class TestMCPServerManager: async def invalidate(self, user_id: str, server_id: str) -> None: raise RuntimeError("redis down") - legacy_deletes: list[tuple[str, str]] = [] - monkeypatch.setattr( - manager_module.mcp_per_user_token_cache, - "delete", - AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))), - ) - manager = MCPServerManager(per_user_oauth_token_store=_Store()) + class _LegacyCache: + def __init__(self) -> None: + self.deletes: list[tuple[str, str]] = [] + + async def delete(self, user_id: str, server_id: str) -> None: + self.deletes.append((user_id, server_id)) + + legacy_cache = _LegacyCache() + manager = MCPServerManager(per_user_oauth_token_store=_Store(), per_user_token_cache=legacy_cache) await manager.invalidate_user_oauth_token_cache("alice", "srv-1") - assert legacy_deletes == [("alice", "srv-1")] + assert legacy_cache.deletes == [("alice", "srv-1")] @pytest.mark.asyncio async def test_resolve_oauth2_headers_no_user_id(self): From c75184bec9b6c5337e5c810e59d34635eb340ecf Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 13:10:24 -0700 Subject: [PATCH 04/12] fix(mcp): make the pre-update identity snapshot advisory so a read failure cannot fail the edit The snapshot read only feeds the stale-token purge decision; leaving it unguarded meant a failed read would 500 an edit whose update would have succeeded, and it broke test_edit_mcp_server_redacts_credentials, whose mocked prisma is not awaitable on the un-patched get_mcp_server path. A failure now logs and skips the purge, consistent with the purge half already being best-effort. Adds the first endpoint-level coverage of the edit purge wiring: purge on a mint-relevant change, no purge when the identity is unchanged, and edit success with purge skipped when the snapshot read raises --- .../mcp_management_endpoints.py | 14 +++- .../test_mcp_management_endpoints.py | 76 +++++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 8f6779b17b9..907a17d76d9 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2320,8 +2320,18 @@ if MCP_AVAILABLE: }, ) - # Snapshot the pre-update identity so we can detect a mint-relevant change below. - old_server_record = await get_mcp_server(prisma_client, payload.server_id) + # Snapshot the pre-update identity so we can detect a mint-relevant change below. The read is + # advisory (it only feeds the stale-token purge decision), so a failure skips the purge with a + # warning instead of failing the edit, whose primary job is the update itself. + try: + old_server_record = await get_mcp_server(prisma_client, payload.server_id) + except Exception as exc: # noqa: BLE001 - advisory read; invalidation is best-effort end-to-end + verbose_logger.warning( + "MCP server %s: could not snapshot the pre-update record; skipping the stale-token check: %s", + payload.server_id, + exc, + ) + old_server_record = None # try to update the mcp server mcp_server_record_updated = await update_mcp_server( 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 86bbce36de3..a8a8ee0fbde 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 @@ -5134,3 +5134,79 @@ def test_stamp_oauth2_flow_ignores_non_oauth2(): payload = _oauth2_create_payload(auth_type="none") mgmt_endpoints.stamp_omitted_oauth2_flow(payload) assert payload.oauth2_flow is None + + +async def _run_edit(old_record, updated_record): + from litellm.proxy.management_endpoints.mcp_management_endpoints import edit_mcp_server + + server_id = updated_record.server_id + with ( + patch("litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE", True), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(side_effect=old_record) + if isinstance(old_record, Exception) + else AsyncMock(return_value=old_record), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", + AsyncMock(return_value=updated_record), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + autospec=True, + ), + patch("litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager") as mock_manager, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.purge_user_oauth_credentials_for_server", + AsyncMock(return_value=1), + ) as mock_purge, + ): + mock_manager.update_server = AsyncMock() + mock_manager.reload_servers_from_database = AsyncMock() + payload = UpdateMCPServerRequest(server_id=server_id, alias=updated_record.alias, url=updated_record.url) + user_auth = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + result = await edit_mcp_server(payload=payload, user_api_key_dict=user_auth) + return result, mock_purge + + +@pytest.mark.asyncio +async def test_edit_mcp_server_purges_user_tokens_on_mint_relevant_change(): + server_id = str(uuid.uuid4()) + old = generate_mock_mcp_server_db_record(server_id=server_id, url="https://old.example.com/mcp") + updated = generate_mock_mcp_server_db_record(server_id=server_id, url="https://new.example.com/mcp") + + result, mock_purge = await _run_edit(old, updated) + + assert result.server_id == server_id + mock_purge.assert_awaited_once() + assert mock_purge.await_args.args[1] == server_id + + +@pytest.mark.asyncio +async def test_edit_mcp_server_skips_purge_when_identity_unchanged(): + server_id = str(uuid.uuid4()) + old = generate_mock_mcp_server_db_record(server_id=server_id, alias="Before") + updated = generate_mock_mcp_server_db_record(server_id=server_id, alias="After") + + result, mock_purge = await _run_edit(old, updated) + + assert result.server_id == server_id + mock_purge.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_edit_mcp_server_snapshot_failure_skips_purge_but_edit_succeeds(): + """The pre-update snapshot read is advisory (it only feeds the purge decision); a read failure + must skip the stale-token check with a warning, never fail the edit itself.""" + server_id = str(uuid.uuid4()) + updated = generate_mock_mcp_server_db_record(server_id=server_id, url="https://new.example.com/mcp") + + result, mock_purge = await _run_edit(RuntimeError("db read failed"), updated) + + assert result.server_id == server_id + mock_purge.assert_not_awaited() From e720b5e25a7dfa3365fd5e32b287b906daaa4216 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 13:53:59 -0700 Subject: [PATCH 05/12] test(ui): drop the vacuous access_token assertion from the preview invalidation test The staged access token never reaches formValues (it is not a registered form field), so the assertion could not fail; the DCR client pair is the leak the test actually pins, proven by the mutation run --- .../src/components/mcp_tools/create_mcp_server.test.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index 7093b7c650e..658116ede1f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -744,7 +744,6 @@ describe("CreateMCPServer", () => { for (const call of vi.mocked(networking.testMCPToolsListRequest).mock.calls) { expect(call[1]?.credentials?.client_id).not.toBe("client-a"); expect(call[1]?.credentials?.client_secret).not.toBe("secret-a"); - expect(call[1]?.credentials?.access_token).not.toBe("stale-tok"); } }); From aa351311c0c7edb7f3d52df7a11a9c4c49af0cd9 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 14:33:39 -0700 Subject: [PATCH 06/12] fix(mcp): spare BYOK rows when purging stale OAuth tokens and invalidate caches on server delete LiteLLM_MCPUserCredentials stores BYOK API keys in the same column as per-user OAuth tokens, so the purge on a mint-relevant config change now deletes only rows whose payload decodes as an OAuth2 credential, each by its (user_id, server_id) pair, instead of every row for the server. An api_key server whose url changes purges nothing. delete_mcp_server now also invalidates each enumerated user's cached token so a re-created server reusing the id cannot serve tokens minted for the deleted one, and both cache drops are best-effort --- litellm/proxy/_experimental/mcp_server/db.py | 64 ++++++-- .../mcp_server/mcp_server_manager.py | 7 +- .../mcp_server/test_db_credentials.py | 143 ++++++++++++++++-- .../mcp_server/test_mcp_server.py | 1 + .../mcp_server/test_mcp_server_manager.py | 19 +++ .../test_mcp_management_endpoints.py | 18 ++- 6 files changed, 222 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 96b28afc093..c6b7620b649 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -558,7 +558,11 @@ async def delete_mcp_server_from_virtualkey(): pass -async def delete_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optional[LiteLLM_MCPServerTable]: +async def delete_mcp_server( + prisma_client: PrismaClient, + server_id: str, + invalidate_token_cache: Optional[Callable[[str, str], Awaitable[None]]] = None, +) -> Optional[LiteLLM_MCPServerTable]: """ Delete the mcp server from the db by server_id @@ -569,6 +573,12 @@ async def delete_mcp_server(prisma_client: PrismaClient, server_id: str) -> Opti caller-visible error. Each table is cleaned independently so a failure on one still attempts the other. + Each enumerated credential row's user also gets their cached per-user token + invalidated (legacy cache + v2 store, via invalidate_token_cache, defaulting + to the manager's shared invalidation): the caches are keyed by + (user_id, server_id), so without this a re-created server reusing the same + server_id would serve tokens minted for the deleted server until TTL. + Returns the deleted mcp server record if it exists, otherwise None """ deleted_server = await MCPServerRepository(prisma_client).table.delete( @@ -577,6 +587,18 @@ async def delete_mcp_server(prisma_client: PrismaClient, server_id: str) -> Opti }, ) if deleted_server is not None: + credential_user_ids: List[str] = [] + try: + credential_rows = await prisma_client.db.litellm_mcpusercredentials.find_many( + where={"server_id": server_id} + ) + credential_user_ids = [row.user_id for row in credential_rows] + except Exception as e: # noqa: BLE001 - enumeration is best-effort; cached tokens expire by TTL + verbose_proxy_logger.warning( + "MCP server %s deleted but per-user credential enumeration failed; cached tokens expire by TTL: %s", + server_id, + e, + ) for model, label in ( (prisma_client.db.litellm_mcpusercredentials, "credential"), (prisma_client.db.litellm_mcpuserenvvars, "env var"), @@ -591,6 +613,15 @@ async def delete_mcp_server(prisma_client: PrismaClient, server_id: str) -> Opti label, e, ) + if credential_user_ids: + if invalidate_token_cache is None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache + for user_id in credential_user_ids: + await invalidate_token_cache(user_id, server_id) return deleted_server @@ -1123,21 +1154,30 @@ async def purge_user_oauth_credentials_for_server( server_id: str, invalidate_token_cache: Optional[Callable[[str, str], Awaitable[None]]] = None, ) -> int: - """Delete every stored per-user OAuth credential for a server and invalidate each user's cached + """Delete every stored per-user OAuth token for a server and invalidate each user's cached token everywhere it can be served from (the legacy per-user token cache and the v2 per-user OAuth token store), so no user keeps a token minted for a superseded configuration. Called when a server update changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows - removed. A row inserted between the find and the delete is removed from the DB but cannot be - evicted from the caches (its user_id was never seen); that case is detected, logged, and bounded - by the cache TTL. + removed. + + LiteLLM_MCPUserCredentials also stores BYOK API keys in the same column; only rows whose payload + decodes as an OAuth2 credential (see _decode_oauth_payload) are deleted, because a config change + only invalidates minted tokens, never a user's own stored key. Rows are therefore deleted per + (user_id, server_id) pair rather than by a blanket server_id filter. An OAuth row inserted while + the purge runs for a user not yet enumerated survives; a re-auth completing in the window for an + already-enumerated user is deleted along with the stale row (the pair delete cannot tell them + apart), which costs that user one extra re-auth and nothing else. invalidate_token_cache is injectable for tests; it defaults to the manager's shared invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens.""" repo = MCPUserCredentialsRepository(prisma_client) rows = await repo.table.find_many(where={"server_id": server_id}) - if not rows: + oauth_rows = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None] + if not oauth_rows: return 0 - deleted_count = await repo.table.delete_many(where={"server_id": server_id}) + deleted_count = sum( + [await repo.table.delete_many(where={"user_id": row.user_id, "server_id": server_id}) for row in oauth_rows] + ) if invalidate_token_cache is None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -1145,15 +1185,15 @@ async def purge_user_oauth_credentials_for_server( invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache - for row in rows: + for row in oauth_rows: await invalidate_token_cache(row.user_id, server_id) - if deleted_count != len(rows): + if deleted_count != len(oauth_rows): verbose_proxy_logger.warning( - "MCP server %s: purge removed %d credential row(s) but %d were enumerated; " - "row(s) raced in during the purge and their cached tokens will expire by TTL", + "MCP server %s: purge removed %d OAuth credential row(s) but %d were enumerated; " + "row(s) were deleted concurrently during the purge", server_id, deleted_count, - len(rows), + len(oauth_rows), ) return deleted_count diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 41a87a17d58..8dd9949e17b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4073,7 +4073,12 @@ class MCPServerManager: verbose_logger.warning( "Failed to invalidate cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc ) - await self._per_user_token_cache.delete(user_id, server_id) + try: + await self._per_user_token_cache.delete(user_id, server_id) + except Exception as exc: # noqa: BLE001 - cache drop is best-effort; TTL is the backstop + verbose_logger.warning( + "Failed to drop legacy cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc + ) async def _resolve_oauth2_headers_for_tool_call( self, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 1615d81fae9..0bdcffe1530 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -148,19 +148,30 @@ def test_mcp_oauth_token_identity_detects_change_under_encryption(): assert mcp_oauth_token_identity(unchanged) != mcp_oauth_token_identity(changed) +def _oauth_row(user_id: str, server_id: str = "srv-1"): + """A stored per-user OAuth token row (payload tagged type=oauth2, legacy plain-base64 encoding).""" + row = _legacy_row(json.dumps({"type": "oauth2", "access_token": "tok-" + user_id})) + row.user_id = user_id + row.server_id = server_id + return row + + +def _byok_row(user_id: str, server_id: str = "srv-1"): + """A stored BYOK API key row: the same column, but the payload is a plain string, not OAuth JSON.""" + row = _legacy_row("sk-byok-" + user_id) + row.user_id = user_id + row.server_id = server_id + return row + + @pytest.mark.asyncio -async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(): - """The purge must route each (user, server) through the injected invalidator (defaulting to the - manager's shared invalidation, the single point covering both the legacy per-user token cache and - the v2 per-user OAuth token store); evicting only one cache lets the other keep serving a token - minted for the old config.""" +async def test_purge_user_oauth_credentials_for_server_invalidates_each_user(): + """The purge must route each (user, server) row through the invalidator exactly once.""" from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server - r1 = MagicMock(user_id="alice", server_id="srv-1") - r2 = MagicMock(user_id="bob", server_id="srv-1") prisma = MagicMock() - prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[r1, r2]) - prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice"), _oauth_row("bob")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1) invalidations = [] @@ -170,29 +181,131 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_every_store() purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) assert purged == 2 - prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + assert prisma.db.litellm_mcpusercredentials.delete_many.await_count == 2 assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_spares_byok_rows(): + """Regression: the purge used to delete_many on server_id alone, wiping BYOK API keys that share + the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted, each by + its (user_id, server_id) pair, and only their users' token caches invalidated.""" + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_byok_row("carol"), _oauth_row("alice")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1) + + invalidations = [] + + async def record_invalidation(user_id: str, server_id: str) -> None: + invalidations.append((user_id, server_id)) + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) + + assert purged == 1 + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with( + where={"user_id": "alice", "server_id": "srv-1"} + ) + assert invalidations == [("alice", "srv-1")] + + +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_all_byok_is_noop(): + """An api_key (BYOK-only) server whose identity tuple changes (e.g. its url) must purge nothing.""" + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_byok_row("carol"), _byok_row("dave")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + + assert purged == 0 + prisma.db.litellm_mcpusercredentials.delete_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_defaults_to_manager_invalidator(monkeypatch): + """When no invalidator is injected, the purge must resolve to the manager's shared + invalidate_user_oauth_token_cache, the single point covering both the legacy per-user token cache + and the v2 per-user OAuth token store; a wrong or no-op default silently leaves every cache + serving tokens minted for the superseded config.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1) + + shared_invalidator = AsyncMock() + monkeypatch.setattr( + mcp_server_manager.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + shared_invalidator, + ) + + purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1") + + assert purged == 1 + shared_invalidator.assert_awaited_once_with("alice", "srv-1") + + @pytest.mark.asyncio async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypatch): from litellm.proxy._experimental.mcp_server import db as db_module from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server prisma = MagicMock() - prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( - return_value=[MagicMock(user_id="alice", server_id="srv-1")] - ) - prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=0) warning = MagicMock() monkeypatch.setattr(db_module.verbose_proxy_logger, "warning", warning) purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=AsyncMock()) - assert purged == 2 + assert purged == 0 warning.assert_called_once() +@pytest.mark.asyncio +async def test_delete_mcp_server_invalidates_cached_tokens_for_enumerated_users(): + """Deleting a server must invalidate each enumerated user's cached per-user token: the caches are + keyed by (user_id, server_id), so a re-created server reusing the same server_id would otherwise + serve tokens minted for the deleted server until TTL.""" + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=MagicMock(server_id="srv-1")) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice"), _byok_row("bob")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) + prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock(return_value=0) + + invalidations = [] + + async def record_invalidation(user_id: str, server_id: str) -> None: + invalidations.append((user_id, server_id)) + + deleted = await delete_mcp_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) + + assert deleted is not None + assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} + + +@pytest.mark.asyncio +async def test_delete_mcp_server_returns_none_without_cleanup_when_server_missing(): + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = MagicMock() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=None) + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock() + + deleted = await delete_mcp_server(prisma, "srv-1", invalidate_token_cache=AsyncMock()) + + assert deleted is None + prisma.db.litellm_mcpusercredentials.find_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_purge_user_oauth_credentials_for_server_noop_when_empty(): from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index bba0eb31cfb..7d25e0ba493 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -6869,6 +6869,7 @@ async def test_call_tool_with_legacy_db_m2m_server_resolves_oauth2_flow(): (None, None), ("", None), ("not a url", None), + ("http://[::1", None), ], ) def test_redact_mcp_resource_url_strips_credentials(url, expected): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 97fdd186dda..a18e1ac2c44 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -3376,6 +3376,25 @@ class TestMCPServerManager: await manager.invalidate_user_oauth_token_cache("alice", "srv-1") assert legacy_cache.deletes == [("alice", "srv-1")] + @pytest.mark.asyncio + async def test_invalidate_user_oauth_token_cache_swallows_legacy_cache_errors(self): + """The legacy cache drop is best-effort like the v2 drop: a failure must be logged, never + raised into the credential write that triggered the invalidation.""" + + class _Store: + async def fetch(self, user_id: str, server_id: str): + return None + + async def invalidate(self, user_id: str, server_id: str) -> None: + return None + + class _RaisingLegacyCache: + async def delete(self, user_id: str, server_id: str) -> None: + raise RuntimeError("redis down") + + manager = MCPServerManager(per_user_oauth_token_store=_Store(), per_user_token_cache=_RaisingLegacyCache()) + await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + @pytest.mark.asyncio async def test_resolve_oauth2_headers_no_user_id(self): """Skip lookup entirely when user_api_key_auth has no user_id.""" 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 a8a8ee0fbde..a02acc02502 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 @@ -5136,7 +5136,7 @@ def test_stamp_oauth2_flow_ignores_non_oauth2(): assert payload.oauth2_flow is None -async def _run_edit(old_record, updated_record): +async def _run_edit(old_record, updated_record, purge_mock=None): from litellm.proxy.management_endpoints.mcp_management_endpoints import edit_mcp_server server_id = updated_record.server_id @@ -5163,7 +5163,7 @@ async def _run_edit(old_record, updated_record): patch("litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager") as mock_manager, patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.purge_user_oauth_credentials_for_server", - AsyncMock(return_value=1), + purge_mock if purge_mock is not None else AsyncMock(return_value=1), ) as mock_purge, ): mock_manager.update_server = AsyncMock() @@ -5199,6 +5199,20 @@ async def test_edit_mcp_server_skips_purge_when_identity_unchanged(): mock_purge.assert_not_awaited() +@pytest.mark.asyncio +async def test_edit_mcp_server_purge_failure_does_not_fail_the_edit(): + """The purge is best-effort: a purge exception after a successful update must be swallowed and + logged, never turned into an error response for an edit whose primary job already succeeded.""" + server_id = str(uuid.uuid4()) + old = generate_mock_mcp_server_db_record(server_id=server_id, url="https://old.example.com/mcp") + updated = generate_mock_mcp_server_db_record(server_id=server_id, url="https://new.example.com/mcp") + + result, mock_purge = await _run_edit(old, updated, purge_mock=AsyncMock(side_effect=RuntimeError("db down"))) + + assert result.server_id == server_id + mock_purge.assert_awaited_once() + + @pytest.mark.asyncio async def test_edit_mcp_server_snapshot_failure_skips_purge_but_edit_succeeds(): """The pre-update snapshot read is advisory (it only feeds the purge decision); a read failure From b304620311b3f449b84d37d20ebdc84cf8d4cb20 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 14:33:51 -0700 Subject: [PATCH 07/12] fix(ui): compare url and spec_path independently in the OAuth authorization identity The identity used to pick the audience from spec_path only when values.transport was OPENAPI, but the create form keeps transport in component state rather than form values, so spec_path edits on OpenAPI servers never invalidated a held token. Comparing url and spec_path independently mirrors the backend's mcp_oauth_token_identity and fires regardless of whether transport is present. Invalidation now also wipes only credentials; the admin-typed endpoint fields are kept --- .../mcp_tools/create_mcp_server.tsx | 3 +- .../mcp_tools/mcp_server_edit.test.tsx | 27 ++++++++++++++ .../src/components/mcp_tools/types.test.tsx | 27 ++++++++++++++ .../src/components/mcp_tools/types.tsx | 35 +++++++++++-------- 4 files changed, 77 insertions(+), 15 deletions(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 24b8e9eaafb..eb48fd02474 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -239,7 +239,8 @@ const CreateMCPServer: React.FC = ({ // Discard the held browser-authorized token and its tool preview when the authorization identity // changes (or the modal closes). The CLEARED_ON_INVALIDATION form fields (shared with the edit form // via types.tsx) are reset too; whatever the admin just changed (passed via changedValues) is - // re-applied so the invalidation never wipes their in-flight edit. + // re-applied so the invalidation never wipes their in-flight edit. Admin-typed endpoint fields are + // left alone (see CLEARED_ON_INVALIDATION). const clearHeldOAuthToken = (changedValues: Record = {}) => { setOauthAccessToken(null); clearTools(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index 4f3d7b69b01..e5daf90e992 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -511,6 +511,33 @@ describe("MCPServerEdit OAuth token invalidation", () => { expect(mockRemoveToken).toHaveBeenCalledWith("oauth_server_1", undefined); }); + it("keeps the admin's in-flight endpoint edits when the token is invalidated", async () => { + // Regression: invalidation used to form.resetFields the endpoint fields; with the edit Form's + // initialValues that silently reverted an admin-corrected token_url back to the saved (wrong) + // value while still looking plausible. Only credentials (the minted material) may be wiped. + renderOAuthEdit(); + + const tokenUrlInput = screen.getByPlaceholderText("https://example.com/oauth/token"); + await act(async () => { + fireEvent.change(tokenUrlInput, { target: { value: "https://corrected.example.com/token" } }); + }); + + act(() => { + mockOauth.onTokenReceived?.({ access_token: "tok-1" }); + }); + mockOauth.reset.mockClear(); + + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://moved.example.com/mcp" } }); + }); + + await waitFor(() => expect(mockOauth.reset).toHaveBeenCalled()); + expect((screen.getByPlaceholderText("https://example.com/oauth/token") as HTMLInputElement).value).toBe( + "https://corrected.example.com/token", + ); + }); + it("keeps a session-authorized token on an http to sse switch with the same url", async () => { // Same url means the same resource/audience (RFC 8707): the minted token is still valid, so a // pure transport swap between the two MCP wire protocols must not force a re-authorize. diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx index 846bcfc9e0b..c6faca1fb51 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx @@ -7,9 +7,36 @@ import { handleTransport, handleAuth, getMcpOAuthMode, + getOAuthAuthorizationIdentity, + isHeldOAuthTokenStale, oauth2FlowToFormValue, } from "./types"; +describe("getOAuthAuthorizationIdentity", () => { + // Regression: the identity used to pick the audience from spec_path only when values.transport was + // OPENAPI, but the create form keeps transport in component state, so values.transport was absent and + // spec_path edits on OpenAPI servers never invalidated a held token. + it("changes when spec_path changes even when transport is absent from form values", () => { + const authorized = { auth_type: AUTH_TYPE.OAUTH2, spec_path: "https://a.example.com/openapi.json" }; + const edited = { auth_type: AUTH_TYPE.OAUTH2, spec_path: "https://b.example.com/openapi.json" }; + expect(getOAuthAuthorizationIdentity(edited)).not.toBe(getOAuthAuthorizationIdentity(authorized)); + expect(isHeldOAuthTokenStale(edited, getOAuthAuthorizationIdentity(authorized))).toBe(true); + }); + + it("changes when url changes", () => { + const authorized = { auth_type: AUTH_TYPE.OAUTH2, url: "https://a.example.com/mcp" }; + const edited = { auth_type: AUTH_TYPE.OAUTH2, url: "https://b.example.com/mcp" }; + expect(getOAuthAuthorizationIdentity(edited)).not.toBe(getOAuthAuthorizationIdentity(authorized)); + }); + + it("is stable across non-mint fields", () => { + const authorized = { auth_type: AUTH_TYPE.OAUTH2, url: "https://a.example.com/mcp", server_name: "one" }; + const renamed = { auth_type: AUTH_TYPE.OAUTH2, url: "https://a.example.com/mcp", server_name: "two" }; + expect(getOAuthAuthorizationIdentity(renamed)).toBe(getOAuthAuthorizationIdentity(authorized)); + expect(isHeldOAuthTokenStale(renamed, getOAuthAuthorizationIdentity(authorized))).toBe(false); + }); +}); + describe("handleTransport", () => { it("should default to SSE when transport is null", () => { expect(handleTransport(null)).toBe(TRANSPORT.SSE); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index e894cf4b8dc..3eba8b30968 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -58,20 +58,24 @@ export const OAUTH_FLOW = { }; // The fields that determine which upstream OAuth token "Authorize & Fetch" mints: the resource/audience -// (url), the OAuth mode/grant (auth_type, oauth_flow_type), the OAuth client and requested scope -// (credentials.client_id / client_secret / scopes), and the authorization-server endpoints -// (authorization_url / token_url / registration_url). Grounded in RFC 8707 / RFC 8693 and the MCP auth -// spec: an access token is bound to exactly this tuple (resource/audience + scope + client + issuer), so +// (url, or spec_path for OpenAPI servers), the OAuth mode/grant (auth_type, oauth_flow_type), the OAuth +// client and requested scope (credentials.client_id / client_secret / scopes), and the authorization-server +// endpoints (authorization_url / token_url / registration_url). Grounded in RFC 8707 / RFC 8693 and the MCP +// auth spec: an access token is bound to exactly this tuple (resource/audience + scope + client + issuer), so // a previously authorized token is stale if and only if this identity changes and must be re-minted. -// Deliberately EXCLUDES: transport (http<->sse on the same url is the same audience; a transport switch -// only matters when it changes the target url, which `target` already captures), delegate_auth_to_upstream -// (a downstream-usage toggle that is never sent to the authorize request), and all metadata/RBAC/routing -// fields. Shared by the create and edit forms so their invalidation logic cannot drift. +// url and spec_path are compared independently rather than selected by transport: the create form keeps +// transport in component state, not in form values, so a transport-conditional target would silently pin the +// audience to a missing url and never fire for spec_path edits on OpenAPI servers. Mirrors the backend's +// mcp_oauth_token_identity. Deliberately EXCLUDES: transport itself (http<->sse on the same url is the same +// audience; a switch to/from OpenAPI shows up as url/spec_path changes because each form clears the field the +// new transport does not use), delegate_auth_to_upstream (a downstream-usage toggle that is never sent to the +// authorize request), and all metadata/RBAC/routing fields. Shared by the create and edit forms so their +// invalidation logic cannot drift. export const getOAuthAuthorizationIdentity = (values: Record): string => { const credentials = (values.credentials ?? {}) as Record; - const target = values.transport === TRANSPORT.OPENAPI ? values.spec_path : values.url; const identity = { - target: typeof target === "string" ? target : null, + url: typeof values.url === "string" ? values.url : null, + spec_path: typeof values.spec_path === "string" ? values.spec_path : null, auth_type: values.auth_type ?? null, oauth_flow_type: values.oauth_flow_type ?? null, client_id: credentials.client_id ?? null, @@ -84,10 +88,13 @@ export const getOAuthAuthorizationIdentity = (values: Record): return JSON.stringify(identity); }; -// The form fields wiped when a held OAuth token is invalidated: the fetched token + DCR client live in -// `credentials`, and the three endpoint fields were discovered by the authorize flow, so all of them are -// stale together with the token. Shared by the create and edit forms so what gets wiped cannot drift. -export const CLEARED_ON_INVALIDATION = ["credentials", "authorization_url", "token_url", "registration_url"] as const; +// The form fields wiped when a held OAuth token is invalidated: only `credentials`, which holds the +// minted material (the fetched token + DCR client). The authorization/token/registration endpoint +// fields are deliberately NOT wiped: nothing programmatic ever writes them (upstream discovery happens +// backend-side), so they only ever hold admin input, and resetting them would wipe it (create) or +// silently revert it to the saved record (edit, whose Form has initialValues). Shared by the create and +// edit forms so what gets wiped cannot drift. +export const CLEARED_ON_INVALIDATION = ["credentials"] as const; // True when a token was authorized in this session (authorizedIdentity recorded at mint time) and the // form's current identity no longer matches it. Every invalidation decision in both forms goes through From 71e0491d37cde6568ca90dca31368d835322bad4 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 15:06:43 -0700 Subject: [PATCH 08/12] fix(ui): preview tools with a staged interactive OAuth token in the edit form For authorization_code the edit preview listed tools by server_id only, relying on the stored per-user DB credential, so a token authorized in the edit session gave an empty preview until the admin saved; the create form previews the identical state through the config-based preview endpoint, which takes the token explicitly. The edit fetch now routes through that same endpoint when a staged interactive token is held, built from the form values with the saved record as fallback, and keeps the by-server_id listing for every other case --- .../mcp_tools/mcp_server_edit.test.tsx | 25 +++++++++ .../components/mcp_tools/mcp_server_edit.tsx | 53 ++++++++++++++++++- 2 files changed, 77 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index e5daf90e992..9a9db2eeb95 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -10,6 +10,7 @@ vi.mock("../networking", () => ({ updateMCPServer: vi.fn(), listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}), + testMCPToolsListRequest: vi.fn().mockResolvedValue({ tools: [], error: null }), })); vi.mock("../molecules/notifications_manager", () => ({ @@ -511,6 +512,30 @@ describe("MCPServerEdit OAuth token invalidation", () => { expect(mockRemoveToken).toHaveBeenCalledWith("oauth_server_1", undefined); }); + it("previews tools with a staged interactive OAuth token before it is saved", async () => { + // Regression: for authorization_code the fetch went by server_id only, relying on the stored DB + // credential, so a token authorized in this edit session gave an empty preview until the admin + // saved; the create form previews the identical state via the config-based preview endpoint. + mockOauth.tokenResponse = { access_token: "staged-obo-tok" }; + + renderOAuthEdit(); + + await waitFor(() => { + expect(vi.mocked(networking.testMCPToolsListRequest)).toHaveBeenCalledWith( + "access-token", + expect.objectContaining({ server_id: "oauth_server_1", url: "https://example.com/mcp" }), + "staged-obo-tok", + ); + }); + expect(networking.listMCPTools).not.toHaveBeenCalled(); + // Previewing must stay stateless: the staged token is committed only by an explicit Save + // (storeMCPOAuthUserCredential for authorization_code, setToken for the client-forwarded modes). + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(mockSetToken).not.toHaveBeenCalled(); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + mockOauth.tokenResponse = null; + }); + it("keeps the admin's in-flight endpoint edits when the token is invalidated", async () => { // Regression: invalidation used to form.resetFields the endpoint fields; with the edit Form's // initialValues that silently reverted an admin-corrected token_url back to the saved (wrong) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index de3528ea6a9..58f63e12019 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -17,7 +17,7 @@ import { getMcpOAuthMode, oauth2FlowToFormValue, } from "./types"; -import { updateMCPServer, listMCPTools, storeMCPOAuthUserCredential } from "../networking"; +import { updateMCPServer, listMCPTools, storeMCPOAuthUserCredential, testMCPToolsListRequest } from "../networking"; import { getToken, isTokenValid, removeToken, setToken } from "@/utils/mcpTokenStore"; import { buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils"; import MCPServerCostConfig from "./mcp_server_cost_config"; @@ -421,6 +421,53 @@ const MCPServerEdit: React.FC = ({ } }; + // A token authorized in this edit session for interactive OAuth (authorization_code) is only + // committed to the DB on save, so a plain by-server_id listing cannot use it and the preview would + // stay empty until the admin saves; the create form previews the identical state through the + // config-based preview endpoint, which takes the staged token explicitly. Returns false when there + // is no staged interactive token so fetchTools falls through to the by-server_id listing. + const previewWithStagedInteractiveToken = async ( + isPassthrough: boolean, + isBrowserHeldTokenMode: boolean, + ): Promise => { + const stagedToken = + !isPassthrough && !isBrowserHeldTokenMode && getEffectiveAuthType() === AUTH_TYPE.OAUTH2 + ? oauthTokenResponse?.access_token + : undefined; + if (!stagedToken) { + return false; + } + setIsLoadingTools(true); + setToolsError(null); + try { + const values = form.getFieldsValue(true); + const rawTransport = values.transport || mcpServer.transport; + const previewConfig = { + server_id: mcpServer.server_id, + server_name: values.server_name || mcpServer.server_name || mcpServer.alias, + url: values.url || mcpServer.url, + transport: rawTransport === TRANSPORT.OPENAPI ? TRANSPORT.HTTP : rawTransport, + auth_type: AUTH_TYPE.OAUTH2, + authorization_url: values.authorization_url, + token_url: values.token_url, + registration_url: values.registration_url, + }; + const toolsResponse = await testMCPToolsListRequest(accessToken, previewConfig, stagedToken); + if (toolsResponse.tools && !toolsResponse.error) { + setTools(toolsResponse.tools); + } else { + setTools([]); + setToolsError(toolsResponse.message || "Failed to load tools"); + } + } catch (error) { + setTools([]); + setToolsError(error instanceof Error ? error.message : "Failed to load tools"); + } finally { + setIsLoadingTools(false); + } + return true; + }; + const fetchTools = async () => { if (!accessToken || !mcpServer.server_id) return; @@ -436,6 +483,10 @@ const MCPServerEdit: React.FC = ({ delegate_auth_to_upstream: mcpServer.delegate_auth_to_upstream, }) === "passthrough"; const isBrowserHeldTokenMode = isClientForwardedTokenMode(getEffectiveAuthType()); + + if (await previewWithStagedInteractiveToken(isPassthrough, isBrowserHeldTokenMode)) { + return; + } if (isPassthrough || isBrowserHeldTokenMode) { const token = oauthTokenResponse?.access_token ?? From 4786e599b0f78b4d7ad4be96782b24202404c5ed Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 15:28:39 -0700 Subject: [PATCH 09/12] test(ui): pin the client-forwarded token contract on create and edit The create and edit submit paths for true_passthrough and oauth_delegate persist only the tool configuration: the parametrized create test authorizes, disables the allowlist, and asserts nothing is persisted before submit, then that the create payload carries allowed_tools but no credentials and no occurrence of the token anywhere in the serialized payload, no per-user DB credential is written, and the token is committed to sessionStorage only, keyed to the created server. The edit save test gains the same serialized-payload assertion --- .../mcp_tools/create_mcp_server.test.tsx | 62 +++++++++++++++++++ .../mcp_tools/mcp_server_edit.test.tsx | 1 + 2 files changed, 63 insertions(+) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index 658116ede1f..02374a3ffa4 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -374,6 +374,68 @@ describe("CreateMCPServer", () => { expect(credentials.access_token).toBeUndefined(); }); + it.each([ + ["true_passthrough", "True Passthrough (no LiteLLM auth)"], + ["oauth_delegate", "OAuth Delegate (client-supplied upstream token)"], + ])("persists only tool config on create for %s; the token stays browser-held", async (_authType, optionLabel) => { + oauthHook.tokenResponse = { access_token: "upstream-tok", token_type: "Bearer" }; + await selectHttpTransport(); + + const user = userEvent.setup({ delay: null }); + await user.type(getServerNameInput(), "CF_Server"); + await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp"); + + await selectAntOption("Authentication", optionLabel); + + await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy()); + await act(async () => { + oauthHook.onTokenReceived!({ access_token: "upstream-tok", token_type: "Bearer" }, undefined); + }); + + fireEvent.click(screen.getByRole("button", { name: "Disable all tools" })); + + // Previewing and configuring must stay stateless: nothing is persisted anywhere (server row, + // per-user DB credential, sessionStorage) until the admin submits. + expect(networking.createMCPServer).not.toHaveBeenCalled(); + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(setToken).not.toHaveBeenCalled(); + + const createdServer = { + server_id: "new-cf-server", + server_name: "CF_Server", + alias: "CF_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: _authType, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }; + vi.mocked(networking.createMCPServer).mockResolvedValue(createdServer); + + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => expect(networking.createMCPServer).toHaveBeenCalledTimes(1)); + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + + // Only the tool configuration persists on the server row; the upstream token appears nowhere + // in the create payload and no per-user DB credential is written. The token is committed to + // sessionStorage only, keyed to the created server. + expect(payload.allowed_tools).toEqual([]); + expect(payload.credentials).toBeUndefined(); + expect(JSON.stringify(payload)).not.toContain("upstream-tok"); + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(setToken).toHaveBeenCalledWith( + "new-cf-server", + expect.objectContaining({ access_token: "upstream-tok" }), + undefined, + ); + }); + it("should not show auth value field when None auth type is selected", async () => { await selectHttpTransport(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index 9a9db2eeb95..93ff1333cd8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -1339,6 +1339,7 @@ describe("MCPServerEdit (OAuth token persistence on save)", () => { expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; expect(payload.credentials).toBeUndefined(); + expect(JSON.stringify(payload)).not.toContain("cf-tok"); }, ); From c46c9d46526077139c7c864d886950ac40c912e6 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 16:04:19 -0700 Subject: [PATCH 10/12] docs(mcp): mcp_server_resource docstring matches the origin-only redaction The field doc still said scheme + host + path while the redactor now strips the path along with userinfo, query, and fragment, since hosted MCP servers routinely embed the credential in the path --- litellm/types/utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6b99cfa3314..c093c213e50 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2533,9 +2533,10 @@ class StandardLoggingMCPToolCall(TypedDict, total=False): mcp_server_resource: Optional[str] """ - The upstream MCP server resource identifier (scheme + host + path) the tool call was - forwarded to. Redacted for logging: userinfo, query string, and fragment are stripped so an - upstream URL carrying an embedded token or secret query parameter never reaches log metadata. + The origin (scheme + host + port) of the upstream MCP server the tool call was forwarded + to. Redacted for logging: userinfo, the path, the query string, and the fragment are all + stripped, because hosted MCP servers routinely embed the credential in the URL path and + this value is readable by callers via request logs. Records which upstream received a relayed request; never a credential. """ From 9dcc21cd48aaee6650e2f8063692d5d1b78a1d41 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 16:18:25 -0700 Subject: [PATCH 11/12] refactor(mcp): batch the purge row deletion into one query The per-row delete_many loop becomes a single delete filtered to the enumerated OAuth users' (user_id IN, server_id) pairs; same rows deleted, same BYOK-sparing precision, same count-mismatch detection, one round-trip instead of N --- litellm/proxy/_experimental/mcp_server/db.py | 4 ++-- .../_experimental/mcp_server/test_db_credentials.py | 12 +++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index c6b7620b649..e4f8b0c331d 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1175,8 +1175,8 @@ async def purge_user_oauth_credentials_for_server( oauth_rows = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None] if not oauth_rows: return 0 - deleted_count = sum( - [await repo.table.delete_many(where={"user_id": row.user_id, "server_id": server_id}) for row in oauth_rows] + deleted_count = await repo.table.delete_many( + where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}} ) if invalidate_token_cache is None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 0bdcffe1530..7269774442b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -171,7 +171,7 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user(): prisma = MagicMock() prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice"), _oauth_row("bob")]) - prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) invalidations = [] @@ -181,15 +181,17 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user(): purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) assert purged == 2 - assert prisma.db.litellm_mcpusercredentials.delete_many.await_count == 2 + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with( + where={"server_id": "srv-1", "user_id": {"in": ["alice", "bob"]}} + ) assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} @pytest.mark.asyncio async def test_purge_user_oauth_credentials_for_server_spares_byok_rows(): """Regression: the purge used to delete_many on server_id alone, wiping BYOK API keys that share - the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted, each by - its (user_id, server_id) pair, and only their users' token caches invalidated.""" + the LiteLLM_MCPUserCredentials table. Only rows holding an OAuth2 payload may be deleted (one + batched query filtered to their user_ids), and only their users' token caches invalidated.""" from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server prisma = MagicMock() @@ -205,7 +207,7 @@ async def test_purge_user_oauth_credentials_for_server_spares_byok_rows(): assert purged == 1 prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with( - where={"user_id": "alice", "server_id": "srv-1"} + where={"server_id": "srv-1", "user_id": {"in": ["alice"]}} ) assert invalidations == [("alice", "srv-1")] From db8c872c7d3cf374b143e1785edc3e8c98194f06 Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 16:26:30 -0700 Subject: [PATCH 12/12] fix(ui): staged edit preview sends explicit oauth2_flow and spec_path; invalidation clears the tool list The preview endpoint infers client_credentials when the inherited client_id, client_secret, and token_url are all present (common once DCR or discovery filled them) and then strips the forwarded bearer to preview as M2M, so the staged interactive token was silently unused; sending oauth2_flow=authorization_code bypasses the inference. spec_path now rides along so OpenAPI servers take the spec-based preview path the create form gets. clearHeldOAuthToken also empties the tool list, mirroring the create form's clearTools, so a preview fetched with the discarded token never lingers while the refetch is in flight --- .../mcp_tools/mcp_server_edit.test.tsx | 36 ++++++++++++++++++- .../components/mcp_tools/mcp_server_edit.tsx | 7 ++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index 93ff1333cd8..adb3e161da5 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -523,7 +523,13 @@ describe("MCPServerEdit OAuth token invalidation", () => { await waitFor(() => { expect(vi.mocked(networking.testMCPToolsListRequest)).toHaveBeenCalledWith( "access-token", - expect.objectContaining({ server_id: "oauth_server_1", url: "https://example.com/mcp" }), + // oauth2_flow must be explicit: the preview endpoint infers client_credentials from + // inherited client_id/client_secret/token_url and would strip the staged bearer. + expect.objectContaining({ + server_id: "oauth_server_1", + url: "https://example.com/mcp", + oauth2_flow: "authorization_code", + }), "staged-obo-tok", ); }); @@ -536,6 +542,34 @@ describe("MCPServerEdit OAuth token invalidation", () => { mockOauth.tokenResponse = null; }); + it("previews an OpenAPI server's staged token against its spec_path", async () => { + mockOauth.tokenResponse = { access_token: "staged-obo-tok" }; + + render( + , + ); + + await waitFor(() => { + expect(vi.mocked(networking.testMCPToolsListRequest)).toHaveBeenCalledWith( + "access-token", + expect.objectContaining({ spec_path: "https://example.com/openapi.json" }), + "staged-obo-tok", + ); + }); + mockOauth.tokenResponse = null; + }); + it("keeps the admin's in-flight endpoint edits when the token is invalidated", async () => { // Regression: invalidation used to form.resetFields the endpoint fields; with the edit Form's // initialValues that silently reverted an admin-corrected token_url back to the saved (wrong) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index 58f63e12019..7446c96c40e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -405,6 +405,7 @@ const MCPServerEdit: React.FC = ({ if (mcpServer.server_id) { removeToken(mcpServer.server_id, userID); } + setTools([]); resetOAuthFlow(); form.resetFields([...CLEARED_ON_INVALIDATION]); const preserved = Object.fromEntries( @@ -442,12 +443,18 @@ const MCPServerEdit: React.FC = ({ try { const values = form.getFieldsValue(true); const rawTransport = values.transport || mcpServer.transport; + // oauth2_flow must be explicit: the preview endpoint infers client_credentials from the + // inherited client_id/client_secret/token_url (common once DCR or discovery filled them) and + // would strip the staged bearer to preview as M2M. spec_path keeps OpenAPI servers on the + // spec-based preview path, mirroring the create form's config. const previewConfig = { server_id: mcpServer.server_id, server_name: values.server_name || mcpServer.server_name || mcpServer.alias, url: values.url || mcpServer.url, + spec_path: values.spec_path || mcpServer.spec_path, transport: rawTransport === TRANSPORT.OPENAPI ? TRANSPORT.HTTP : rawTransport, auth_type: AUTH_TYPE.OAUTH2, + oauth2_flow: MCP_OAUTH2_FLOW_INTERACTIVE, authorization_url: values.authorization_url, token_url: values.token_url, registration_url: values.registration_url,