From 2dd151d4fcc5170321155d1a745325515e152ae8 Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 10 Nov 2026 14:47:49 -0800 Subject: [PATCH] feat(ui): send secret headers in the MCP tool preview and connection test --- .../_components/createServerPayload.ts | 3 ++ .../_components/editToolPreview.test.ts | 54 +++++++++++++++++++ .../_components/editToolPreview.ts | 37 ++++++++++--- .../src/hooks/useTestMCPConnection.test.tsx | 54 +++++++++++++++++++ .../src/hooks/useTestMCPConnection.tsx | 16 +++++- 5 files changed, 155 insertions(+), 9 deletions(-) create mode 100644 ui/litellm-dashboard/src/hooks/useTestMCPConnection.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts index 5de20ca8ffa..6d112f3bf4f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/createServerPayload.ts @@ -76,6 +76,9 @@ export const reduceStaticHeaders = (list: unknown): Record => { }, {}); }; +export const hasUnfinishedHeaderRow = (list: unknown): boolean => + Array.isArray(list) && list.some((entry) => !entry?.header?.trim() || !entry?.value?.trim()); + export const withSecretHeaders = ( credentials: Readonly> | undefined, authType: string | undefined, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts index c7147466d43..3cd2763014a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.test.ts @@ -102,4 +102,58 @@ describe("getEditToolPreview", () => { expect(getEditToolPreview({ ...saved, transport }, saved)).toEqual({ kind: "saved" }); } }); + + describe("secret headers", () => { + const savedCustom = { ...saved, auth_type: "secret_headers", static_headers: [] }; + const rows = [{ header: "DD-API-KEY", value: "rotated-secret" }]; + + it("keeps saved discovery while no header rows are entered", () => { + expect(getEditToolPreview({ ...savedCustom, secret_headers: [] }, savedCustom)).toEqual({ kind: "saved" }); + }); + + it("previews entered rows as credentials.headers", () => { + expect(getEditToolPreview({ ...savedCustom, secret_headers: rows }, savedCustom)).toEqual({ + kind: "preview", + config: { + url: "https://example.com/mcp", + transport: "http", + auth_type: "secret_headers", + static_headers: {}, + credentials: { headers: { "DD-API-KEY": "rotated-secret" } }, + }, + }); + }); + + it("needs rows before previewing a switch to secret headers", () => { + expect(getEditToolPreview({ ...saved, auth_type: "secret_headers" }, saved).kind).toBe("incomplete"); + expect(getEditToolPreview({ ...saved, auth_type: "secret_headers", secret_headers: rows }, saved).kind).toBe( + "preview", + ); + }); + + it("does not reuse the stored headers on a new origin", () => { + expect(getEditToolPreview({ ...savedCustom, url: "https://other.example/mcp" }, savedCustom).kind).toBe( + "incomplete", + ); + }); + + it("waits for a value on every row", () => { + expect( + getEditToolPreview({ ...savedCustom, secret_headers: [{ header: "DD-API-KEY", value: " " }] }, savedCustom) + .kind, + ).toBe("incomplete"); + }); + }); + + it("waits while any secret header row is unnamed or has no value", () => { + const edited = { ...saved, auth_type: "secret_headers", url: "https://example.com/corrected-mcp" }; + const complete = { header: "X-Api-Key", value: "first" }; + expect(getEditToolPreview({ ...edited, secret_headers: [complete] }, saved).kind).toBe("preview"); + for (const unfinished of [ + { header: "", value: "second" }, + { header: "X-Other", value: "" }, + ]) { + expect(getEditToolPreview({ ...edited, secret_headers: [complete, unfinished] }, saved).kind).toBe("incomplete"); + } + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts index 4ffc7dc600b..7f17183388d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/editToolPreview.ts @@ -1,7 +1,17 @@ import type { TFunction } from "i18next"; import i18n from "@/i18n/i18n"; import { AUTH_TYPE, TRANSPORT } from "@/components/mcp_tools/types"; -import { AUTH_TYPES_REQUIRING_AUTH_VALUE, reduceStaticHeaders } from "./createServerPayload"; +import { AUTH_TYPES_REQUIRING_AUTH_VALUE, hasUnfinishedHeaderRow, reduceStaticHeaders } from "./createServerPayload"; + +const previewCredentials = ( + authValue: unknown, + needsAuthValue: boolean, + secretHeaders: Readonly>, +) => { + if (needsAuthValue && typeof authValue === "string" && authValue.trim()) return { auth_value: authValue }; + if (Object.keys(secretHeaders).length > 0) return { headers: secretHeaders }; + return undefined; +}; const connectionConfig = (values: Readonly>) => { const credentials = values.credentials; @@ -9,6 +19,12 @@ const connectionConfig = (values: Readonly>) => { credentials && typeof credentials === "object" && "auth_value" in credentials ? credentials.auth_value : undefined; const needsAuthValue = typeof values.auth_type === "string" && AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(values.auth_type); + const secretHeaders = + values.auth_type === AUTH_TYPE.SECRET_HEADERS + ? Object.fromEntries( + Object.entries(reduceStaticHeaders(values.secret_headers)).sort(([a], [b]) => a.localeCompare(b)), + ) + : {}; return { url: typeof values.url === "string" ? values.url : "", transport: typeof values.transport === "string" ? values.transport : "", @@ -16,8 +32,7 @@ const connectionConfig = (values: Readonly>) => { static_headers: Object.fromEntries( Object.entries(reduceStaticHeaders(values.static_headers)).sort(([a], [b]) => a.localeCompare(b)), ), - credentials: - needsAuthValue && typeof authValue === "string" && authValue.trim() ? { auth_value: authValue } : undefined, + credentials: previewCredentials(authValue, needsAuthValue, secretHeaders), }; }; @@ -33,6 +48,7 @@ export const getEditToolPreview = ( ): EditToolPreview => { const staticAuth = values.auth_type === AUTH_TYPE.NONE || + values.auth_type === AUTH_TYPE.SECRET_HEADERS || (typeof values.auth_type === "string" && AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(values.auth_type)); if (!staticAuth || ![TRANSPORT.HTTP, TRANSPORT.SSE].includes(String(values.transport))) { return { kind: "saved" }; @@ -45,11 +61,16 @@ export const getEditToolPreview = ( const missingNewCredential = config.auth_type !== initialValues.auth_type && - AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) && + (AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) || config.auth_type === AUTH_TYPE.SECRET_HEADERS) && config.credentials === undefined; const validUrl = URL.canParse(config.url) && ["http:", "https:"].includes(new URL(config.url).protocol); - const incompleteHeaders = Object.values(config.static_headers).some((value) => !value.trim()); - if (!validUrl || missingNewCredential || incompleteHeaders) { + const incompleteHeaders = [ + ...Object.values(config.static_headers), + ...Object.values(config.credentials?.headers ?? {}), + ].some((value) => !value.trim()); + const unfinishedSecretRow = + config.auth_type === AUTH_TYPE.SECRET_HEADERS && hasUnfinishedHeaderRow(values.secret_headers); + if (!validUrl || missingNewCredential || incompleteHeaders || unfinishedSecretRow) { return { kind: "incomplete" }; } const savedConfig = connectionConfig(initialValues); @@ -58,7 +79,9 @@ export const getEditToolPreview = ( const reusesHeader = Object.entries(config.static_headers).some( ([key, value]) => savedConfig.static_headers[key] === value, ); - const needsSavedCredential = AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) && !config.credentials; + const needsSavedCredential = + (AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(config.auth_type) || config.auth_type === AUTH_TYPE.SECRET_HEADERS) && + !config.credentials; if (changedOrigin && (needsSavedCredential || reusesHeader)) { return { kind: "incomplete", diff --git a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.test.tsx b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.test.tsx new file mode 100644 index 00000000000..9b74adfb77a --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.test.tsx @@ -0,0 +1,54 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import "@/i18n/i18n"; +import { testMCPToolsListRequest } from "../components/networking"; +import { useTestMCPConnection } from "./useTestMCPConnection"; + +vi.mock("../components/networking", () => ({ testMCPToolsListRequest: vi.fn() })); + +const formValues = { + url: "https://upstream.example/mcp", + transport: "http", + auth_type: "secret_headers", + static_headers: [{ header: "X-Upstream-Site", value: "example.com" }], +}; +const rows = [ + { header: "DD-API-KEY", value: "api-secret" }, + { header: "DD-APPLICATION-KEY", value: "app-secret" }, +]; + +describe("useTestMCPConnection with secret headers", () => { + beforeEach(() => { + vi.mocked(testMCPToolsListRequest).mockReset().mockResolvedValue({ tools: [] }); + }); + + it("sends the header rows as credentials.headers once every row is complete", async () => { + const { result } = renderHook(() => + useTestMCPConnection({ accessToken: "sk-admin", formValues: { ...formValues, secret_headers: rows } }), + ); + await waitFor(() => expect(testMCPToolsListRequest).toHaveBeenCalledTimes(1)); + expect(vi.mocked(testMCPToolsListRequest).mock.calls[0][1]).toMatchObject({ + auth_type: "secret_headers", + static_headers: { "X-Upstream-Site": "example.com" }, + credentials: { headers: { "DD-API-KEY": "api-secret", "DD-APPLICATION-KEY": "app-secret" } }, + }); + expect(result.current.canFetchTools).toBe(true); + }); + + it.each([ + [[]], + [[{ header: "DD-API-KEY", value: "" }]], + [ + [ + { header: "DD-API-KEY", value: "api-secret" }, + { header: "", value: "second-secret" }, + ], + ], + ])("does not call upstream for incomplete rows %j", (custom) => { + const { result } = renderHook(() => + useTestMCPConnection({ accessToken: "sk-admin", formValues: { ...formValues, secret_headers: custom } }), + ); + expect(result.current.canFetchTools).toBe(false); + expect(testMCPToolsListRequest).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx index d01174d2de6..cd80f2cb1fe 100644 --- a/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx +++ b/ui/litellm-dashboard/src/hooks/useTestMCPConnection.tsx @@ -2,6 +2,10 @@ import { useState, useEffect, useCallback } from "react"; import { useTranslation } from "react-i18next"; import { testMCPToolsListRequest } from "../components/networking"; import { AUTH_TYPE, OAUTH_FLOW, TRANSPORT, isClientForwardedTokenMode } from "@/components/mcp_tools/types"; +import { + hasUnfinishedHeaderRow, + withSecretHeaders, +} from "@/app/(dashboard)/mcp-servers/_components/createServerPayload"; interface MCPServerConfig { server_id?: string; @@ -20,6 +24,7 @@ interface MCPServerConfig { client_id?: string; client_secret?: string; scopes?: string[]; + headers?: Record; }; } @@ -62,6 +67,9 @@ export const useTestMCPConnection = ({ const requiresOAuthToken = (formValues.auth_type === AUTH_TYPE.OAUTH2 && !isM2MOAuth) || isBrowserHeldTokenMode; const isOpenAPITransport = formValues.transport === TRANSPORT.OPENAPI; const hasEndpoint = isOpenAPITransport ? !!formValues.spec_path : !!formValues.url; + const lacksSecretHeaders = + formValues.auth_type === AUTH_TYPE.SECRET_HEADERS && + (!formValues.secret_headers?.length || hasUnfinishedHeaderRow(formValues.secret_headers)); // For OpenAPI: tools are derived from the spec itself — no auth needed to load them. // For other transports: auth_type is required, and OAuth interactive flows need a token. @@ -72,11 +80,13 @@ export const useTestMCPConnection = ({ formValues.transport && formValues.auth_type && accessToken && + !lacksSecretHeaders && (!requiresOAuthToken || oauthAccessToken) ); const staticHeadersKey = JSON.stringify(formValues.static_headers ?? {}); const credentialsKey = JSON.stringify(formValues.credentials ?? {}); + const secretHeadersKey = JSON.stringify(formValues.secret_headers ?? []); const fetchTools = async () => { if (!accessToken || (!formValues.url && !formValues.spec_path)) { @@ -152,8 +162,9 @@ export const useTestMCPConnection = ({ static_headers: staticHeaders, }; - if (credentials && Object.keys(credentials).length > 0) { - mcpServerConfig.credentials = credentials; + const credentialsWithHeaders = withSecretHeaders(credentials, formValues.auth_type, formValues.secret_headers); + if (credentialsWithHeaders && Object.keys(credentialsWithHeaders).length > 0) { + mcpServerConfig.credentials = credentialsWithHeaders; } const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig, oauthAccessToken); @@ -217,6 +228,7 @@ export const useTestMCPConnection = ({ canFetchTools, staticHeadersKey, credentialsKey, + secretHeadersKey, ]); return {