diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260410120000_add_oauth2_flow_to_mcp_server_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260410120000_add_oauth2_flow_to_mcp_server_table/migration.sql new file mode 100644 index 00000000000..03abf3c47a4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260410120000_add_oauth2_flow_to_mcp_server_table/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +-- OAuth2 flow discriminator: "client_credentials" (M2M) | "authorization_code" (interactive); nullable for legacy rows / inference +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth2_flow" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fce95465b55..45c82746117 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -316,6 +316,8 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + // OAuth2: "client_credentials" (M2M) | "authorization_code" (interactive); null = infer on read/write when safe + oauth2_flow String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) is_byok Boolean @default(false) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index e9bd41bb951..4ce41ac95b3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -41,6 +41,9 @@ def _prepare_mcp_server_data( Dict with properly serialized JSON fields """ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + from litellm.proxy._experimental.mcp_server.oauth2_flow_utils import ( + infer_oauth2_flow_for_storage, + ) # Convert model to dict data_dict = data.model_dump(exclude_none=True) @@ -48,6 +51,21 @@ def _prepare_mcp_server_data( if "alias" not in data_dict: data_dict["alias"] = getattr(data, "alias", None) + # Persist oauth2_flow for OAuth2 M2M when the UI omits it (same rule as test MCP client). + _creds_plain = ( + data_dict["credentials"] + if isinstance(data_dict.get("credentials"), dict) + else None + ) + _inferred_flow = infer_oauth2_flow_for_storage( + auth_type=data_dict.get("auth_type") or getattr(data, "auth_type", None), + oauth2_flow=data_dict.get("oauth2_flow") or getattr(data, "oauth2_flow", None), + token_url=data_dict.get("token_url") or getattr(data, "token_url", None), + credentials_plain=_creds_plain, + ) + if _inferred_flow is not None: + data_dict["oauth2_flow"] = _inferred_flow + # Handle credentials serialization credentials = data_dict.get("credentials") if credentials is not None: diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8d3831e75fb..738114f5b52 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -46,6 +46,9 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) +from litellm.proxy._experimental.mcp_server.oauth2_flow_utils import ( + resolve_oauth2_flow_for_runtime, +) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mcp_auth from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, @@ -640,6 +643,18 @@ class MCPServerManager: mcp_oauth_metadata.scopes if mcp_oauth_metadata else None ) + token_url_for_oauth_flow = mcp_server.token_url or ( + getattr(mcp_oauth_metadata, "token_url", None) if mcp_oauth_metadata else None + ) + oauth2_flow_effective = resolve_oauth2_flow_for_runtime( + auth_type=auth_type, + stored_oauth2_flow=getattr(mcp_server, "oauth2_flow", None), + token_url=token_url_for_oauth_flow, + client_id=client_id_value or getattr(mcp_server, "client_id", None), + client_secret=client_secret_value + or getattr(mcp_server, "client_secret", None), + ) + new_server = MCPServer( server_id=mcp_server.server_id, name=name_for_prefix, @@ -656,12 +671,11 @@ class MCPServerManager: client_id=client_id_value or getattr(mcp_server, "client_id", None), client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), - oauth2_flow=getattr(mcp_server, "oauth2_flow", None), + oauth2_flow=oauth2_flow_effective, scopes=resolved_scopes, authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), - token_url=mcp_server.token_url - or getattr(mcp_oauth_metadata, "token_url", None), + token_url=token_url_for_oauth_flow, registration_url=mcp_server.registration_url or getattr(mcp_oauth_metadata, "registration_url", None), command=getattr(mcp_server, "command", None), @@ -2875,12 +2889,15 @@ class MCPServerManager: if server.requires_per_user_auth: should_skip_health_check = True # Skip if auth_type is not none and authentication_token is missing - # (except aws_sigv4 which uses its own credential fields) + # (except aws_sigv4 which uses its own credential fields). + # OAuth2 M2M (client_credentials) uses client_id/client_secret — no static + # authentication_token; _create_mcp_client can still obtain a token, so do not skip. elif ( server.auth_type and server.auth_type != MCPAuth.none and server.auth_type != MCPAuth.aws_sigv4 and not server.authentication_token + and not server.has_client_credentials ): should_skip_health_check = True @@ -2945,6 +2962,7 @@ class MCPServerManager: authorization_url=server.authorization_url, token_url=server.token_url, registration_url=server.registration_url, + oauth2_flow=getattr(server, "oauth2_flow", None), allow_all_keys=server.allow_all_keys, ) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_flow_utils.py b/litellm/proxy/_experimental/mcp_server/oauth2_flow_utils.py new file mode 100644 index 00000000000..0076161fe77 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/oauth2_flow_utils.py @@ -0,0 +1,83 @@ +""" +Persistence and runtime resolution for MCP OAuth2 ``oauth2_flow`` (M2M vs interactive). + +Mirrors the inference used in ``rest_endpoints._execute_with_mcp_client`` so DB-backed +servers created from the UI (which omits ``oauth2_flow``) still persist and load as M2M +when ``client_id``, ``client_secret``, and ``token_url`` are present. + +Interactive-only servers that reuse the same three fields (e.g. some GitHub Enterprise setups) +should set ``oauth2_flow`` to ``authorization_code`` explicitly via API or config. +""" + +from __future__ import annotations + +from typing import Any, Dict, Literal, Optional, Union + +from litellm.types.mcp import MCPAuth + +OAuth2FlowLiteral = Literal["client_credentials", "authorization_code"] + + +def _auth_type_str(auth_type: Optional[Any]) -> Optional[str]: + if auth_type is None: + return None + return auth_type.value if hasattr(auth_type, "value") else str(auth_type) + + +def infer_oauth2_flow_for_storage( + *, + auth_type: Optional[Any], + oauth2_flow: Optional[str], + token_url: Optional[str], + credentials_plain: Optional[Dict[str, Any]], +) -> Optional[OAuth2FlowLiteral]: + """ + Resolve ``oauth2_flow`` to store on ``LiteLLM_MCPServerTable``. + + - Honors explicit ``client_credentials`` or ``authorization_code``. + - If unset and auth is OAuth2, infers ``client_credentials`` when + ``client_id``, ``client_secret``, and ``token_url`` are all present + (same rule as ``_execute_with_mcp_client``). + """ + if oauth2_flow in ("client_credentials", "authorization_code"): + return oauth2_flow # type: ignore[return-value] + if oauth2_flow is not None and str(oauth2_flow).strip() != "": + return None + + if _auth_type_str(auth_type) != MCPAuth.oauth2.value: + return None + + creds = credentials_plain or {} + client_id = creds.get("client_id") + client_secret = creds.get("client_secret") + if client_id and client_secret and token_url: + return "client_credentials" + return None + + +def resolve_oauth2_flow_for_runtime( + *, + auth_type: Optional[Any], + stored_oauth2_flow: Optional[str], + token_url: Optional[str], + client_id: Optional[str], + client_secret: Optional[str], +) -> Optional[OAuth2FlowLiteral]: + """ + Effective ``oauth2_flow`` when constructing ``MCPServer`` from a DB row. + + Uses the stored column when set; otherwise applies the same inference as + :func:`infer_oauth2_flow_for_storage` using decrypted credential fields. + """ + if stored_oauth2_flow in ("client_credentials", "authorization_code"): + return stored_oauth2_flow # type: ignore[return-value] + return infer_oauth2_flow_for_storage( + auth_type=auth_type, + oauth2_flow=None, + token_url=token_url, + credentials_plain=( + {"client_id": client_id, "client_secret": client_secret} + if (client_id or client_secret) + else None + ), + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 99578d006e1..581958979fa 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2589,17 +2589,15 @@ if MCP_AVAILABLE: f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response + # Pre-emptive www-authenticate 401 is only for OAuth2 servers that are neither + # M2M (proxy fetches client_credentials tokens) nor interactive-without-headers + # (those skip 401 here so list_tools/call_tool can return a proper protocol error). for server_name in mcp_servers or []: server = global_mcp_server_manager.get_mcp_server_by_name( server_name, client_ip=_client_ip ) if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: - # For servers that store per-user tokens server-side, skip the - # pre-emptive 401 — the call_tool / list_tools dispatch will look - # up the stored token from Redis / DB and only fail at the MCP - # protocol level if none is found, giving the client a proper - # tool-execution error rather than an HTTP 401. - if server.needs_user_oauth_token: + if server.has_client_credentials or server.needs_user_oauth_token: continue request = StarletteRequest(scope) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 96e221a9ac0..1cdc62644c4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1222,6 +1222,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None allow_all_keys: bool = False available_on_public_internet: bool = True is_byok: bool = False @@ -1292,6 +1293,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None allow_all_keys: bool = False available_on_public_internet: bool = True is_byok: bool = False diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 40d00adeb0b..b810a43f4f4 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -461,9 +461,24 @@ if MCP_AVAILABLE: payload: NewMCPServerRequest, created_by: Optional[str], ) -> LiteLLM_MCPServerTable: + from litellm.proxy._experimental.mcp_server.oauth2_flow_utils import ( + infer_oauth2_flow_for_storage, + ) + now = datetime.utcnow() server_id = payload.server_id or str(uuid.uuid4()) server_name = payload.server_name or payload.alias or server_id + _creds = ( + payload.credentials + if isinstance(payload.credentials, dict) + else None + ) + _oauth2_flow = infer_oauth2_flow_for_storage( + auth_type=payload.auth_type, + oauth2_flow=payload.oauth2_flow, + token_url=payload.token_url, + credentials_plain=_creds, + ) return LiteLLM_MCPServerTable( server_id=server_id, server_name=server_name, @@ -489,6 +504,7 @@ if MCP_AVAILABLE: authorization_url=payload.authorization_url, token_url=payload.token_url, registration_url=payload.registration_url, + oauth2_flow=_oauth2_flow, allow_all_keys=payload.allow_all_keys, available_on_public_internet=payload.available_on_public_internet, ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fce95465b55..45c82746117 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -316,6 +316,8 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + // OAuth2: "client_credentials" (M2M) | "authorization_code" (interactive); null = infer on read/write when safe + oauth2_flow String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) is_byok Boolean @default(false) diff --git a/schema.prisma b/schema.prisma index fce95465b55..45c82746117 100644 --- a/schema.prisma +++ b/schema.prisma @@ -316,6 +316,8 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + // OAuth2: "client_credentials" (M2M) | "authorization_code" (interactive); null = infer on read/write when safe + oauth2_flow String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) is_byok Boolean @default(false) 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 656a9c616e8..b4437fcd3a6 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 @@ -959,6 +959,37 @@ class TestMCPServerManager: assert result.health_check_error is None assert result.last_health_check is not None + @pytest.mark.asyncio + async def test_health_check_server_oauth2_m2m_runs_check(self): + """OAuth2 M2M uses client credentials — health check should run (not stay unknown).""" + manager = MCPServerManager() + + server = MCPServer( + server_id="oauth2-m2m", + name="oauth2-m2m", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + client_id="cid", + client_secret="sec", + token_url="https://idp.example.com/token", + url="http://mcp.example.com", + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = AsyncMock(return_value=mock_client) + + result = await manager.health_check_server("oauth2-m2m") + + manager._create_mcp_client.assert_called_once() + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "oauth2-m2m" + assert result.status == "healthy" + assert result.health_check_error is None + @pytest.mark.asyncio async def test_health_check_server_no_token_skips_check(self): """Test that health check is skipped when auth_type is set but authentication_token is missing""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_utils.py new file mode 100644 index 00000000000..d47d44f2054 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_utils.py @@ -0,0 +1,99 @@ +"""Tests for MCP OAuth2 flow persistence helpers.""" + +import pytest + +from litellm.proxy._experimental.mcp_server.oauth2_flow_utils import ( + infer_oauth2_flow_for_storage, + resolve_oauth2_flow_for_runtime, +) +from litellm.types.mcp import MCPAuth + + +def test_infer_explicit_client_credentials(): + assert ( + infer_oauth2_flow_for_storage( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://idp.example.com/token", + credentials_plain={"client_id": "a", "client_secret": "b"}, + ) + == "client_credentials" + ) + + +def test_infer_explicit_authorization_code(): + assert ( + infer_oauth2_flow_for_storage( + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + token_url="https://github.example.com/token", + credentials_plain={"client_id": "a", "client_secret": "b"}, + ) + == "authorization_code" + ) + + +def test_infer_from_credentials_and_token_url_ui_pattern(): + """UI omits oauth2_flow but sends M2M-style fields (same as _execute_with_mcp_client).""" + assert ( + infer_oauth2_flow_for_storage( + auth_type=MCPAuth.oauth2, + oauth2_flow=None, + token_url="http://127.0.0.1:18901/oauth/token", + credentials_plain={ + "client_id": "cid", + "client_secret": "sec", + }, + ) + == "client_credentials" + ) + + +def test_infer_not_oauth2(): + assert ( + infer_oauth2_flow_for_storage( + auth_type=MCPAuth.api_key, + oauth2_flow=None, + token_url="https://x/token", + credentials_plain={"client_id": "a", "client_secret": "b"}, + ) + is None + ) + + +def test_infer_partial_credentials(): + assert ( + infer_oauth2_flow_for_storage( + auth_type=MCPAuth.oauth2, + oauth2_flow=None, + token_url="https://x/token", + credentials_plain={"client_id": "only-id"}, + ) + is None + ) + + +def test_resolve_runtime_uses_stored(): + assert ( + resolve_oauth2_flow_for_runtime( + auth_type=MCPAuth.oauth2, + stored_oauth2_flow="authorization_code", + token_url="https://gh/token", + client_id="a", + client_secret="b", + ) + == "authorization_code" + ) + + +def test_resolve_runtime_infers_when_stored_null(): + assert ( + resolve_oauth2_flow_for_runtime( + auth_type=MCPAuth.oauth2, + stored_oauth2_flow=None, + token_url="https://idp/token", + client_id="x", + client_secret="y", + ) + == "client_credentials" + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts index 681bf4161ad..01753bb7490 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -36,10 +36,13 @@ export const useMCPServerHealth = () => { { queryKey: mcpServerHealthKeys.lists() }, (oldData) => { if (!oldData) return result; - return oldData.map((h) => { - const updated = result.find((r) => r.server_id === h.server_id); - return updated ?? h; - }); + // Merge by id: update existing rows and append any new server_ids from this fetch + // (single-server recheck previously dropped rows not present in oldData). + const byId = new Map(oldData.map((h) => [h.server_id, h])); + for (const r of result) { + byId.set(r.server_id, r); + } + return Array.from(byId.values()); }, ); } finally { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx index 4a808ca489d..435d879c6c9 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx @@ -45,8 +45,8 @@ const OAuthFormFields: React.FC = ({ } name="oauth_flow_type" @@ -55,14 +55,14 @@ const OAuthFormFields: React.FC = ({ @@ -73,21 +73,33 @@ const OAuthFormFields: React.FC = ({ } name={["credentials", "client_id"]} - rules={[{ required: true, message: "Client ID is required for M2M OAuth" }]} + rules={ + isEditing + ? [] + : [{ required: true, message: "Client ID is required for M2M OAuth" }] + } > } name={["credentials", "client_secret"]} - rules={[{ required: true, message: "Client Secret is required for M2M OAuth" }]} + rules={ + isEditing + ? [] + : [{ required: true, message: "Client Secret is required for M2M OAuth" }] + } > } name="token_url" - rules={[{ required: true, message: "Token URL is required for M2M OAuth" }]} + rules={ + isEditing + ? [] + : [{ required: true, message: "Token URL is required for M2M OAuth" }] + } > 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 b4251267137..dd9a64ea963 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 @@ -349,12 +349,14 @@ describe("CreateMCPServer", () => { await selectAntOption("Authentication", "OAuth"); - // Wait for OAuthFormFields to render (OAuth Flow Type selector is the sentinel) + // Wait for OAuthFormFields to render (OAuth grant type selector is the sentinel) await waitFor(() => { - expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + expect(screen.getByText("OAuth grant type")).toBeInTheDocument(); }); - // OAuthFormFields defaults to INTERACTIVE, so the new fields should appear + // Default grant is M2M; switch to interactive so token validation / PKCE fields appear (tests target this path) + await selectAntOption("OAuth grant type", "Interactive"); + await waitFor(() => { expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); 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 17bcd59c43e..f87fdcf8f44 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 @@ -3,7 +3,15 @@ import { Modal, Tooltip, Form, Select, Input, Switch, Collapse } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { createMCPServer, registerMCPServer } from "../networking"; -import { AUTH_TYPE, DiscoverableMCPServer, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types"; +import { + AUTH_TYPE, + DiscoverableMCPServer, + OAUTH_FLOW, + mapUiOAuthFlowToApi, + MCPServer, + MCPServerCostInfo, + TRANSPORT, +} from "./types"; import OAuthFormFields from "./OAuthFormFields"; import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPConnectionStatus from "./mcp_connection_status"; @@ -87,7 +95,11 @@ const CreateMCPServer: React.FC = ({ const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false; const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2; const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4; - const isM2MFlow = isOAuthAuthType && formValues.oauth_flow_type === OAUTH_FLOW.M2M; + // Use Form.useWatch — formValues from onValuesChange does not update on initialValue mount, so the + // grant-type Select could show M2M while formValues.oauth_flow_type was still undefined (PKCE fields shown). + const oauthFlowTypeWatched = Form.useWatch("oauth_flow_type", form) as string | undefined; + const isM2MFlow = + isOAuthAuthType && (oauthFlowTypeWatched ?? OAUTH_FLOW.M2M) === OAUTH_FLOW.M2M; const persistCreateUiState = () => { if (typeof window === "undefined") { @@ -146,6 +158,7 @@ const CreateMCPServer: React.FC = ({ authorization_url: values.authorization_url, token_url: values.token_url, registration_url: values.registration_url, + oauth2_flow: mapUiOAuthFlowToApi(values.oauth_flow_type), mcp_access_groups: values.mcp_access_groups, static_headers: staticHeaders, command: values.command, @@ -285,6 +298,7 @@ const CreateMCPServer: React.FC = ({ allow_all_keys: allowAllKeysRaw, available_on_public_internet: availableOnPublicInternetRaw, token_validation_json: rawTokenValidationJson, + oauth_flow_type: oauthFlowType, ...restValues } = values; @@ -390,6 +404,9 @@ const CreateMCPServer: React.FC = ({ available_on_public_internet: Boolean(availableOnPublicInternetRaw), static_headers: staticHeaders, ...(tokenValidation !== null && { token_validation: tokenValidation }), + ...(restValues.auth_type === AUTH_TYPE.OAUTH2 && mapUiOAuthFlowToApi(oauthFlowType) + ? { oauth2_flow: mapUiOAuthFlowToApi(oauthFlowType) } + : {}), }; payload.static_headers = staticHeaders; @@ -818,7 +835,7 @@ const CreateMCPServer: React.FC = ({ {isOAuthAuthType && ( ({ updateMCPServer: vi.fn(), - testMCPToolsListRequest: vi.fn().mockResolvedValue({ tools: [], error: null }), + listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), })); vi.mock("../molecules/notifications_manager", () => ({ @@ -48,7 +48,8 @@ const interactiveOAuthServer = { transport: "http", url: "https://example.com/mcp", auth_type: "oauth2", - // No token_url → edit form defaults to INTERACTIVE flow + oauth2_flow: "authorization_code", + // No token_url → interactive flow; oauth2_flow disambiguates from inferred M2M token_url: null, authorization_url: null, registration_url: null, 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 574e7871759..e5518ef064a 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 @@ -2,14 +2,23 @@ import React, { useState, useEffect } from "react"; import { Form, Select, Button as AntdButton, Tooltip, Input, InputNumber } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; -import { AUTH_TYPE, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types"; -import { updateMCPServer, testMCPToolsListRequest } from "../networking"; +import { + AUTH_TYPE, + OAUTH_FLOW, + mapApiOAuthFlowToUi, + mapUiOAuthFlowToApi, + MCPServer, + MCPServerCostInfo, + TRANSPORT, +} from "./types"; +import { updateMCPServer, listMCPTools } from "../networking"; import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPPermissionManagement from "./MCPPermissionManagement"; import MCPToolConfiguration from "./mcp_tool_configuration"; import StdioConfiguration from "./StdioConfiguration"; import MCPLogoSelector from "./MCPLogoSelector"; import { validateMCPServerUrl, validateMCPServerName } from "./utils"; +import OAuthFormFields from "./OAuthFormFields"; import NotificationsManager from "../molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; @@ -37,6 +46,7 @@ const MCPServerEdit: React.FC = ({ const [costConfig, setCostConfig] = useState({}); const [tools, setTools] = useState([]); const [isLoadingTools, setIsLoadingTools] = useState(false); + const [toolsListError, setToolsListError] = useState(null); const [searchValue, setSearchValue] = useState(""); const [aliasManuallyEdited, setAliasManuallyEdited] = useState(false); const [allowedTools, setAllowedTools] = useState([]); @@ -125,6 +135,10 @@ const MCPServerEdit: React.FC = ({ transport, auth_type: AUTH_TYPE.OAUTH2, credentials: values.credentials, + authorization_url: values.authorization_url ?? mcpServer.authorization_url, + token_url: values.token_url ?? mcpServer.token_url, + registration_url: values.registration_url ?? mcpServer.registration_url, + oauth2_flow: mapUiOAuthFlowToApi(values.oauth_flow_type), mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups, static_headers: staticHeaders, command: values.command, @@ -189,7 +203,9 @@ const MCPServerEdit: React.FC = ({ ...mcpServer, transport: effectiveTransport, static_headers: initialStaticHeaders, - oauth_flow_type: mcpServer.token_url ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE, + oauth_flow_type: + mapApiOAuthFlowToUi(mcpServer.oauth2_flow) ?? + (mcpServer.token_url ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE), token_validation_json: mcpServer.token_validation ? JSON.stringify(mcpServer.token_validation, null, 2) : undefined, @@ -271,57 +287,37 @@ const MCPServerEdit: React.FC = ({ } }, [mcpServer]); - // Fetch tools when component mounts or when OAuth token is received - // But only if the server has been properly saved (has a permanent server_id) + // Load tools from the registry (same path as MCP Tools tab) so OAuth M2M and other + // secrets stored only in the DB work — /test/tools/list only sees the request body. useEffect(() => { - // Don't fetch if server hasn't been saved yet (no permanent server_id) if (!mcpServer.server_id || mcpServer.server_id.trim() === "") { return; } fetchTools(); - }, [mcpServer, accessToken, oauthAccessToken]); + }, [mcpServer.server_id, accessToken]); const fetchTools = async () => { if (!accessToken) return; - // HTTP/SSE requires a URL (unless spec_path is set); stdio does not. - if (mcpServer.transport !== "stdio" && !mcpServer.url && !mcpServer.spec_path) return; - - const isM2M = mcpServer.auth_type === AUTH_TYPE.OAUTH2 && !!mcpServer.token_url; - if (mcpServer.auth_type === AUTH_TYPE.OAUTH2 && !isM2M && !oauthAccessToken) { - return; - } - setIsLoadingTools(true); + setToolsListError(null); try { - // Prepare the MCP server config from existing server data - const mcpServerConfig = { - server_id: mcpServer.server_id, - server_name: mcpServer.server_name, - url: mcpServer.url, - transport: mcpServer.transport, - auth_type: mcpServer.auth_type, - mcp_info: mcpServer.mcp_info, - authorization_url: mcpServer.authorization_url, - token_url: mcpServer.token_url, - registration_url: mcpServer.registration_url, - command: mcpServer.command, - args: mcpServer.args, - env: mcpServer.env, - }; - - const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig, oauthAccessToken); + const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id); if (toolsResponse.tools && !toolsResponse.error) { setTools(toolsResponse.tools); + setToolsListError(null); } else { - console.error("Failed to fetch tools:", toolsResponse.message); + const msg = toolsResponse.message || "Failed to fetch MCP tools"; + console.error("Failed to fetch tools:", msg); setTools([]); + setToolsListError(msg); } } catch (error) { console.error("Tools fetch error:", error); setTools([]); + setToolsListError(error instanceof Error ? error.message : String(error)); } finally { setIsLoadingTools(false); } @@ -404,6 +400,7 @@ const MCPServerEdit: React.FC = ({ allow_all_keys: allowAllKeysRaw, available_on_public_internet: availableOnPublicInternetRaw, token_validation_json: rawTokenValidationJson, + oauth_flow_type: oauthFlowType, ...restValues } = values; @@ -575,6 +572,9 @@ const MCPServerEdit: React.FC = ({ ...(tokenValidation !== null || mcpServer.token_validation ? { token_validation: tokenValidation } : {}), + ...(restValues.auth_type === AUTH_TYPE.OAUTH2 && mapUiOAuthFlowToApi(oauthFlowType) + ? { oauth2_flow: mapUiOAuthFlowToApi(oauthFlowType) } + : {}), }; const includeCredentials = restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); @@ -782,179 +782,16 @@ const MCPServerEdit: React.FC = ({ )} {!isStdioTransport && isOAuthAuthType && ( - <> - - OAuth Client ID (optional) - - - - - } - name={["credentials", "client_id"]} - > - - - - OAuth Client Secret (optional) - - - - - } - name={["credentials", "client_secret"]} - > - - - - OAuth Scopes (optional) - - - - - } - name={["credentials", "scopes"]} - > - - - - Token URL Override (optional) - - - - - } - name="token_url" - > - - - - Registration URL Override (optional) - - - - - } - name="registration_url" - > - - - {!isM2MFlow && ( - <> - - Token Validation Rules (optional) - - - - - } - name="token_validation_json" - rules={[ - { - validator: (_: any, value: string) => { - if (!value || value.trim() === "") return Promise.resolve(); - try { - JSON.parse(value); - return Promise.resolve(); - } catch { - return Promise.reject(new Error("Must be valid JSON")); - } - }, - }, - ]} - > - - - - Token Storage TTL (seconds, optional) - - - - - } - name="token_storage_ttl_seconds" - > - - - - )} -
-

Use OAuth to fetch a fresh access token and temporarily save it in the session as the authentication value.

- - {oauthError &&

{oauthError}

} - {oauthStatus === "success" && oauthTokenResponse?.access_token && ( -

- Token fetched. Expires in {oauthTokenResponse.expires_in ?? "?"} seconds. -

- )} -
- + )} {!isStdioTransport && isAwsSigV4AuthType && ( @@ -1099,6 +936,10 @@ const MCPServerEdit: React.FC = ({ = ({ transport: transportType ?? mcpServer.transport, auth_type: currentAuthType ?? mcpServer.auth_type, mcp_info: mcpServer.mcp_info, - oauth_flow_type: (currentTokenUrl ?? mcpServer.token_url) ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE, + oauth_flow_type: + oauthFlowTypeValue ?? + mapApiOAuthFlowToUi(mcpServer.oauth2_flow) ?? + ((currentTokenUrl ?? mcpServer.token_url) ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE), static_headers: currentStaticHeaders ?? mcpServer.static_headers, credentials: currentCredentials, authorization_url: currentAuthorizationUrl ?? mcpServer.authorization_url, diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx index c70d188bdbb..b8a560b0fea 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { ArrowLeftIcon, EyeIcon, EyeOffIcon } from "@heroicons/react/outline"; import { Title, Card, Button, Text, Grid, TabGroup, TabList, TabPanel, TabPanels, Tab, Icon } from "@tremor/react"; -import { MCPServer, handleTransport, handleAuth } from "./types"; +import { MCPServer, formatOAuth2FlowForDisplay, handleTransport, handleAuth, AUTH_TYPE } from "./types"; // TODO: Move Tools viewer from index file import { MCPToolsViewer } from "."; import MCPServerEdit from "./mcp_server_edit"; @@ -141,6 +141,15 @@ export const MCPServerView: React.FC = ({ + {mcpServer.auth_type === AUTH_TYPE.OAUTH2 && ( + + OAuth grant (stored) +
+ {formatOAuth2FlowForDisplay(mcpServer.oauth2_flow)} +
+ + )} + Host URL
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 72d5e4b5aa8..703ea959ec7 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -28,10 +28,16 @@ const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; const { Option } = Select; const MCPServers: React.FC = ({ accessToken, userRole, userID }) => { - const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers(); + const { data: mcpServers, isLoading: isLoadingServers, refetch: refetchMcpServers } = useMCPServers(); // Fetch health status for all servers - const { data: healthStatuses, isLoading: isLoadingHealth, recheckServerHealth, recheckingServerIds } = useMCPServerHealth(); + const { + data: healthStatuses, + isLoading: isLoadingHealth, + recheckServerHealth, + recheckingServerIds, + refetch: refetchMcpHealth, + } = useMCPServerHealth(); // Merge health status data into servers const serversWithHealth = useMemo(() => { @@ -190,7 +196,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) setIsDeletingServer(true); await deleteMCPServer(accessToken, serverIdToDelete); NotificationsManager.success("Deleted MCP Server successfully"); - refetch(); + refetchMcpServers(); } catch (error) { console.error("Error deleting the mcp server:", error); } finally { @@ -210,10 +216,12 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) ? (mcpServers || []).find((server) => server.server_id === serverIdToDelete) : null; - const handleCreateSuccess = (newMcpServer: MCPServer) => { + const handleCreateSuccess = async (newMcpServer: MCPServer) => { setFilteredServers((prev) => [...prev, newMcpServer]); setModalVisible(false); - refetch(); + // Refresh list + health so the new row shows real status (not "unknown" until full page reload). + await refetchMcpServers(); + await refetchMcpHealth(); }; // Memoize the selected server to prevent unnecessary re-renders @@ -236,8 +244,8 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const handleBack = React.useCallback(() => { setEditServer(false); setSelectedServerId(null); - refetch(); - }, [refetch]); + refetchMcpServers(); + }, [refetchMcpServers]); if (!accessToken || !userRole || !userID) { console.log("Missing required authentication parameters", { accessToken, userRole, userID }); @@ -455,7 +463,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) open={!!byokModalServer} onClose={() => setByokModalServer(null)} onSuccess={(_serverId) => { - refetch(); + refetchMcpServers(); setByokModalServer(null); }} accessToken={accessToken || ""} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx index af6890a6d83..5f2aa5d6954 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tool_configuration.tsx @@ -384,14 +384,6 @@ const MCPToolConfiguration: React.FC = ({ )}
- {/* Description */} -
- - Select which tools users can call: Only checked tools will be available for users to - invoke. Unchecked tools will be blocked from execution. - -
- {/* Loading state */} {isLoadingTools && (
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 155abe27ae2..fa2ecd93573 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.test.tsx @@ -1,5 +1,14 @@ import { describe, it, expect } from "vitest"; -import { AUTH_TYPE, OAUTH_FLOW, TRANSPORT, handleTransport, handleAuth } from "./types"; +import { + AUTH_TYPE, + OAUTH_FLOW, + TRANSPORT, + handleTransport, + handleAuth, + mapUiOAuthFlowToApi, + mapApiOAuthFlowToUi, + formatOAuth2FlowForDisplay, +} from "./types"; describe("handleTransport", () => { it("should default to SSE when transport is null", () => { @@ -57,3 +66,22 @@ describe("constants", () => { expect(OAUTH_FLOW.M2M).toBe("m2m"); }); }); + +describe("OAuth2 flow API mapping", () => { + it("maps UI M2M to client_credentials", () => { + expect(mapUiOAuthFlowToApi(OAUTH_FLOW.M2M)).toBe("client_credentials"); + }); + it("maps UI interactive to authorization_code", () => { + expect(mapUiOAuthFlowToApi(OAUTH_FLOW.INTERACTIVE)).toBe("authorization_code"); + }); + it("maps API values back to UI constants", () => { + expect(mapApiOAuthFlowToUi("client_credentials")).toBe(OAUTH_FLOW.M2M); + expect(mapApiOAuthFlowToUi("authorization_code")).toBe(OAUTH_FLOW.INTERACTIVE); + expect(mapApiOAuthFlowToUi(null)).toBeUndefined(); + }); + it("formats display strings", () => { + expect(formatOAuth2FlowForDisplay("client_credentials")).toContain("Machine-to-machine"); + expect(formatOAuth2FlowForDisplay("authorization_code")).toContain("Interactive"); + expect(formatOAuth2FlowForDisplay(null)).toContain("infer"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 814c5d74f46..47fb85be533 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -47,6 +47,29 @@ export const OAUTH_FLOW = { M2M: "m2m", }; +/** Values persisted by the proxy / API (`NewMCPServerRequest.oauth2_flow`). */ +export type OAuth2FlowApi = "client_credentials" | "authorization_code"; + +export function mapUiOAuthFlowToApi(ui: string | null | undefined): OAuth2FlowApi | undefined { + if (ui === OAUTH_FLOW.M2M) return "client_credentials"; + if (ui === OAUTH_FLOW.INTERACTIVE) return "authorization_code"; + return undefined; +} + +export function mapApiOAuthFlowToUi(api: string | null | undefined): typeof OAUTH_FLOW.M2M | typeof OAUTH_FLOW.INTERACTIVE | undefined { + if (api === "client_credentials") return OAUTH_FLOW.M2M; + if (api === "authorization_code") return OAUTH_FLOW.INTERACTIVE; + return undefined; +} + +/** Human-readable label for overview / read-only UI. */ +export function formatOAuth2FlowForDisplay(api: string | null | undefined): string { + if (api === "client_credentials") return "Machine-to-machine (client credentials)"; + if (api === "authorization_code") return "Interactive (authorization code)"; + if (api == null || api === "") return "Not set — proxy may infer from credentials"; + return api; +} + export const TRANSPORT = { SSE: "sse", HTTP: "http", @@ -185,6 +208,8 @@ export interface MCPServer { authorization_url?: string | null; token_url?: string | null; registration_url?: string | null; + /** OAuth2 grant: `client_credentials` (M2M) or `authorization_code` (interactive). */ + oauth2_flow?: string | null; mcp_info?: MCPInfo | null; created_at: string; created_by: string;