diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokFields.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokFields.tsx new file mode 100644 index 00000000000..ee793c62ddc --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokFields.tsx @@ -0,0 +1,89 @@ +import React from "react"; +import { Form, Select, Input, Switch, Tooltip } from "antd"; +import type { FormInstance } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { AUTH_TYPE } from "./types"; + +export const BYOK_AUTH_FORMAT_HINT: Record = { + [AUTH_TYPE.BEARER_TOKEN]: "Authorization: Bearer {key}", + [AUTH_TYPE.TOKEN]: "Authorization: token {key}", + [AUTH_TYPE.API_KEY]: "x-api-key: {key}", + [AUTH_TYPE.BASIC]: "Authorization: Basic {key}", +}; + +interface ByokFieldsProps { + form: FormInstance; +} + +const ByokFields: React.FC = ({ form }) => { + const isByok = Form.useWatch("is_byok", form); + const authType = Form.useWatch("auth_type", form) as string | undefined; + const formatHint = authType ? BYOK_AUTH_FORMAT_HINT[authType] : undefined; + + return ( + <> + + BYOK (Bring Your Own Key) + + + + + } + name="is_byok" + valuePropName="checked" + > + + + + {isByok && ( + <> + {formatHint && ( +
+ + + User keys will be sent as: {formatHint} + +
+ )} + + + Access Description + + + + + } + name="byok_description" + > + + + + )} + + ); +}; + +export default ByokFields; 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 e70548d6a96..70f6c3a3a84 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 @@ -1,4 +1,4 @@ -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import * as networking from "../networking"; @@ -773,3 +773,127 @@ describe("CreateMCPServer", () => { }); }); }); + +const BYOK_LABEL = "BYOK (Bring Your Own Key)"; + +function getByokSwitch(): HTMLButtonElement | null { + const label = screen.queryByText(BYOK_LABEL); + const formItem = label?.closest(".ant-form-item"); + return (formItem?.querySelector('button[role="switch"]') as HTMLButtonElement) ?? null; +} + +describe("CreateMCPServer BYOK toggle", () => { + beforeEach(() => { + vi.clearAllMocks(); + oauthHook.tokenResponse = null; + oauthHook.onTokenReceived = null; + }); + + async function selectHttpTransport() { + render(); + await selectAntOption("Transport Type", "Streamable HTTP"); + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + } + + it("shows the BYOK toggle for each static-credential auth type", async () => { + for (const authLabel of ["API Key", "Bearer Token", "Token", "Basic Auth"]) { + await selectHttpTransport(); + await selectAntOption("Authentication", authLabel); + await waitFor(() => { + expect(screen.getByText(BYOK_LABEL)).toBeInTheDocument(); + }); + cleanup(); + } + }); + + it("hides the BYOK toggle for non-static auth types (none, oauth, sigv4)", async () => { + for (const authLabel of ["None", "OAuth", "AWS SigV4 (Bedrock AgentCore MCPs)"]) { + await selectHttpTransport(); + await selectAntOption("Authentication", authLabel); + await waitFor(() => { + expect(screen.queryByText(BYOK_LABEL)).not.toBeInTheDocument(); + }); + cleanup(); + } + }); + + it("swaps the shared Authentication Value field for per-user fields when BYOK is on", async () => { + await selectHttpTransport(); + await selectAntOption("Authentication", "Bearer Token"); + + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); + + const byokSwitch = getByokSwitch(); + expect(byokSwitch).toBeTruthy(); + await act(async () => { + fireEvent.click(byokSwitch!); + }); + + await waitFor(() => { + expect(screen.queryByText("Authentication Value")).not.toBeInTheDocument(); + expect(screen.getByText("Access Description")).toBeInTheDocument(); + }); + }); + + it("sends is_byok=true and omits the shared credential when BYOK is enabled", async () => { + await selectHttpTransport(); + + const user = userEvent.setup({ delay: null }); + await user.type(getServerNameInput(), "Byok_Server"); + await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp"); + + await selectAntOption("Authentication", "Bearer Token"); + await waitFor(() => expect(screen.getByText(BYOK_LABEL)).toBeInTheDocument()); + + await act(async () => { + fireEvent.click(getByokSwitch()!); + }); + await waitFor(() => expect(screen.queryByText("Authentication Value")).not.toBeInTheDocument()); + + vi.mocked(networking.createMCPServer).mockResolvedValue({ server_id: "byok-1" } as any); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: "Add MCP Server" })); + }); + + await waitFor(() => expect(networking.createMCPServer).toHaveBeenCalledTimes(1)); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.is_byok).toBe(true); + expect(payload.auth_type).toBe("bearer_token"); + expect(payload.credentials).toBeUndefined(); + }); + + it("forces is_byok=false when auth_type is switched away from a static type after enabling it", async () => { + await selectHttpTransport(); + + const user = userEvent.setup({ delay: null }); + await user.type(getServerNameInput(), "Switched_Server"); + await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp"); + + await selectAntOption("Authentication", "Bearer Token"); + await waitFor(() => expect(screen.getByText(BYOK_LABEL)).toBeInTheDocument()); + await act(async () => { + fireEvent.click(getByokSwitch()!); + }); + + // Switching to OAuth removes BYOK eligibility; a stale is_byok=true must not persist. + await selectAntOption("Authentication", "None"); + await waitFor(() => expect(screen.queryByText(BYOK_LABEL)).not.toBeInTheDocument()); + + vi.mocked(networking.createMCPServer).mockResolvedValue({ server_id: "switched-1" } as any); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: "Add MCP Server" })); + }); + + await waitFor(() => expect(networking.createMCPServer).toHaveBeenCalledTimes(1)); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.is_byok).toBe(false); + }); +}); 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 ddcc9f65d38..8adfa376853 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 @@ -1,5 +1,5 @@ import React, { useState } from "react"; -import { Modal, Tooltip, Form, Select, Input, Switch, Collapse } from "antd"; +import { Modal, Tooltip, Form, Select, Input, Collapse } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { createMCPServer, registerMCPServer, storeMCPOAuthUserCredential } from "../networking"; @@ -15,6 +15,7 @@ import { MCP_OAUTH2_FLOW_M2M, } from "./types"; import OAuthFormFields from "./OAuthFormFields"; +import ByokFields from "./ByokFields"; import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPConnectionStatus from "./mcp_connection_status"; import MCPToolConfiguration from "./mcp_tool_configuration"; @@ -100,6 +101,7 @@ const CreateMCPServer: React.FC = ({ const authType = formValues.auth_type as string | undefined; const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false; + const isByok = shouldShowAuthValueField && Boolean(formValues.is_byok); const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2; const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4; const isM2MFlow = isOAuthAuthType && formValues.oauth_flow_type === OAUTH_FLOW.M2M; @@ -382,6 +384,14 @@ const CreateMCPServer: React.FC = ({ restValues.transport = "http"; } + // BYOK is only valid for static-credential auth types; drop a stale ``true`` + // (and its metadata) left behind by switching auth_type after toggling it on. + const isByokEnabled = + AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(restValues.auth_type) && Boolean(restValues.is_byok); + restValues.is_byok = isByokEnabled; + restValues.byok_description = isByokEnabled ? restValues.byok_description || [] : []; + restValues.byok_api_key_help_url = isByokEnabled ? restValues.byok_api_key_help_url || null : null; + // Parse token_validation JSON if provided let tokenValidation: Record | null = null; if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") { @@ -422,7 +432,7 @@ const CreateMCPServer: React.FC = ({ }; const includeCredentials = - restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); + !isByokEnabled && restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); if (includeCredentials && credentialsPayload && Object.keys(credentialsPayload).length > 0) { payload.credentials = credentialsPayload; @@ -742,96 +752,6 @@ const CreateMCPServer: React.FC = ({ /> )} - {/* BYOK toggle - only for OpenAPI */} - {transportType === TRANSPORT.OPENAPI && ( - <> - - BYOK (Bring Your Own Key) - - - - - } - name="is_byok" - valuePropName="checked" - > - - - - prev.is_byok !== cur.is_byok || prev.auth_type !== cur.auth_type} - > - {({ getFieldValue }) => - getFieldValue("is_byok") ? ( - <> - {/* Auth format hint */} - {getFieldValue("auth_type") && getFieldValue("auth_type") !== "none" && ( -
- - - User keys will be sent as:{" "} - - {getFieldValue("auth_type") === "bearer_token" && "Authorization: Bearer {key}"} - {getFieldValue("auth_type") === "token" && "Authorization: token {key}"} - {getFieldValue("auth_type") === "api_key" && "x-api-key: {key}"} - {getFieldValue("auth_type") === "basic" && "Authorization: Basic {key}"} - {getFieldValue("auth_type") === "authorization" && "Authorization: {key}"} - - {!getFieldValue("auth_type") && "Set Authentication Type below to specify the format."} - -
- )} - {!getFieldValue("auth_type") && ( -
- - - Set the Authentication Type below to specify how user keys are sent - (e.g., Bearer Token, API Key header). - -
- )} - - Access Description - - - - - } - name="byok_description" - > - - - - ) : null - } - - - )} - {/* Authentication - show for HTTP, SSE, and OpenAPI */} {transportType !== "stdio" && transportType !== "" && ( = ({ - {shouldShowAuthValueField && ( + {shouldShowAuthValueField && } + + {shouldShowAuthValueField && !isByok && ( diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index d7c1241b044..a6e142bf103 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -940,3 +940,101 @@ describe("MCPServerEdit (OAuth token persistence on save)", () => { expect(mockSetToken).not.toHaveBeenCalled(); }); }); + +describe("MCPServerEdit (BYOK)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + const byokServer = { + server_id: "byok_server_1", + server_name: "ByokServer", + alias: "byok_server", + description: "BYOK MCP server", + transport: "http", + url: "https://example.com/mcp", + auth_type: "bearer_token", + is_byok: true, + byok_description: ["Create and manage issues"], + byok_api_key_help_url: "https://docs.example.com/keys", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + mcp_access_groups: [], + }; + + const getByokSwitch = () => { + const formItem = screen.queryByText("BYOK (Bring Your Own Key)")?.closest(".ant-form-item"); + return (formItem?.querySelector('button[role="switch"]') as HTMLButtonElement) ?? null; + }; + + const save = async () => { + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + await waitFor(() => expect(networking.updateMCPServer).toHaveBeenCalledTimes(1)); + return vi.mocked(networking.updateMCPServer).mock.calls[0][1]; + }; + + it("renders the BYOK toggle on for an existing BYOK server and hides the shared auth value field", async () => { + render( + , + ); + + await waitFor(() => expect(getByokSwitch()).toBeTruthy()); + expect(getByokSwitch()).toHaveAttribute("aria-checked", "true"); + expect(screen.queryByText("Authentication Value")).not.toBeInTheDocument(); + }); + + it("persists is_byok=true when saving an existing BYOK server", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue(byokServer as any); + + render( + , + ); + + await waitFor(() => expect(getByokSwitch()).toBeTruthy()); + + const payload = await save(); + expect(payload.is_byok).toBe(true); + expect(payload.credentials).toBeUndefined(); + }); + + it("persists is_byok=false and restores the shared auth value field when BYOK is toggled off", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue({ ...byokServer, is_byok: false } as any); + + render( + , + ); + + await waitFor(() => expect(getByokSwitch()).toBeTruthy()); + await act(async () => { + fireEvent.click(getByokSwitch()!); + }); + + await waitFor(() => expect(screen.getByText("Authentication Value")).toBeInTheDocument()); + + const payload = await save(); + expect(payload.is_byok).toBe(false); + }); +}); 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 b992313d4a8..dbc788243f1 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 @@ -18,6 +18,7 @@ import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPPermissionManagement from "./MCPPermissionManagement"; import MCPToolConfiguration from "./mcp_tool_configuration"; import StdioConfiguration from "./StdioConfiguration"; +import ByokFields from "./ByokFields"; import MCPLogoSelector from "./MCPLogoSelector"; import EnvVarsSection from "./EnvVarsSection"; import { validateMCPServerUrl, validateMCPServerName, normalizeEnvVars } from "./utils"; @@ -65,6 +66,8 @@ const MCPServerEdit: React.FC = ({ const isOpenAPITransport = transportType === TRANSPORT.OPENAPI; const isMCPTransport = !isStdioTransport && !isOpenAPITransport; const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false; + const isByokRaw = Form.useWatch("is_byok", form); + const isByok = shouldShowAuthValueField && Boolean(isByokRaw); const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2; const isAwsSigV4AuthType = authType === AUTH_TYPE.AWS_SIGV4; const oauthFlowTypeValue = Form.useWatch("oauth_flow_type", form) as string | undefined; @@ -587,6 +590,14 @@ const MCPServerEdit: React.FC = ({ restValues.transport = "http"; } + // BYOK is only valid for static-credential auth types; drop a stale ``true`` + // (and its metadata) left behind by switching auth_type after toggling it on. + const isByokEnabled = + AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(restValues.auth_type) && Boolean(restValues.is_byok); + restValues.is_byok = isByokEnabled; + restValues.byok_description = isByokEnabled ? restValues.byok_description || [] : []; + restValues.byok_api_key_help_url = isByokEnabled ? restValues.byok_api_key_help_url || null : null; + // Parse token_validation JSON if provided let tokenValidation: Record | null = null; if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") { @@ -670,7 +681,7 @@ const MCPServerEdit: React.FC = ({ }; const includeCredentials = - restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); + !isByokEnabled && restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); if (includeCredentials && credentialsPayload && Object.keys(credentialsPayload).length > 0) { payload.credentials = credentialsPayload; @@ -824,6 +835,8 @@ const MCPServerEdit: React.FC = ({ )} + {!isStdioTransport && shouldShowAuthValueField && } + {isStdioTransport && (

@@ -884,7 +897,7 @@ const MCPServerEdit: React.FC = ({

)} - {!isStdioTransport && shouldShowAuthValueField && ( + {!isStdioTransport && shouldShowAuthValueField && !isByok && (