feat(mcp): Add tool call and tool list support via UI for Oauth mcps (#28454)

* feat(mcp): cache OAuth token client-side so Tools tab loads without re-auth

After a user creates an OAuth MCP server and completes the authorization
flow, the resulting access token is now stored in sessionStorage keyed by
server_id.  The MCP Tools tab reads this cached token and includes it as
an MCP auth header when listing and invoking tools, so the user never sees
an empty tool list.  When the session ends (tab close / new browser) an
Authorize button re-triggers the flow without leaving the Tools screen.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

* fix(ui/mcp): surface listMCPTools 401 errors so auth gate reappears

listMCPTools previously swallowed all errors (including HTTP 401) by
returning a synthetic { tools: [], error: 'network_error', ... } payload.
That made the useQuery retry-on-401 guard and mcpToolsError dead code,
so expired OAuth tokens never re-triggered the auth gate.

- Throw an enhanced Error with .status attached on non-2xx responses
  (still preserves the legacy shape for true network failures so the
  caller can render a generic message without crashing).
- Clear the cached OAuth session token when the tools query fails with
  401, mirroring callMCPTool's onError handler so the Authorize button
  is shown again.
- Surface mcpToolsError in the existing error banner.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(mcp-tools): stable onSuccess + reuse parsed flow state

- Pass stable setOauthToken setter directly as onSuccess to avoid
  recreating useToolsOAuthFlow's resumeOAuthFlow on every render.
- Reuse the already-parsed FLOW_STATE_KEY value (peeked) instead of
  re-reading and re-parsing sessionStorage in resumeOAuthFlow.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(ui/mcp): restore listMCPTools never-throws contract

The previous fix made listMCPTools throw on HTTP errors while still
returning a synthetic object on network errors. This inconsistent
contract broke existing callers (MCPToolPermissions, MCPAppsPanel,
MCPConnectPicker) which inspect result.error / result.message and
expect the function to never throw.

- Return a normalized { tools: [], error, message, status, ... }
  object on HTTP errors (instead of throwing) so all callers see a
  consistent shape and the user-visible error text from
  result.message is preserved.
- Convert the returned error object into a thrown Error inside the
  one caller that needs it — the useQuery in mcp_tools.tsx — so the
  401 retry/onError handlers still trigger and clear the cached
  OAuth token.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix greptile

* fix(mcp): align OAuth header alias lookup with dashboard sanitization

Backend auth header resolution now matches x-mcp-{alias} keys produced by
the dashboard sanitizer, and the Tools tab re-syncs OAuth tokens when
serverId changes.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(mcp): widen auth header lookup types for list_tools

Accept legacy str | dict server auth maps and annotate list_tools
server_auth_header as Union[str, dict] for mypy.

Co-authored-by: Cursor <cursoragent@cursor.com>

* refactor(ui): extract shared buildCallbackUrl/clearStorage for MCP OAuth hooks

Hoist the duplicate buildCallbackUrl and clearStorage helpers out of
useToolsOAuthFlow and useUserMcpOAuthFlow into a new shared module
src/hooks/mcpOAuthUtils.ts so the two hooks cannot drift if the URL
construction or storage cleanup logic needs to change.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(ui): don't gate M2M OAuth MCP servers behind interactive authorize

M2M (client_credentials) OAuth servers share auth_type="oauth2" with
interactive PKCE servers, but the backend fetches their token internally
and they typically lack a user authorization endpoint. Gating tool
listing on them rendered an Authorize button that would fail or redirect
incorrectly. Detect M2M via the presence of token_url (matching the
existing heuristic in mcp_server_edit.tsx) and skip the auth gate.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(ui/mcp): return error shape when listMCPTools JSON parse fails

Restore the never-throws contract when response.json() fails on a 2xx
body so callers do not receive null and crash on result.tools.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
This commit is contained in:
Sameer Kankute 2026-05-22 21:34:04 +05:30 committed by GitHub
parent 7a93cceb9f
commit ef36e89638
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 846 additions and 125 deletions

View File

@ -1212,11 +1212,17 @@ class MCPServerManager:
return []
# Get server-specific auth header if available
server_auth_header = None
if mcp_server_auth_headers and server.alias:
server_auth_header = mcp_server_auth_headers.get(server.alias)
elif mcp_server_auth_headers and server.server_name:
server_auth_header = mcp_server_auth_headers.get(server.server_name)
server_auth_header: Optional[Union[str, Dict[str, str]]] = None
if mcp_server_auth_headers:
from litellm.proxy._experimental.mcp_server.utils import (
lookup_mcp_server_auth_in_headers,
)
server_auth_header = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
)
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
@ -2707,16 +2713,15 @@ class MCPServerManager:
server_auth_header: Optional[Union[Dict[str, str], str]] = None
if mcp_server_auth_headers:
# Normalize keys for case-insensitive lookup
normalized_headers = {
k.lower(): v for k, v in mcp_server_auth_headers.items()
}
from litellm.proxy._experimental.mcp_server.utils import (
lookup_mcp_server_auth_in_headers,
)
if mcp_server.alias:
server_auth_header = normalized_headers.get(mcp_server.alias.lower())
if server_auth_header is None and mcp_server.server_name:
server_auth_header = normalized_headers.get(
mcp_server.server_name.lower()
)
server_auth_header = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=mcp_server.alias,
server_name=mcp_server.server_name,
)
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:

View File

@ -62,20 +62,16 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str],
) -> Optional[Union[Dict[str, str], str]]:
"""Helper function to get server-specific auth header with case-insensitive matching."""
if mcp_server_auth_headers and server.alias:
normalized_server_alias = server.alias.lower()
normalized_headers = {
k.lower(): v for k, v in mcp_server_auth_headers.items()
}
server_auth = normalized_headers.get(normalized_server_alias)
if server_auth is not None:
return server_auth
elif mcp_server_auth_headers and server.server_name:
normalized_server_name = server.server_name.lower()
normalized_headers = {
k.lower(): v for k, v in mcp_server_auth_headers.items()
}
server_auth = normalized_headers.get(normalized_server_name)
from litellm.proxy._experimental.mcp_server.utils import (
lookup_mcp_server_auth_in_headers,
)
if mcp_server_auth_headers:
server_auth = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
)
if server_auth is not None:
return server_auth
return mcp_auth_header

View File

@ -1114,10 +1114,16 @@ if MCP_AVAILABLE:
) -> Tuple[Optional[Union[Dict[str, str], str]], Optional[Dict[str, str]]]:
"""Build auth and extra headers for a server."""
server_auth_header: Optional[Union[Dict[str, str], str]] = None
if mcp_server_auth_headers and server.alias is not None:
server_auth_header = mcp_server_auth_headers.get(server.alias)
elif mcp_server_auth_headers and server.server_name is not None:
server_auth_header = mcp_server_auth_headers.get(server.server_name)
if mcp_server_auth_headers:
from litellm.proxy._experimental.mcp_server.utils import (
lookup_mcp_server_auth_in_headers,
)
server_auth_header = lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers,
alias=server.alias,
server_name=server.server_name,
)
extra_headers: Optional[Dict[str, str]] = None
if server.auth_type == MCPAuth.oauth2:

View File

@ -2,7 +2,8 @@
MCP Server Utilities
"""
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple
import re
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple, Union
import hashlib
import importlib
@ -117,6 +118,50 @@ def normalize_server_name(server_name: str) -> str:
return server_name.replace(" ", "_")
_MCP_ALIAS_HEADER_INVALID_RE = re.compile(r"[^a-z0-9_]")
def sanitize_mcp_alias_for_header(alias: str) -> str:
"""
Sanitize an MCP server alias for x-mcp-{alias}-{header} HTTP headers.
Must stay in sync with ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts.
"""
sanitized = _MCP_ALIAS_HEADER_INVALID_RE.sub("_", alias.lower().strip())
sanitized = re.sub(r"_+", "_", sanitized)
return sanitized.strip("_")
def lookup_mcp_server_auth_in_headers(
mcp_server_auth_headers: Mapping[str, Union[str, Dict[str, str]]],
*,
alias: Optional[str] = None,
server_name: Optional[str] = None,
) -> Optional[Union[str, Dict[str, str]]]:
"""
Resolve server-specific auth headers with case-insensitive matching.
Tries the raw alias/server_name (lowercased) and the header-safe sanitized
alias so dashboard clients using sanitize_mcp_alias_for_header() still match.
"""
if not mcp_server_auth_headers:
return None
normalized_headers = {k.lower(): v for k, v in mcp_server_auth_headers.items()}
for identifier in (alias, server_name):
if not identifier:
continue
keys_to_try = [identifier.lower()]
sanitized = sanitize_mcp_alias_for_header(identifier)
if sanitized and sanitized not in keys_to_try:
keys_to_try.append(sanitized)
for key in keys_to_try:
if key in normalized_headers:
return normalized_headers[key]
return None
def validate_and_normalize_mcp_server_payload(payload: Any) -> None:
"""
Validate and normalize MCP server payload fields (server_name and alias).

View File

@ -1770,6 +1770,26 @@ def test_get_server_auth_header_fallback_to_default():
assert result == "Bearer default_token"
def test_get_server_auth_header_hyphenated_alias_sanitized_header_key():
"""Header keys use sanitized alias; lookup must match legacy hyphenated aliases."""
from litellm.proxy._experimental.mcp_server.rest_endpoints import (
_get_server_auth_header,
)
mock_server = MagicMock()
mock_server.alias = "GitHub-MCP"
mock_server.server_name = "github_mcp_server"
mcp_server_auth_headers = {
"github_mcp": {"Authorization": "Bearer github-mcp-token"},
}
result = _get_server_auth_header(
mock_server, mcp_server_auth_headers, "Bearer default_token"
)
assert result == {"Authorization": "Bearer github-mcp-token"}
def test_get_server_auth_header_no_auth_headers():
"""Test _get_server_auth_header function with no auth headers."""
from litellm.proxy._experimental.mcp_server.rest_endpoints import (

View File

@ -0,0 +1,18 @@
"""Tests for MCP header alias sanitization and auth header lookup."""
from litellm.proxy._experimental.mcp_server.utils import (
lookup_mcp_server_auth_in_headers,
sanitize_mcp_alias_for_header,
)
def test_sanitize_mcp_alias_for_header():
assert sanitize_mcp_alias_for_header("My Server") == "my_server"
assert sanitize_mcp_alias_for_header("GitHub-MCP!") == "github_mcp"
assert sanitize_mcp_alias_for_header("github_mcp2") == "github_mcp2"
def test_lookup_mcp_server_auth_in_headers_sanitized_alias():
headers = {"github_mcp": {"Authorization": "Bearer token"}}
result = lookup_mcp_server_auth_in_headers(headers, alias="GitHub-MCP")
assert result == {"Authorization": "Bearer token"}

View File

@ -4,11 +4,12 @@ import { Suspense, useEffect, useMemo } from "react";
import { useSearchParams } from "next/navigation";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
// Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the
// user hook (useUserMcpOAuthFlow) can pick up the result. Each hook reads
// its own namespace to avoid cross-flow collisions.
// Written to sessionStorage so the admin hook (useMcpOAuthFlow), the user hook
// (useUserMcpOAuthFlow), and the tools re-auth hook (useToolsOAuthFlow) can each
// pick up the result. Each hook reads its own namespace to avoid cross-flow collisions.
const ADMIN_RESULT_KEY = "litellm-mcp-oauth-result";
const USER_RESULT_KEY = "litellm-user-mcp-oauth-result";
const TOOLS_RESULT_KEY = "litellm-tools-mcp-oauth-result";
const RETURN_URL_STORAGE_KEY = "litellm-mcp-oauth-return-url";
const resolveDefaultRedirect = () => {
@ -50,11 +51,12 @@ const McpOAuthCallbackContent = () => {
}
try {
// Write to both namespace keys (admin and user) so whichever hook is
// active can consume the result. sessionStorage only — no localStorage.
// Write to all namespace keys so whichever hook is active can consume
// the result. sessionStorage only — no localStorage.
const serialized = JSON.stringify(payload);
setSecureItem(ADMIN_RESULT_KEY, serialized);
setSecureItem(USER_RESULT_KEY, serialized);
setSecureItem(TOOLS_RESULT_KEY, serialized);
} catch (err) {
// Silently ignore storage errors
}

View File

@ -3,6 +3,7 @@ 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 { setToken } from "@/utils/mcpTokenStore";
import { AUTH_TYPE, DiscoverableMCPServer, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types";
import OAuthFormFields from "./OAuthFormFields";
import MCPServerCostConfig from "./mcp_server_cost_config";
@ -24,6 +25,7 @@ export const mcpLogoImg = `${asset_logos_folder}mcp_logo.png`;
interface CreateMCPServerProps {
userRole: string;
userID?: string | null;
accessToken: string | null;
onCreateSuccess: (newMcpServer: MCPServer) => void;
isModalVisible: boolean;
@ -47,6 +49,7 @@ const reduceStaticHeaders = (list: unknown): Record<string, string> => {
};
const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
userID,
userRole,
accessToken,
onCreateSuccess,
@ -409,6 +412,21 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
? await createMCPServer(accessToken, payload)
: await registerMCPServer(accessToken, payload);
// Cache the OAuth token in sessionStorage so the Tools tab can use it
// immediately without re-authenticating. No backend DB write.
if (oauthTokenResponse?.access_token && response?.server_id) {
setToken(
response.server_id,
{
access_token: oauthTokenResponse.access_token,
expires_in: oauthTokenResponse.expires_in,
refresh_token: oauthTokenResponse.refresh_token,
token_type: oauthTokenResponse.token_type,
},
userID,
);
}
NotificationsManager.success(
isAdmin
? "MCP Server created successfully"

View File

@ -174,6 +174,7 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
serverId={mcpServer.server_id}
accessToken={accessToken}
auth_type={mcpServer.auth_type}
tokenUrl={mcpServer.token_url}
userRole={userRole}
userID={userID}
serverAlias={mcpServer.alias}

View File

@ -287,6 +287,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
</Modal>
<CreateMCPServer
userRole={userRole}
userID={userID}
accessToken={accessToken}
onCreateSuccess={handleCreateSuccess}
isModalVisible={isModalVisible}

View File

@ -1,17 +1,21 @@
import React, { useState } from "react";
import React, { useEffect, useState } from "react";
import { useQuery, useMutation } from "@tanstack/react-query";
import { ToolTestPanel } from "./ToolTestPanel";
import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse } from "./types";
import { listMCPTools, callMCPTool } from "../networking";
import { isTokenValid, getToken, removeToken } from "@/utils/mcpTokenStore";
import { sanitizeMcpAliasForHeader } from "@/utils/mcpHeaderUtils";
import { useToolsOAuthFlow } from "@/hooks/useToolsOAuthFlow";
import { Card, Title, Text } from "@tremor/react";
import { RobotOutlined, ToolOutlined, SearchOutlined, KeyOutlined } from "@ant-design/icons";
import { RobotOutlined, ToolOutlined, SearchOutlined, KeyOutlined, LockOutlined } from "@ant-design/icons";
import { Input, Button as AntdButton } from "antd";
const MCPToolsViewer = ({
serverId,
accessToken,
auth_type,
tokenUrl,
userRole,
userID,
serverAlias,
@ -21,29 +25,85 @@ const MCPToolsViewer = ({
const [toolResult, setToolResult] = useState<MCPContent[] | null>(null);
const [toolError, setToolError] = useState<Error | null>(null);
const [toolSearchTerm, setToolSearchTerm] = useState("");
// State for passthrough headers
const [passthroughHeaders, setPassthroughHeaders] = useState<Record<string, string>>({});
const [showHeaderInput, setShowHeaderInput] = useState(false);
// OAuth session token (sessionStorage-backed, cleared on tab/browser close).
// Only the interactive (authorization_code/PKCE) flow needs a user-facing
// auth gate. M2M (client_credentials) servers are also `auth_type === "oauth2"`,
// but the backend fetches their token internally — gating tool listing on
// them would force users through a non-existent authorization endpoint.
// We detect M2M via the presence of `tokenUrl`, matching the heuristic in
// `mcp_server_edit.tsx`.
const isOAuth = auth_type === "oauth2" && !tokenUrl;
const [oauthToken, setOauthToken] = useState<string | null>(() =>
isOAuth && isTokenValid(serverId, userID)
? (getToken(serverId, userID)?.access_token ?? null)
: null
);
// Re-sync token when serverId/userID changes (useState initializer only runs on mount).
useEffect(() => {
if (!isOAuth) {
setOauthToken(null);
return;
}
setOauthToken(
isTokenValid(serverId, userID)
? (getToken(serverId, userID)?.access_token ?? null)
: null
);
}, [serverId, userID, isOAuth]);
const { startOAuthFlow, status: oauthStatus, error: oauthError } = useToolsOAuthFlow({
accessToken: accessToken ?? "",
serverId,
serverAlias,
userId: userID,
onSuccess: setOauthToken,
});
// Check if this server has extra headers configured
const hasExtraHeaders = extraHeaders && extraHeaders.length > 0;
// Build custom headers for MCP server requests
const buildCustomHeaders = () => {
if (!serverAlias || !hasExtraHeaders) return undefined;
const customHeaders: Record<string, string> = {};
// Add passthrough headers with server-specific prefix
Object.entries(passthroughHeaders).forEach(([headerName, headerValue]) => {
if (headerValue && headerValue.trim()) {
// Format: x-mcp-{alias}-{header_name}
const mcpHeaderName = `x-mcp-${serverAlias}-${headerName.toLowerCase()}`;
customHeaders[mcpHeaderName] = headerValue;
// Include the session OAuth token using MCP-specific headers so it doesn't
// conflict with the Authorization header used by the LiteLLM proxy itself.
// The backend's _get_mcp_server_auth_headers_from_headers() picks up the
// x-mcp-{alias}-{header} pattern and forwards it to the upstream MCP server.
// When no alias is available, fall back to x-mcp-auth (legacy but still supported).
if (oauthToken) {
if (serverAlias) {
const safeAlias = sanitizeMcpAliasForHeader(serverAlias);
if (safeAlias) {
customHeaders[`x-mcp-${safeAlias}-authorization`] = `Bearer ${oauthToken}`;
} else {
customHeaders["x-mcp-auth"] = `Bearer ${oauthToken}`;
}
} else {
customHeaders["x-mcp-auth"] = `Bearer ${oauthToken}`;
}
});
}
// Add passthrough headers with server-specific prefix
if (serverAlias && hasExtraHeaders) {
const safeAlias = sanitizeMcpAliasForHeader(serverAlias);
if (safeAlias) {
Object.entries(passthroughHeaders).forEach(([headerName, headerValue]) => {
if (headerValue && headerValue.trim()) {
// Format: x-mcp-{alias}-{header_name}
const mcpHeaderName = `x-mcp-${safeAlias}-${headerName.toLowerCase()}`;
customHeaders[mcpHeaderName] = headerValue;
}
});
}
}
return Object.keys(customHeaders).length > 0 ? customHeaders : undefined;
};
@ -54,15 +114,55 @@ const MCPToolsViewer = ({
error: mcpToolsError,
refetch: refetchTools,
} = useQuery({
queryKey: ["mcpTools", serverId, passthroughHeaders],
queryFn: () => {
queryKey: ["mcpTools", serverId, passthroughHeaders, oauthToken],
queryFn: async () => {
if (!accessToken) throw new Error("Access Token required");
return listMCPTools(accessToken, serverId, buildCustomHeaders());
const result = await listMCPTools(accessToken, serverId, buildCustomHeaders());
// listMCPTools never throws — surface error responses as thrown errors
// here so useQuery's retry/onError can react (e.g. clear the cached
// OAuth token on 401).
if (result?.error) {
const status = (result as { status?: number }).status;
if (status === 401) {
removeToken(serverId, userID);
}
const enhancedError = new Error(
result.message || result.error || "Failed to fetch MCP tools",
) as Error & {
status?: number;
statusText?: string;
details?: any;
};
enhancedError.status = status;
enhancedError.statusText = (result as any).statusText;
enhancedError.details = (result as any).details;
throw enhancedError;
}
return result;
},
enabled: !!accessToken,
// For OAuth servers, block the query until a session token is available
enabled: !!accessToken && (!isOAuth || oauthToken !== null),
staleTime: 30000, // Consider data fresh for 30 seconds
retry: (failureCount, error: any) => {
// Don't retry on 401 — token is invalid, user must re-authenticate
if (error?.status === 401 || error?.response?.status === 401) return false;
return failureCount < 2;
},
});
// If the tools query fails with 401, the cached OAuth token is invalid —
// clear it so the auth gate is shown again and the user can re-authenticate.
useEffect(() => {
const err = mcpToolsError as
| (Error & { status?: number; response?: { status?: number } })
| null;
const status = err?.status ?? err?.response?.status;
if (status === 401) {
removeToken(serverId, userID);
setOauthToken(null);
}
}, [mcpToolsError, serverId, userID]);
// Mutation for calling a tool
const { mutate: executeTool, isPending: isCallingTool } = useMutation({
mutationFn: async (args: { tool: MCPTool; arguments: Record<string, any> }) => {
@ -85,9 +185,14 @@ const MCPToolsViewer = ({
setToolResult(data.content);
setToolError(null);
},
onError: (error: Error) => {
onError: (error: Error & { status?: number; response?: { status?: number } }) => {
setToolError(error);
setToolResult(null);
// On 401, clear the cached token so the auth gate is shown again
if (error?.status === 401 || (error as any)?.response?.status === 401) {
removeToken(serverId, userID);
setOauthToken(null);
}
},
});
@ -197,7 +302,31 @@ const MCPToolsViewer = ({
)}
</Text>
{/* Search Bar */}
{/* OAuth Auth Gate — shown when token is absent for OAuth servers */}
{isOAuth && !oauthToken && (
<div className="p-4 text-center bg-white border border-gray-200 rounded-lg">
<LockOutlined className="text-2xl text-gray-400 mb-2" />
<p className="text-xs font-medium text-gray-700 mb-1">Authentication required</p>
<p className="text-xs text-gray-500 mb-3">
Authenticate to view available tools
</p>
<AntdButton
size="small"
type="primary"
loading={oauthStatus === "authorizing" || oauthStatus === "exchanging"}
onClick={startOAuthFlow}
disabled={!accessToken}
>
Authorize
</AntdButton>
{oauthError && (
<p className="text-xs text-red-500 mt-2">{oauthError}</p>
)}
</div>
)}
{/* Search Bar — only shown when tools are loaded */}
{!isOAuth || oauthToken ? <>
{toolsData.length > 0 && (
<div className="mb-3">
<Input
@ -224,14 +353,16 @@ const MCPToolsViewer = ({
)}
{/* Error State */}
{mcpToolsResponse?.error && !isLoadingTools && !toolsData.length && (
{(mcpToolsResponse?.error || mcpToolsError) && !isLoadingTools && !toolsData.length && (
<div className="p-3 text-xs text-red-800 rounded-lg bg-red-50 border border-red-200">
<p className="font-medium">Error: {mcpToolsResponse.message}</p>
<p className="font-medium">
Error: {mcpToolsResponse?.message || (mcpToolsError as Error)?.message}
</p>
</div>
)}
{/* No Tools State */}
{!isLoadingTools && !mcpToolsResponse?.error && (!toolsData || toolsData.length === 0) && (
{!isLoadingTools && !mcpToolsResponse?.error && !mcpToolsError && (!toolsData || toolsData.length === 0) && (
<div className="p-4 text-center bg-white border border-gray-200 rounded-lg">
<div className="mx-auto w-8 h-8 bg-gray-200 rounded-full flex items-center justify-center mb-2">
<svg className="w-4 h-4 text-gray-400" fill="none" stroke="currentColor" viewBox="0 0 24 24">
@ -315,6 +446,7 @@ const MCPToolsViewer = ({
)}
</>
)}
</> : null}
</div>
</div>
</div>

View File

@ -163,6 +163,13 @@ export interface MCPToolsViewerProps {
serverId: string;
accessToken: string | null;
auth_type?: string | null;
/**
* When set, indicates the server uses the OAuth2 M2M (client_credentials)
* flow the backend handles token acquisition internally, so the UI must
* not gate tool listing behind an interactive PKCE authorization. Mirrors
* the heuristic used in `mcp_server_edit.tsx` (`token_url` set => M2M).
*/
tokenUrl?: string | null;
userRole: string | null;
userID: string | null;
serverAlias?: string | null;

View File

@ -7068,46 +7068,33 @@ export const testSearchToolConnection = async (accessToken: string, litellmParam
};
export const listMCPTools = async (
accessToken: string,
accessToken: string,
serverId: string,
customHeaders?: Record<string, string>
customHeaders?: Record<string, string>,
) => {
// Construct base URL
let url = proxyBaseUrl
? `${proxyBaseUrl}/mcp-rest/tools/list?server_id=${serverId}`
: `/mcp-rest/tools/list?server_id=${serverId}`;
console.log("Fetching MCP tools from:", url);
const headers: Record<string, string> = {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
...customHeaders, // Merge custom headers for passthrough auth
};
let response: Response;
try {
// Construct base URL
let url = proxyBaseUrl
? `${proxyBaseUrl}/mcp-rest/tools/list?server_id=${serverId}`
: `/mcp-rest/tools/list?server_id=${serverId}`;
console.log("Fetching MCP tools from:", url);
const headers: Record<string, string> = {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
...customHeaders, // Merge custom headers for passthrough auth
};
const response = await fetch(url, {
response = await fetch(url, {
method: "GET",
headers,
});
const data = await response.json();
console.log("Fetched MCP tools response:", data);
if (!response.ok) {
// If the server returned an error response, use it
if (data.error && data.message) {
throw new Error(data.message);
}
// Otherwise use a generic error
throw new Error("Failed to fetch MCP tools");
}
// Return the full response object which includes tools, error, message, and stack_trace
return data;
} catch (error) {
console.error("Failed to fetch MCP tools:", error);
// Return an error response in the same format as the API
// Network-level failure (no HTTP response). Preserve legacy shape so the
// caller can render a generic error message without crashing.
console.error("Failed to fetch MCP tools (network error):", error);
return {
tools: [],
error: "network_error",
@ -7115,6 +7102,44 @@ export const listMCPTools = async (
stack_trace: null,
};
}
let data: any = null;
try {
data = await response.json();
} catch (parseError) {
console.error("Failed to parse MCP tools response:", parseError);
return {
tools: [],
error: "parse_error",
message: "Failed to parse MCP tools response",
status: response.status,
statusText: response.statusText,
stack_trace: null,
};
}
console.log("Fetched MCP tools response:", data);
if (!response.ok) {
// Preserve the legacy "never throws" contract so existing callers
// (e.g. MCPToolPermissions, MCPAppsPanel, MCPConnectPicker) can continue
// to inspect `result.error` / `result.message`. Attach `status` so
// callers that need to react to auth failures (e.g. the useQuery in
// mcp_tools.tsx) can still detect 401s from the returned object.
const errorMessage =
(data && (data.message || data.error)) || "Failed to fetch MCP tools";
return {
tools: [],
error: (data && data.error) || `http_${response.status}`,
message: errorMessage,
status: response.status,
statusText: response.statusText,
details: data,
stack_trace: null,
};
}
// Return the full response object which includes tools, error, message, and stack_trace
return data;
};
interface CallMCPToolOptions {

View File

@ -0,0 +1,39 @@
/**
* Shared utilities for MCP OAuth2 PKCE flow hooks.
*
* These helpers are used by both useToolsOAuthFlow and useUserMcpOAuthFlow
* to avoid divergence in URL construction and storage cleanup logic.
*/
import { getProxyBaseUrl, serverRootPath } from "@/components/networking";
/**
* Build the OAuth callback URL for the current UI deployment.
*
* In the browser, derive the `/ui` prefix from the current pathname so the
* callback works regardless of how the proxy is mounted. Outside the browser
* (SSR), fall back to the configured proxy base URL and server root path.
*/
export const buildCallbackUrl = (): string => {
if (typeof window !== "undefined") {
const path = window.location.pathname || "";
const idx = path.indexOf("/ui");
const prefix = idx >= 0 ? path.slice(0, idx + 3).replace(/\/+$/, "") : "";
return `${window.location.origin}${prefix}/mcp/oauth/callback`;
}
const base = (getProxyBaseUrl() || "").replace(/\/+$/, "");
const root = serverRootPath && serverRootPath !== "/" ? serverRootPath : "";
return `${base}${root}/ui/mcp/oauth/callback`;
};
/**
* Remove the given keys from sessionStorage, ignoring errors (e.g. storage
* disabled by browser privacy settings).
*/
export const clearStorage = (...keys: string[]): void => {
keys.forEach((k) => {
try {
window.sessionStorage.removeItem(k);
} catch (_) {}
});
};

View File

@ -224,12 +224,21 @@ export const useMcpOAuthFlow = ({
if (!storedPayload) {
return;
}
// Guard: the callback page writes to the admin result key for *all* OAuth
// flows (including the tools re-auth flow). Only proceed if this hook's
// own flow state exists, meaning startOAuthFlow() was actually called here.
// Without this guard, a tools re-auth redirect triggers a spurious
// "OAuth session state was lost" error from this hook.
const storedFlowState = getStorageItem(FLOW_STATE_KEY);
if (!storedFlowState) {
return;
}
// Mark as processing
processingRef.current = true;
payload = JSON.parse(storedPayload);
const storedFlowState = getStorageItem(FLOW_STATE_KEY);
flowState = storedFlowState ? JSON.parse(storedFlowState) : null;
flowState = JSON.parse(storedFlowState);
} catch (err) {
clearStoredFlow();
processingRef.current = false;

View File

@ -0,0 +1,232 @@
"use client";
/**
* OAuth2 PKCE flow for the Tools screen re-authentication path.
*
* Unlike useUserMcpOAuthFlow (used in the chat panel), this hook:
* - stores the resulting token in sessionStorage via mcpTokenStore only
* - does NOT call storeMCPOAuthUserCredential (no backend DB write)
* - uses "litellm-tools-mcp-oauth-result" as its result key to avoid
* collisions with the admin and user flows
*
* The OAuth callback page (src/app/mcp/oauth/callback/page.tsx) writes
* to this key so this hook can pick up the result after the redirect.
*/
import { useCallback, useEffect, useRef, useState } from "react";
import {
buildMcpOAuthAuthorizeUrl,
exchangeMcpOAuthToken,
registerMcpOAuthClient,
} from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { extractErrorMessage } from "@/utils/errorUtils";
import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
import { setToken } from "@/utils/mcpTokenStore";
import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils";
export type ToolsOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
interface UseToolsOAuthFlowOptions {
accessToken: string;
serverId: string;
serverAlias?: string | null;
userId?: string | null;
scopes?: string[];
clientId?: string | null;
onSuccess: (accessToken: string) => void;
}
interface UseToolsOAuthFlowResult {
startOAuthFlow: () => Promise<void>;
status: ToolsOAuthStatus;
error: string | null;
}
const FLOW_STATE_KEY = "litellm-tools-mcp-oauth-flow-state";
const RESULT_KEY = "litellm-tools-mcp-oauth-result";
const RETURN_URL_KEY = "litellm-mcp-oauth-return-url";
type StoredFlowState = {
state: string;
codeVerifier: string;
serverId: string;
redirectUri: string;
clientId?: string;
clientSecret?: string;
scopes?: string[];
};
export const useToolsOAuthFlow = ({
accessToken,
serverId,
serverAlias,
userId,
scopes,
clientId: preClientId,
onSuccess,
}: UseToolsOAuthFlowOptions): UseToolsOAuthFlowResult => {
const [status, setStatus] = useState<ToolsOAuthStatus>("idle");
const [error, setError] = useState<string | null>(null);
const processingRef = useRef(false);
const onSuccessRef = useRef(onSuccess);
onSuccessRef.current = onSuccess;
const startOAuthFlow = useCallback(async () => {
if (typeof window === "undefined") return;
try {
setStatus("authorizing");
setError(null);
let clientId: string | undefined = preClientId ?? undefined;
let clientSecret: string | undefined;
if (!clientId) {
try {
const reg = await registerMcpOAuthClient(accessToken, serverId, {
client_name: serverAlias || serverId,
grant_types: ["authorization_code", "refresh_token"],
response_types: ["code"],
token_endpoint_auth_method: "none",
});
clientId = reg?.client_id;
clientSecret = reg?.client_secret;
} catch (_) {
// Registration is optional; proceed without client_id
}
}
const verifier = generateCodeVerifier();
const challenge = await generateCodeChallenge(verifier);
const state = crypto.randomUUID();
const redirectUri = buildCallbackUrl();
const scopeString = scopes?.filter((s) => s.trim()).join(" ");
const authorizeUrl = buildMcpOAuthAuthorizeUrl({
serverId,
clientId,
redirectUri,
state,
codeChallenge: challenge,
scope: scopeString,
});
const flowState: StoredFlowState = {
state,
codeVerifier: verifier,
serverId,
redirectUri,
clientId,
clientSecret,
scopes,
};
setSecureItem(FLOW_STATE_KEY, JSON.stringify(flowState));
// Return to the current page (Tools tab) after the OAuth redirect
setSecureItem(RETURN_URL_KEY, window.location.href);
window.location.href = authorizeUrl;
} catch (err) {
const msg = extractErrorMessage(err);
setError(msg);
setStatus("error");
NotificationsManager.error(msg);
}
}, [accessToken, serverId, serverAlias, scopes, preClientId]);
const resumeOAuthFlow = useCallback(async () => {
if (typeof window === "undefined" || processingRef.current) return;
const storedResult = getSecureItem(RESULT_KEY);
if (!storedResult) return;
// The callback page writes to this result key for every OAuth flow (including
// the admin server-creation flow). Guard: only proceed if *this* hook's flow
// state exists, meaning startOAuthFlow() was actually called from the Tools screen.
// Without this guard, a stale result written during server creation would trigger
// "OAuth session state was lost" when the user navigates to the Tools tab.
const rawFlowState = getSecureItem(FLOW_STATE_KEY);
if (!rawFlowState) return;
let peeked: StoredFlowState | null = null;
try {
peeked = JSON.parse(rawFlowState) as StoredFlowState;
if (peeked.serverId && peeked.serverId !== serverId) return;
} catch (_) {}
processingRef.current = true;
clearStorage(RESULT_KEY);
let payload: Record<string, unknown> | null = null;
let flowState: StoredFlowState | null = null;
try {
payload = JSON.parse(storedResult);
flowState = peeked;
} catch (_) {
setError("Failed to resume OAuth flow. Please retry.");
setStatus("error");
processingRef.current = false;
clearStorage(FLOW_STATE_KEY);
return;
}
try {
if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) {
throw new Error("OAuth session state was lost. Please retry.");
}
if (!payload?.state || payload.state !== flowState.state) {
throw new Error("OAuth state mismatch. Please retry.");
}
if (payload.error) {
throw new Error((payload.error_description as string) || (payload.error as string));
}
if (!payload.code) {
throw new Error("Authorization code missing in callback.");
}
setStatus("exchanging");
const token = await exchangeMcpOAuthToken({
serverId: flowState.serverId,
code: payload.code as string,
clientId: flowState.clientId,
clientSecret: flowState.clientSecret,
codeVerifier: flowState.codeVerifier,
redirectUri: flowState.redirectUri,
accessToken,
});
// Store in sessionStorage only — no backend DB write
setToken(
flowState.serverId,
{
access_token: token.access_token,
expires_in: token.expires_in,
refresh_token: token.refresh_token,
token_type: token.token_type,
},
userId,
);
setStatus("success");
setError(null);
NotificationsManager.success("Connected successfully");
onSuccessRef.current(token.access_token);
} catch (err) {
const msg = extractErrorMessage(err);
setError(msg);
setStatus("error");
NotificationsManager.error(msg);
} finally {
clearStorage(FLOW_STATE_KEY);
setTimeout(() => { processingRef.current = false; }, 1000);
}
}, [accessToken, serverId, userId]);
useEffect(() => {
resumeOAuthFlow();
}, [resumeOAuthFlow]);
return { startOAuthFlow, status, error };
};

View File

@ -16,15 +16,14 @@ import { useCallback, useEffect, useRef, useState } from "react";
import {
buildMcpOAuthAuthorizeUrl,
exchangeMcpOAuthToken,
getProxyBaseUrl,
registerMcpOAuthClient,
serverRootPath,
storeMCPOAuthUserCredential,
} from "@/components/networking";
import NotificationsManager from "@/components/molecules/notifications_manager";
import { extractErrorMessage } from "@/utils/errorUtils";
import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce";
import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils";
export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
@ -69,26 +68,6 @@ const getStorage = (key: string): string | null => {
return getSecureItem(key);
};
const clearStorage = (...keys: string[]) => {
keys.forEach((k) => {
try {
window.sessionStorage.removeItem(k);
} catch (_) {}
});
};
const buildCallbackUrl = (): string => {
if (typeof window !== "undefined") {
const path = window.location.pathname || "";
const idx = path.indexOf("/ui");
const prefix = idx >= 0 ? path.slice(0, idx + 3).replace(/\/+$/, "") : "";
return `${window.location.origin}${prefix}/mcp/oauth/callback`;
}
const base = (getProxyBaseUrl() || "").replace(/\/+$/, "");
const root = serverRootPath && serverRootPath !== "/" ? serverRootPath : "";
return `${base}${root}/ui/mcp/oauth/callback`;
};
export const useUserMcpOAuthFlow = ({
accessToken,
serverId,
@ -176,13 +155,17 @@ export const useUserMcpOAuthFlow = ({
// mount and would compete for the same RESULT_KEY. Peek at the stored
// flow state first: only the hook instance whose serverId matches the one
// that initiated the OAuth flow should consume the result.
// Guard: only proceed if this hook's flow state exists (startOAuthFlow was
// called from this hook). Without the guard, a tools re-auth redirect writes
// to the user result key too, and every OAuth2ConnectButton instance would try
// to resume a flow that was never started here.
const rawFlowState = getStorage(FLOW_STATE_KEY);
if (rawFlowState) {
try {
const peeked = JSON.parse(rawFlowState) as StoredFlowState;
if (peeked.serverId && peeked.serverId !== serverId) return;
} catch (_) {}
}
if (!rawFlowState) return;
try {
const peeked = JSON.parse(rawFlowState) as StoredFlowState;
if (peeked.serverId && peeked.serverId !== serverId) return;
} catch (_) {}
processingRef.current = true;
clearStorage(RESULT_KEY);

View File

@ -1,5 +1,6 @@
import { describe, it, expect, beforeEach, vi } from "vitest";
import { clearTokenCookies, getCookie, storeLoginToken } from "./cookieUtils";
import { getToken, setToken } from "./mcpTokenStore";
describe("cookieUtils", () => {
beforeEach(() => {
@ -20,6 +21,15 @@ describe("cookieUtils", () => {
expect(getCookie("token")).toBeNull();
});
it("should clear MCP session tokens on logout", () => {
setToken("server-1", { access_token: "mcp-tok" }, "user-a");
expect(getToken("server-1", "user-a")).not.toBeNull();
clearTokenCookies();
expect(getToken("server-1", "user-a")).toBeNull();
});
it("should clear token cookie from /ui path", () => {
document.cookie = "token=test-token-value; path=/ui";
clearTokenCookies();

View File

@ -2,6 +2,8 @@
* Utility functions for managing cookies
*/
import { clearAllMcpTokens } from "./mcpTokenStore";
/**
* Returns the cookie path for the UI.
* Derives the path from window.location.pathname so it works when
@ -67,6 +69,7 @@ export function clearTokenCookies() {
// sessionStorage may be unavailable
}
clearAllMcpTokens();
}
/**

View File

@ -0,0 +1,16 @@
import { describe, expect, it } from "vitest";
import { sanitizeMcpAliasForHeader } from "./mcpHeaderUtils";
describe("sanitizeMcpAliasForHeader", () => {
it("lowercases and replaces spaces with underscores", () => {
expect(sanitizeMcpAliasForHeader("My Server")).toBe("my_server");
});
it("replaces invalid characters for header token segments", () => {
expect(sanitizeMcpAliasForHeader("GitHub-MCP!")).toBe("github_mcp");
});
it("preserves underscores and digits", () => {
expect(sanitizeMcpAliasForHeader("github_mcp2")).toBe("github_mcp2");
});
});

View File

@ -0,0 +1,14 @@
/**
* Sanitize an MCP server alias for use in HTTP header names (x-mcp-{alias}-...).
* RFC 7230 tchar allows token chars; aliases with spaces or hyphens break parsing
* because the backend splits on the first dash after the x-mcp- prefix.
* Keep in sync with litellm.proxy._experimental.mcp_server.utils.sanitize_mcp_alias_for_header.
*/
export function sanitizeMcpAliasForHeader(alias: string): string {
return alias
.toLowerCase()
.trim()
.replace(/[^a-z0-9_]/g, "_")
.replace(/_+/g, "_")
.replace(/^_|_$/g, "");
}

View File

@ -0,0 +1,46 @@
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import {
clearAllMcpTokens,
getToken,
isTokenValid,
removeToken,
setToken,
} from "./mcpTokenStore";
describe("mcpTokenStore", () => {
beforeEach(() => {
sessionStorage.clear();
});
afterEach(() => {
sessionStorage.clear();
});
it("scopes tokens by user id", () => {
setToken("server-a", { access_token: "user1-token" }, "user-1");
setToken("server-a", { access_token: "user2-token" }, "user-2");
expect(getToken("server-a", "user-1")?.access_token).toBe("user1-token");
expect(getToken("server-a", "user-2")?.access_token).toBe("user2-token");
expect(getToken("server-a", "user-3")).toBeNull();
});
it("validates expiry per user scope", () => {
setToken("server-a", { access_token: "tok", expires_in: 3600 }, "user-1");
expect(isTokenValid("server-a", "user-1")).toBe(true);
removeToken("server-a", "user-1");
expect(isTokenValid("server-a", "user-1")).toBe(false);
});
it("clearAllMcpTokens removes every mcp-session-token entry", () => {
setToken("s1", { access_token: "a" }, "u1");
setToken("s2", { access_token: "b" }, "u2");
sessionStorage.setItem("unrelated", "keep");
clearAllMcpTokens();
expect(getToken("s1", "u1")).toBeNull();
expect(getToken("s2", "u2")).toBeNull();
expect(sessionStorage.getItem("unrelated")).toBe("keep");
});
});

View File

@ -0,0 +1,93 @@
/**
* Session-storage-backed OAuth token store for MCP servers.
* Tokens are keyed by LiteLLM user id + server_id and cleared when the browser
* session ends (tab/window close). Never written to localStorage.
*/
const KEY_PREFIX = "mcp-session-token:";
interface StoredToken {
access_token: string;
expires_at: number;
refresh_token?: string;
token_type: string;
}
interface TokenInput {
access_token: string;
expires_in?: number;
refresh_token?: string;
token_type?: string;
}
const DEFAULT_TTL_MS = 3600 * 1000; // 1 hour
function storageKey(serverId: string, userId?: string | null): string {
const userPart = userId?.trim() || "_anonymous";
return `${KEY_PREFIX}${userPart}:${serverId}`;
}
export function setToken(
serverId: string,
data: TokenInput,
userId?: string | null,
): void {
if (typeof window === "undefined") return;
const stored: StoredToken = {
access_token: data.access_token,
expires_at: Date.now() + (data.expires_in != null ? data.expires_in * 1000 : DEFAULT_TTL_MS),
token_type: data.token_type ?? "bearer",
...(data.refresh_token ? { refresh_token: data.refresh_token } : {}),
};
try {
window.sessionStorage.setItem(storageKey(serverId, userId), JSON.stringify(stored));
} catch {
// Silently ignore storage errors (private browsing, quota exceeded, etc.)
}
}
export function getToken(
serverId: string,
userId?: string | null,
): StoredToken | null {
if (typeof window === "undefined") return null;
try {
const raw = window.sessionStorage.getItem(storageKey(serverId, userId));
if (!raw) return null;
return JSON.parse(raw) as StoredToken;
} catch {
return null;
}
}
export function removeToken(serverId: string, userId?: string | null): void {
if (typeof window === "undefined") return;
try {
window.sessionStorage.removeItem(storageKey(serverId, userId));
} catch {
// Silently ignore
}
}
export function isTokenValid(serverId: string, userId?: string | null): boolean {
const token = getToken(serverId, userId);
if (!token) return false;
return token.expires_at > Date.now();
}
/** Remove all MCP session tokens (e.g. on logout or user switch). */
export function clearAllMcpTokens(): void {
if (typeof window === "undefined") return;
try {
const keysToRemove: string[] = [];
for (let i = 0; i < window.sessionStorage.length; i++) {
const key = window.sessionStorage.key(i);
if (key?.startsWith(KEY_PREFIX)) {
keysToRemove.push(key);
}
}
keysToRemove.forEach((key) => window.sessionStorage.removeItem(key));
} catch {
// Silently ignore
}
}