From bafad619ab81bcc0fe8160bb18da65c553a68c48 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 8 Jun 2026 19:58:51 -0700 Subject: [PATCH] fix(mcp): load MCP tool configuration tools via the OBO/passthrough-aware GET path (#29960) * fix(ui): load MCP tool configuration tools via the OBO/passthrough-aware GET path * fix(mcp): admin-only include_disabled_tools so the settings UI shows toggled-off tools * fix(ui): repopulate MCP server edit form when server data loads after mount (OAuth return) * fix(ui): persist MCP OAuth token on save and return to the Settings tab after authorize * fix(ui): scope MCP OAuth callback to the initiating form so create and edit flows don't cross-talk * fix(ui): derive OAuth-return Settings tab via lazy state init instead of setState-in-effect * Fix MCP OAuth edit token handling --------- Co-authored-by: Cursor Agent --- .../mcp_server/rest_endpoints.py | 38 +- .../mcp_server/test_rest_endpoints.py | 197 ++++++++++ .../mcp_tools/create_mcp_server.tsx | 2 +- .../mcp_tools/mcp_server_edit.test.tsx | 345 +++++++++++++++++- .../components/mcp_tools/mcp_server_edit.tsx | 123 ++++++- .../components/mcp_tools/mcp_server_view.tsx | 29 +- .../mcp_tools/mcp_tool_configuration.tsx | 23 +- .../src/components/mcp_tools/mcp_tools.tsx | 13 +- .../src/components/networking.tsx | 9 +- .../src/hooks/useMcpOAuthFlow.tsx | 18 + .../src/utils/mcpHeaderUtils.ts | 15 + 11 files changed, 755 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 924f2cd5dcb..05d7aeb5bab 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -355,8 +355,15 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, extra_headers: Optional[Dict[str, str]] = None, + apply_tool_filters: bool = True, ): - """Helper function to get tools for a single server.""" + """Helper function to get tools for a single server. + + When ``apply_tool_filters`` is False the raw server catalog is returned + without the allowed_tools/disallowed_tools gate or the per-key tool + permissions. This is the admin-only configuration view; every runtime + path keeps the default True so callable tools stay filtered. + """ tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, @@ -366,10 +373,12 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) - # Filter tools based on allowed_tools configuration - # Only filter if allowed_tools is explicitly configured (not None and not empty) - if server.allowed_tools is not None and len(server.allowed_tools) > 0: - tools = filter_tools_by_allowed_tools(tools, server) + if not apply_tool_filters: + return _create_tool_response_objects(tools, server.mcp_info) + + # Always apply allowed_tools/disallowed_tools so the blacklist is + # enforced even when no allowlist is set (matches the SSE/HTTP path). + tools = filter_tools_by_allowed_tools(tools, server) # Filter tools based on user_api_key_auth.object_permission.mcp_tool_permissions # This provides per-key/team/org control over which tools can be accessed @@ -528,6 +537,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str], raw_headers_from_request: dict, user_api_key_dict: UserAPIKeyAuth, + apply_tool_filters: bool = True, ) -> dict: """Handle tool listing for a single server_id request.""" # Resolve a server name to its UUID if needed @@ -592,6 +602,7 @@ if MCP_AVAILABLE: raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) except Exception as e: verbose_logger.exception(f"Error getting tools from {server.name}: {e}") @@ -612,6 +623,14 @@ if MCP_AVAILABLE: server_id: Optional[str] = Query( None, description="The server id to list tools for" ), + include_disabled_tools: bool = Query( + False, + description=( + "Admin only. Return the full server tool catalog without the " + "allowed_tools filter or per-key tool permissions, so the MCP " + "settings UI can configure the allowlist. Ignored for non-admins." + ), + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> dict: """ @@ -639,6 +658,13 @@ if MCP_AVAILABLE: ) try: + # The full catalog (allowlist filter skipped) is admin-only so the + # REST endpoint can't be used to enumerate deliberately-disabled tools. + apply_tool_filters = not ( + include_disabled_tools + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + # Extract auth headers from request headers = request.headers raw_headers_from_request = dict(headers) @@ -680,6 +706,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, raw_headers_from_request=raw_headers_from_request, user_api_key_dict=user_api_key_dict, + apply_tool_filters=apply_tool_filters, ) else: if not allowed_server_ids: @@ -737,6 +764,7 @@ if MCP_AVAILABLE: raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) list_tools_result.extend(tools_result) except Exception as e: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 593facd9279..bc0af933eb5 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -500,6 +500,7 @@ class TestListToolsRestAPI: raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server"] = server @@ -544,6 +545,150 @@ class TestListToolsRestAPI: assert result["error"] is None assert result["message"] == "Successfully retrieved tools" + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): + """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; + a non-admin passing it stays filtered so the REST endpoint can't be used + to enumerate deliberately-disabled tools.""" + from litellm.proxy._types import LitellmUserRoles + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = ["tool1"] + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + + stub_server = StubServer() + captured = {} + + async def fake_get_tools( + server, server_auth_header, *args, apply_tool_filters=True, **kwargs + ): + captured["apply_tool_filters"] = apply_tool_filters + return ["tool-1"] + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert captured["apply_tool_filters"] is False + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(), + ) + assert captured["apply_tool_filters"] is True + + @pytest.mark.parametrize("upstream_status", [401, 403]) + async def test_upstream_auth_failure_surfaces_status_and_challenge( + self, monkeypatch, upstream_status + ): + """A single-server pass-through request whose upstream rejects the token + must surface the upstream status (401 or 403) plus its WWW-Authenticate + challenge, not collapse into a 200 ``unexpected_error`` body.""" + from litellm.proxy._experimental.mcp_server.exceptions import ( + MCPUpstreamAuthError, + ) + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "passthrough" + allowed_tools = None + mcp_info = {"server_name": "passthrough"} + available_on_public_internet = True + + stub_server = StubServer() + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + challenge = 'Bearer resource_metadata="https://upstream/.well-known"' + + async def fake_get_tools(*args, **kwargs): + raise MCPUpstreamAuthError( + status_code=upstream_status, + www_authenticate=challenge, + server_name="passthrough", + ) + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + with pytest.raises(HTTPException) as exc_info: + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert exc_info.value.status_code == upstream_status + assert exc_info.value.headers == {"www-authenticate": challenge} + async def test_name_resolution_finds_server_by_uuid(self, monkeypatch): """When server_id is a name string, it should be resolved to its UUID and used for the tools lookup when the UUID is in allowed_server_ids.""" @@ -576,6 +721,7 @@ class TestListToolsRestAPI: raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server_arg"] = server @@ -719,6 +865,7 @@ class TestListToolsRestAPI: raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["server"] = server captured["auth_header"] = server_auth_header @@ -1211,6 +1358,56 @@ class TestGetToolsForSingleServer: assert "tool1" not in tool_names assert "tool4" not in tool_names + async def test_apply_tool_filters_false_returns_full_catalog(self, monkeypatch): + """apply_tool_filters=False returns the raw catalog without the server + allowed_tools gate, so the config UI can render disabled tools as off.""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + class MockTool: + def __init__(self, name): + self.name = name + self.description = name + self.inputSchema = {} + + mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")] + + async def fake_get_tools_from_server(**kwargs): + return mock_tools + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + fake_get_tools_from_server, + raising=False, + ) + + # Server enforces an allowlist of just tool1. + server = MCPServer( + server_id="test-server-id", + name="test-server", + transport=MCPTransport.sse, + allowed_tools=["tool1"], + ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", object_permission=None) + + # Runtime default: only the allowed tool comes back. + filtered = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + ) + assert [t.name for t in filtered] == ["tool1"] + + # Config view: full catalog, including the disabled tools. + full = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + apply_tool_filters=False, + ) + assert {t.name for t in full} == {"tool1", "tool2", "tool3"} + class TestStdioCommandAllowlist: """Tests for MCP stdio command allowlist validation.""" 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 d69280d6808..47736413a02 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 @@ -184,6 +184,7 @@ const CreateMCPServer: React.FC = ({ } }, onBeforeRedirect: persistCreateUiState, + flowSource: "create", }); React.useEffect(() => { @@ -1061,7 +1062,6 @@ const CreateMCPServer: React.FC = ({
({ updateMCPServer: vi.fn(), listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), + storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}), })); vi.mock("../molecules/notifications_manager", () => ({ @@ -17,12 +18,13 @@ vi.mock("../molecules/notifications_manager", () => ({ }, })); +const mockOauth: { tokenResponse: any } = { tokenResponse: null }; vi.mock("@/hooks/useMcpOAuthFlow", () => ({ useMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null, - tokenResponse: null, + tokenResponse: mockOauth.tokenResponse, }), })); @@ -35,7 +37,50 @@ vi.mock("./MCPPermissionManagement", () => ({ })); vi.mock("./mcp_tool_configuration", () => ({ - default: () =>
, + default: ({ + existingAllowedTools, + externalTools, + externalError, + onAllowedToolsChange, + onToolAllowlistInteraction, + onToolNameToDisplayNameChange, + onToolNameToDescriptionChange, + }: any) => ( +
+ + +
+ ), +})); + +const mockGetToken = vi.fn(); +const mockIsTokenValid = vi.fn(); +const mockSetToken = vi.fn(); +vi.mock("@/utils/mcpTokenStore", () => ({ + getToken: (...args: any[]) => mockGetToken(...args), + isTokenValid: (...args: any[]) => mockIsTokenValid(...args), + setToken: (...args: any[]) => mockSetToken(...args), })); // ── fixtures ────────────────────────────────────────────────────────────────── @@ -442,3 +487,299 @@ describe("MCPServerEdit (interactive OAuth)", () => { expect(payload.token_storage_ttl_seconds).toBe(7200); }); }); + +describe("MCPServerEdit (tool list fetch)", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: [], error: null }); + mockOauth.tokenResponse = null; + }); + + it("loads an OBO server's tools via GET listMCPTools with no passthrough headers", async () => { + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: [{ name: "read_user" }], + error: null, + }); + + render( + , + ); + + await waitFor(() => { + // includeDisabledTools=true so the config screen gets the full catalog. + expect(networking.listMCPTools).toHaveBeenCalledWith("access-token", "oauth_server_1", undefined, true); + }); + // OBO uses the backend-stored token; the browser passthrough store is never consulted. + expect(mockIsTokenValid).not.toHaveBeenCalled(); + + await waitFor(() => { + expect(screen.getByTestId("mcp-tool-config")).toHaveAttribute( + "data-external-tools", + JSON.stringify([{ name: "read_user" }]), + ); + }); + }); + + it("forwards the sessionStorage token as the x-mcp passthrough header for a passthrough server", async () => { + mockIsTokenValid.mockReturnValue(true); + mockGetToken.mockReturnValue({ access_token: "browser-token" }); + + render( + , + ); + + await waitFor(() => { + expect(networking.listMCPTools).toHaveBeenCalledWith( + "access-token", + "oauth_server_1", + { "x-mcp-oauth_server-authorization": "Bearer browser-token" }, + true, + ); + }); + expect(mockGetToken).toHaveBeenCalledWith("oauth_server_1", "user-1"); + }); + + it("uses the staged OAuth token to load passthrough tools after authorize", async () => { + const passthroughServer = { ...interactiveOAuthServer, delegate_auth_to_upstream: true }; + mockIsTokenValid.mockReturnValue(false); + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: [{ name: "read_user" }], + error: null, + }); + + const props = { + mcpServer: passthroughServer, + accessToken: "access-token", + userID: "user-1", + onCancel: vi.fn(), + onSuccess: vi.fn(), + availableAccessGroups: [], + }; + + const { rerender } = render(); + + await waitFor(() => { + expect(screen.getByTestId("mcp-tool-config").getAttribute("data-external-error")).toContain( + "Authenticate with this server in the Tools tab", + ); + }); + expect(networking.listMCPTools).not.toHaveBeenCalled(); + + mockOauth.tokenResponse = { access_token: "staged-token", expires_in: 1800 }; + rerender(); + + await waitFor(() => { + expect(networking.listMCPTools).toHaveBeenCalledWith( + "access-token", + "oauth_server_1", + { "x-mcp-oauth_server-authorization": "Bearer staged-token" }, + true, + ); + }); + expect(mockGetToken).not.toHaveBeenCalled(); + }); + + it("prompts to authenticate and does not fetch when a passthrough server has no session token", async () => { + mockIsTokenValid.mockReturnValue(false); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByTestId("mcp-tool-config").getAttribute("data-external-error")).toContain( + "Authenticate with this server in the Tools tab", + ); + }); + expect(networking.listMCPTools).not.toHaveBeenCalled(); + }); +}); + +describe("MCPServerEdit (form resync)", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: [], error: null }); + }); + + it("repopulates the form when the server data arrives after mount", async () => { + const props = { + accessToken: "access-token", + onCancel: vi.fn(), + onSuccess: vi.fn(), + availableAccessGroups: [], + }; + + // Mount before the server is loaded (mirrors landing on the page mid OAuth return). + const { rerender } = render(); + expect(screen.queryByDisplayValue("https://example.com/mcp")).not.toBeInTheDocument(); + + // Server data arrives; the form must repopulate rather than staying blank. + rerender(); + + await waitFor(() => { + expect(screen.getByDisplayValue("https://example.com/mcp")).toBeInTheDocument(); + }); + }); +}); + +describe("MCPServerEdit (OAuth token persistence on save)", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(networking.listMCPTools).mockResolvedValue({ tools: [], error: null }); + mockOauth.tokenResponse = null; + }); + + it("persists the OBO token to the DB on save after authorize", async () => { + mockOauth.tokenResponse = { + access_token: "obo-tok", + refresh_token: "obo-refresh", + expires_in: 3600, + scope: "read write", + }; + vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer }); + + render( + , + ); + + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]); + }); + + await waitFor(() => { + expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalledWith( + "access-token", + "oauth_server_1", + expect.objectContaining({ + access_token: "obo-tok", + refresh_token: "obo-refresh", + expires_in: 3600, + scopes: ["read", "write"], + }), + ); + }); + expect(mockSetToken).not.toHaveBeenCalled(); + }); + + it("does not show success when OBO token persistence fails after update", async () => { + mockOauth.tokenResponse = { + access_token: "obo-tok", + refresh_token: "obo-refresh", + expires_in: 3600, + scope: "read write", + }; + vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer }); + vi.mocked(networking.storeMCPOAuthUserCredential).mockRejectedValueOnce(new Error("write failed")); + const onSuccess = vi.fn(); + + render( + , + ); + + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]); + }); + + await waitFor(() => { + expect(networking.storeMCPOAuthUserCredential).toHaveBeenCalled(); + }); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + "MCP Server updated, but failed to persist OAuth token: write failed", + ); + expect(NotificationsManager.success).not.toHaveBeenCalledWith("MCP Server updated successfully"); + expect(onSuccess).not.toHaveBeenCalled(); + }); + + it("persists the passthrough token to sessionStorage on save after authorize", async () => { + mockOauth.tokenResponse = { access_token: "pt-tok", expires_in: 1800, token_type: "bearer" }; + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + delegate_auth_to_upstream: true, + }); + + render( + , + ); + + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]); + }); + + await waitFor(() => { + expect(mockSetToken).toHaveBeenCalledWith( + "oauth_server_1", + expect.objectContaining({ access_token: "pt-tok", expires_in: 1800, token_type: "bearer" }), + "user-1", + ); + }); + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + }); + + it("persists nothing on save when no token was fetched", async () => { + mockOauth.tokenResponse = null; + vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...interactiveOAuthServer }); + + render( + , + ); + + await act(async () => { + fireEvent.click(screen.getAllByRole("button", { name: "Save Changes" })[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalled(); + }); + expect(networking.storeMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(mockSetToken).not.toHaveBeenCalled(); + }); +}); 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 9278d41c3e3..0d8f312622a 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,8 +2,18 @@ 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, listMCPTools } from "../networking"; +import { + AUTH_TYPE, + OAUTH_FLOW, + MCP_OAUTH2_FLOW_M2M, + MCPServer, + MCPServerCostInfo, + TRANSPORT, + getMcpOAuthMode, +} from "./types"; +import { updateMCPServer, listMCPTools, storeMCPOAuthUserCredential } from "../networking"; +import { getToken, isTokenValid, setToken } from "@/utils/mcpTokenStore"; +import { buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils"; import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPPermissionManagement from "./MCPPermissionManagement"; import MCPToolConfiguration from "./mcp_tool_configuration"; @@ -17,6 +27,7 @@ import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; interface MCPServerEditProps { mcpServer: MCPServer; accessToken: string | null; + userID?: string | null; onCancel: () => void; onSuccess: (server: MCPServer) => void; availableAccessGroups: string[]; @@ -24,11 +35,12 @@ interface MCPServerEditProps { const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC]; const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2, AUTH_TYPE.AWS_SIGV4]; -const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; +export const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; const MCPServerEdit: React.FC = ({ mcpServer, accessToken, + userID, onCancel, onSuccess, availableAccessGroups, @@ -56,8 +68,6 @@ const MCPServerEdit: React.FC = ({ const oauthFlowTypeValue = Form.useWatch("oauth_flow_type", form) as string | undefined; const isM2MFlow = isOAuthAuthType && oauthFlowTypeValue === OAUTH_FLOW.M2M; - const [oauthAccessToken, setOauthAccessToken] = useState(null); - // Watch form fields that affect tool fetching const currentUrl = Form.useWatch("url", form); const currentSpecPath = Form.useWatch("spec_path", form); @@ -134,8 +144,6 @@ const MCPServerEdit: React.FC = ({ }; }, onTokenReceived: (token) => { - setOauthAccessToken(token?.access_token ?? null); - if (token?.access_token) { const credentials = { access_token: token.access_token, @@ -152,6 +160,7 @@ const MCPServerEdit: React.FC = ({ } }, onBeforeRedirect: persistEditUiState, + flowSource: "edit", }); const initialStaticHeaders = React.useMemo(() => { @@ -199,6 +208,20 @@ const MCPServerEdit: React.FC = ({ [mcpServer, effectiveTransport, initialStaticHeaders, initialEnvJson], ); + // antd applies `initialValues` only at first mount. When the server loads after + // mount (e.g. returning from the OAuth redirect lands on Overview and the form + // mounts before the server data is ready), the form would stay blank. Re-sync it + // from the loaded server once per server_id so it always reflects the saved config; + // the OAuth-restore effect below then overlays any in-progress edits on top. + const syncedServerIdRef = React.useRef(null); + useEffect(() => { + if (!mcpServer.server_id || syncedServerIdRef.current === mcpServer.server_id) { + return; + } + syncedServerIdRef.current = mcpServer.server_id; + form.setFieldsValue(initialValues); + }, [mcpServer.server_id, initialValues, form]); + // Initialize cost config from existing server data useEffect(() => { if (mcpServer.mcp_info?.mcp_server_cost_info) { @@ -255,6 +278,9 @@ const MCPServerEdit: React.FC = ({ if (!pendingRestoredValues) { return; } + // Set transport first so transport-dependent fields render, then apply the rest + // on the re-run triggered by the transportType watch (without it the effect's + // deps never change and the second pass never runs, leaving fields blank). const transport = pendingRestoredValues.transport || mcpServer.transport; if (transport && transport !== form.getFieldValue("transport")) { form.setFieldsValue({ transport }); @@ -262,7 +288,7 @@ const MCPServerEdit: React.FC = ({ } form.setFieldsValue(pendingRestoredValues); setPendingRestoredValues(null); - }, [pendingRestoredValues, form, mcpServer.transport]); + }, [pendingRestoredValues, form, mcpServer.transport, transportType]); // Transform string array to object array for initial form values useEffect(() => { @@ -279,28 +305,52 @@ const MCPServerEdit: React.FC = ({ return; } fetchTools(); - }, [mcpServer, accessToken]); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [mcpServer, accessToken, userID, oauthTokenResponse?.access_token]); const fetchTools = async () => { if (!accessToken || !mcpServer.server_id) return; + // OBO/M2M/static auth is attached server-side from the stored credential, so + // a plain GET /tools/list?server_id suffices. PKCE passthrough holds the token + // in the browser, so forward it from sessionStorage as the x-mcp header the + // same way the Tools playground does. + let customHeaders: Record | undefined; + const isPassthrough = + getMcpOAuthMode({ + auth_type: mcpServer.auth_type, + oauth2_flow: mcpServer.oauth2_flow, + delegate_auth_to_upstream: mcpServer.delegate_auth_to_upstream, + }) === "passthrough"; + if (isPassthrough) { + const token = + oauthTokenResponse?.access_token ?? + (isTokenValid(mcpServer.server_id, userID) + ? getToken(mcpServer.server_id, userID)?.access_token ?? null + : null); + if (!token) { + setTools([]); + setToolsError("Authenticate with this server in the Tools tab to load and configure its tools."); + return; + } + customHeaders = buildMcpPassthroughAuthHeader(mcpServer.alias, token); + } + setIsLoadingTools(true); setToolsError(null); try { - // Use the GET endpoint which looks up stored credentials by server_id, - // rather than POST /test/tools/list which requires inline credentials. - const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id); + // include_disabled_tools: configuring the allowlist needs the full server + // catalog, so tools toggled off still render (as unchecked) instead of vanishing. + const toolsResponse = await listMCPTools(accessToken, mcpServer.server_id, customHeaders, true); if (toolsResponse.tools && !toolsResponse.error) { setTools(toolsResponse.tools); } else { - console.error("Failed to fetch tools:", toolsResponse.message); setTools([]); setToolsError(toolsResponse.message || "Failed to load tools"); } } catch (error) { - console.error("Tools fetch error:", error); setTools([]); setToolsError(error instanceof Error ? error.message : "Failed to load tools"); } finally { @@ -575,6 +625,46 @@ const MCPServerEdit: React.FC = ({ } const updated = await updateMCPServer(accessToken, payload); + + // Persist the token staged via "Authorize & Fetch" (mirrors the create flow's + // commit-on-submit): OBO writes the per-user token to the DB, passthrough keeps + // it in sessionStorage. M2M/static auth resolve server-side and need neither. + if (oauthTokenResponse?.access_token) { + const oauthMode = getMcpOAuthMode({ + auth_type: restValues.auth_type, + oauth2_flow: isM2MFlow ? MCP_OAUTH2_FLOW_M2M : null, + delegate_auth_to_upstream: Boolean(delegateAuthToUpstreamRaw ?? mcpServer.delegate_auth_to_upstream), + }); + try { + if (oauthMode === "obo") { + const scope = oauthTokenResponse.scope; + await storeMCPOAuthUserCredential(accessToken, mcpServer.server_id, { + access_token: oauthTokenResponse.access_token, + refresh_token: oauthTokenResponse.refresh_token, + expires_in: oauthTokenResponse.expires_in, + scopes: typeof scope === "string" && scope ? scope.split(" ") : undefined, + }); + } else if (oauthMode === "passthrough") { + setToken( + mcpServer.server_id, + { + access_token: oauthTokenResponse.access_token, + expires_in: oauthTokenResponse.expires_in, + refresh_token: oauthTokenResponse.refresh_token, + token_type: oauthTokenResponse.token_type, + }, + userID, + ); + } + } catch (error: unknown) { + const message = error instanceof Error ? error.message : ""; + NotificationsManager.fromBackend( + "MCP Server updated, but failed to persist OAuth token" + (message ? `: ${message}` : ""), + ); + return; + } + } + NotificationsManager.success("MCP Server updated successfully"); onSuccess(updated); } catch (error: any) { @@ -1089,7 +1179,6 @@ const MCPServerEdit: React.FC = ({
= ({ toolNameToDescription={toolNameToDescription} onToolNameToDisplayNameChange={setToolNameToDisplayName} onToolNameToDescriptionChange={setToolNameToDescription} + externalTools={tools} + externalIsLoading={isLoadingTools} + externalError={toolsError} + externalCanFetch={true} />
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 0d9d562f794..c10f25022d5 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 @@ -5,7 +5,8 @@ import { Title, Card, Button, Text, Grid, TabGroup, TabList, TabPanel, TabPanels import { MCPServer, handleTransport, handleAuth } from "./types"; // TODO: Move Tools viewer from index file import { MCPToolsViewer } from "."; -import MCPServerEdit from "./mcp_server_edit"; +import MCPServerEdit, { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit"; +import { getSecureItem } from "@/utils/secureStorage"; import MCPServerCostDisplay from "./mcp_server_cost_display"; import { getMaskedAndFullUrl } from "./utils"; import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; @@ -24,6 +25,24 @@ interface MCPServerViewProps { initialTabIndex?: number; } +// True when this render is the return from the edit-settings OAuth redirect for this +// server: the edit form wrote its UI-state snapshot before redirecting. Used to open +// the editing Settings tab on first render instead of defaulting to Overview. +function isReturningFromEditOAuth(isProxyAdmin: boolean, serverId: string): boolean { + if (typeof window === "undefined" || !isProxyAdmin) { + return false; + } + const stored = getSecureItem(EDIT_OAUTH_UI_STATE_KEY); + if (!stored) { + return false; + } + try { + return JSON.parse(stored)?.serverId === serverId; + } catch { + return false; + } +} + export const MCPServerView: React.FC = ({ mcpServer, onBack, @@ -35,10 +54,13 @@ export const MCPServerView: React.FC = ({ availableAccessGroups, initialTabIndex = 0, }) => { - const [editing, setEditing] = useState(isEditing); + // Open the editing Settings tab on first render when returning from the edit OAuth + // redirect, so the "token fetched" feedback shows where the user left off (Settings=2). + const returningFromEditOAuth = isReturningFromEditOAuth(isProxyAdmin, mcpServer.server_id); + const [editing, setEditing] = useState(isEditing || returningFromEditOAuth); const [showFullUrl, setShowFullUrl] = useState(false); const [copiedStates, setCopiedStates] = useState>({}); - const [selectedTabIndex, setSelectedTabIndex] = useState(initialTabIndex); + const [selectedTabIndex, setSelectedTabIndex] = useState(returningFromEditOAuth ? 2 : initialTabIndex); const handleSuccess = (updated: MCPServer) => { setEditing(false); @@ -201,6 +223,7 @@ export const MCPServerView: React.FC = ({ setEditing(false)} onSuccess={handleSuccess} availableAccessGroups={availableAccessGroups} 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..d4e36f31261 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 @@ -2,7 +2,6 @@ import React, { useEffect, useMemo, useRef, useState } from "react"; import { Card, Title, Text } from "@tremor/react"; import { ToolOutlined, CheckCircleOutlined, SearchOutlined, EditOutlined } from "@ant-design/icons"; import { Badge, Spin, Checkbox, Input, Radio } from "antd"; -import { useTestMCPConnection } from "../../hooks/useTestMCPConnection"; import McpCrudPermissionPanel from "./McpCrudPermissionPanel"; interface KeyTool { @@ -12,7 +11,6 @@ interface KeyTool { interface MCPToolConfigurationProps { accessToken: string | null; - oauthAccessToken?: string | null; formValues: Record; allowedTools: string[]; existingAllowedTools: string[] | null; @@ -144,7 +142,6 @@ const ToolRow: React.FC = ({ const MCPToolConfiguration: React.FC = ({ accessToken, - oauthAccessToken, formValues, allowedTools, existingAllowedTools, @@ -166,19 +163,13 @@ const MCPToolConfiguration: React.FC = ({ const previousSuggestedToolNamesRef = useRef(""); const [expandedTools, setExpandedTools] = useState>(new Set()); - // Use external tool state when provided (avoids duplicate fetch with MCPConnectionStatus). - // Fall back to internal hook when used standalone (e.g., edit flow). - const hasExternalState = externalTools !== undefined; - const internalHook = useTestMCPConnection({ - accessToken, - oauthAccessToken, - formValues, - enabled: !hasExternalState, - }); - const tools: ToolEntry[] = hasExternalState ? externalTools : internalHook.tools; - const isLoadingTools = hasExternalState ? (externalIsLoading ?? false) : internalHook.isLoadingTools; - const toolsError = hasExternalState ? (externalError ?? null) : internalHook.toolsError; - const canFetchTools = hasExternalState ? (externalCanFetch ?? false) : internalHook.canFetchTools; + // Tool list is fetched by the parent (create/edit flow) and passed in. This + // component renders that state; it never fetches on its own, so there is a + // single source of truth and no risk of falling back to a different endpoint. + const tools: ToolEntry[] = externalTools ?? []; + const isLoadingTools = externalIsLoading ?? false; + const toolsError = externalError ?? null; + const canFetchTools = externalCanFetch ?? false; // Fuzzy-match curated key tool names against actual loaded tool names const suggestedTools = useMemo(() => { diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx index 8957c7b30b8..d94dfbb966f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx @@ -4,7 +4,7 @@ import { ToolTestPanel } from "./ToolTestPanel"; import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse, getMcpOAuthMode } from "./types"; import { listMCPTools, callMCPTool, getMCPOAuthUserCredentialStatus } from "../networking"; import { isTokenValid, getToken, removeToken } from "@/utils/mcpTokenStore"; -import { sanitizeMcpAliasForHeader } from "@/utils/mcpHeaderUtils"; +import { sanitizeMcpAliasForHeader, buildMcpPassthroughAuthHeader } from "@/utils/mcpHeaderUtils"; import { useToolsOAuthFlow } from "@/hooks/useToolsOAuthFlow"; import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; @@ -106,16 +106,7 @@ const MCPToolsViewer = ({ // When no alias is available, fall back to x-mcp-auth (legacy but still supported). // Passthrough only: OBO/M2M tokens are attached server-side, not from the browser. if (isPassthrough && 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}`; - } + Object.assign(customHeaders, buildMcpPassthroughAuthHeader(serverAlias, oauthToken)); } // Add passthrough headers with server-specific prefix diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index aa5e08d195f..774a946b2a4 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -7144,11 +7144,12 @@ export const listMCPTools = async ( accessToken: string, serverId: string, customHeaders?: Record, + includeDisabledTools?: boolean, ) => { - // Construct base URL - let url = proxyBaseUrl - ? `${proxyBaseUrl}/mcp-rest/tools/list?server_id=${serverId}` - : `/mcp-rest/tools/list?server_id=${serverId}`; + // Construct base URL. include_disabled_tools returns the full server catalog + // (admin-only, backend-enforced) so the settings UI can configure the allowlist. + const query = `server_id=${serverId}${includeDisabledTools ? "&include_disabled_tools=true" : ""}`; + let url = proxyBaseUrl ? `${proxyBaseUrl}/mcp-rest/tools/list?${query}` : `/mcp-rest/tools/list?${query}`; console.log("Fetching MCP tools from:", url); diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index 11efaba53b3..4c94db47ea8 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -26,6 +26,11 @@ interface UseMcpOAuthFlowOptions { getTemporaryPayload: () => Record | null; onTokenReceived: (tokenResponse: Record) => void; onBeforeRedirect?: () => void; + // Distinguishes which form started the flow (e.g. "create" vs "edit"). Both forms + // mount this hook with shared storage keys, so the return handler only processes a + // callback whose stored flowSource matches, preventing one form from grabbing the + // other's OAuth result. + flowSource: string; } interface UseMcpOAuthFlowResult { @@ -41,6 +46,7 @@ export const useMcpOAuthFlow = ({ getTemporaryPayload, onTokenReceived, onBeforeRedirect, + flowSource, }: UseMcpOAuthFlowOptions): UseMcpOAuthFlowResult => { const [status, setStatus] = useState("idle"); const [error, setError] = useState(null); @@ -58,6 +64,7 @@ export const useMcpOAuthFlow = ({ clientSecret?: string; serverId: string; redirectUri: string; + flowSource?: string; }; const setStorageItem = (key: string, value: string) => { @@ -175,6 +182,7 @@ export const useMcpOAuthFlow = ({ clientSecret: registeredClient.clientSecret || credentials.client_secret, serverId, redirectUri: callbackUrl(), + flowSource, }; if (typeof window === "undefined") { @@ -253,6 +261,16 @@ export const useMcpOAuthFlow = ({ return; } + // Only the form that started this redirect should consume the result. The create + // form and the edit form both mount this hook with shared storage keys, so without + // this another instance (e.g. the always-mounted create form) would grab and handle + // an edit-page authorization. Bail out without clearing RESULT_KEY so the matching + // instance can still process it. + if (flowState?.flowSource !== flowSource) { + processingRef.current = false; + return; + } + // Clear the result key after reading it if (typeof window !== "undefined") { try { diff --git a/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts b/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts index 76c752a4b61..ebe7116c18f 100644 --- a/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts +++ b/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts @@ -12,3 +12,18 @@ export function sanitizeMcpAliasForHeader(alias: string): string { .replace(/_+/g, "_") .replace(/^_|_$/g, ""); } + +/** + * Build the passthrough auth header that forwards a browser-held PKCE token to + * the upstream MCP server. The backend picks up the x-mcp-{alias}-authorization + * pattern and forwards it; without a usable alias it falls back to the legacy + * x-mcp-auth header. + */ +export function buildMcpPassthroughAuthHeader( + serverAlias: string | null | undefined, + token: string, +): Record { + const safeAlias = serverAlias ? sanitizeMcpAliasForHeader(serverAlias) : ""; + const headerName = safeAlias ? `x-mcp-${safeAlias}-authorization` : "x-mcp-auth"; + return { [headerName]: `Bearer ${token}` }; +}