Add support for custom headers and migrate from legacy auth

- Introduced CustomHeaders component for managing custom headers.
- Updated Sidebar to include CustomHeaders component.
- Enhanced useConnection hook to handle custom headers and legacy auth migration.
- Added tests for custom headers functionality.
- Updated package.json and package-lock.json for new dependencies.
This commit is contained in:
jeetul.samaiya 2025-08-26 18:51:06 +05:30
commit c5acc63129
8 changed files with 606 additions and 48 deletions

View file

@ -74,6 +74,7 @@ import ElicitationTab, {
PendingElicitationRequest,
ElicitationResponse,
} from "./components/ElicitationTab";
import { CustomHeaders, migrateFromLegacyAuth } from "./lib/types/customHeaders";
const CONFIG_LOCAL_STORAGE_KEY = "inspectorConfig_v1";
@ -127,6 +128,28 @@ const App = () => {
return localStorage.getItem("lastOauthScope") || "";
});
// Custom headers state with migration from legacy auth
const [customHeaders, setCustomHeaders] = useState<CustomHeaders>(() => {
const savedHeaders = localStorage.getItem("lastCustomHeaders");
if (savedHeaders) {
try {
return JSON.parse(savedHeaders);
} catch {
// Fall back to migration if JSON parsing fails
}
}
// Migrate from legacy auth if available
const legacyToken = localStorage.getItem("lastBearerToken") || "";
const legacyHeaderName = localStorage.getItem("lastHeaderName") || "";
if (legacyToken) {
return migrateFromLegacyAuth(legacyToken, legacyHeaderName);
}
return [];
});
const [pendingSampleRequests, setPendingSampleRequests] = useState<
Array<
PendingRequest & {
@ -215,6 +238,7 @@ const App = () => {
env,
bearerToken,
headerName,
customHeaders,
oauthClientId,
oauthScope,
config,
@ -310,6 +334,23 @@ const App = () => {
localStorage.setItem("lastHeaderName", headerName);
}, [headerName]);
useEffect(() => {
localStorage.setItem("lastCustomHeaders", JSON.stringify(customHeaders));
}, [customHeaders]);
// Auto-migrate from legacy auth when custom headers are empty but legacy auth exists
useEffect(() => {
if (customHeaders.length === 0 && (bearerToken || headerName)) {
const migratedHeaders = migrateFromLegacyAuth(bearerToken, headerName);
if (migratedHeaders.length > 0) {
setCustomHeaders(migratedHeaders);
// Clear legacy auth after migration
setBearerToken("");
setHeaderName("");
}
}
}, [bearerToken, headerName, customHeaders, setCustomHeaders]);
useEffect(() => {
localStorage.setItem("lastOauthClientId", oauthClientId);
}, [oauthClientId]);
@ -814,6 +855,8 @@ const App = () => {
setBearerToken={setBearerToken}
headerName={headerName}
setHeaderName={setHeaderName}
customHeaders={customHeaders}
setCustomHeaders={setCustomHeaders}
oauthClientId={oauthClientId}
setOauthClientId={setOauthClientId}
oauthScope={oauthScope}

View file

@ -0,0 +1,215 @@
import { useState } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Textarea } from "@/components/ui/textarea";
import { Switch } from "@/components/ui/switch";
import { Plus, Trash2, Eye, EyeOff } from "lucide-react";
import { CustomHeaders as CustomHeadersType, CustomHeader, createEmptyHeader } from "@/lib/types/customHeaders";
interface CustomHeadersProps {
headers: CustomHeadersType;
onChange: (headers: CustomHeadersType) => void;
className?: string;
}
const CustomHeaders = ({ headers, onChange, className }: CustomHeadersProps) => {
const [isJsonMode, setIsJsonMode] = useState(false);
const [jsonValue, setJsonValue] = useState("");
const [jsonError, setJsonError] = useState<string | null>(null);
const [visibleValues, setVisibleValues] = useState<Set<number>>(new Set());
const updateHeader = (index: number, field: keyof CustomHeader, value: string | boolean) => {
const newHeaders = [...headers];
newHeaders[index] = { ...newHeaders[index], [field]: value };
onChange(newHeaders);
};
const addHeader = () => {
onChange([...headers, createEmptyHeader()]);
};
const removeHeader = (index: number) => {
const newHeaders = headers.filter((_, i) => i !== index);
onChange(newHeaders);
};
const toggleValueVisibility = (index: number) => {
const newVisible = new Set(visibleValues);
if (newVisible.has(index)) {
newVisible.delete(index);
} else {
newVisible.add(index);
}
setVisibleValues(newVisible);
};
const switchToJsonMode = () => {
const jsonObject: Record<string, string> = {};
headers.forEach((header) => {
if (header.enabled && header.name.trim() && header.value.trim()) {
jsonObject[header.name.trim()] = header.value.trim();
}
});
setJsonValue(JSON.stringify(jsonObject, null, 2));
setJsonError(null);
setIsJsonMode(true);
};
const switchToFormMode = () => {
try {
const parsed = JSON.parse(jsonValue);
if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
setJsonError("JSON must be an object with string key-value pairs");
return;
}
const newHeaders: CustomHeadersType = Object.entries(parsed).map(([name, value]) => ({
name,
value: String(value),
enabled: true,
}));
onChange(newHeaders);
setJsonError(null);
setIsJsonMode(false);
} catch (error) {
setJsonError("Invalid JSON format");
}
};
const handleJsonChange = (value: string) => {
setJsonValue(value);
setJsonError(null);
};
if (isJsonMode) {
return (
<div className={`space-y-3 ${className}`}>
<div className="flex justify-between items-center gap-2">
<h4 className="text-sm font-semibold flex-shrink-0">Custom Headers (JSON)</h4>
<Button
type="button"
variant="outline"
size="sm"
onClick={switchToFormMode}
className="flex-shrink-0"
>
Switch to Form
</Button>
</div>
<div className="space-y-2">
<Textarea
value={jsonValue}
onChange={(e) => handleJsonChange(e.target.value)}
placeholder='{\n "Authorization": "Bearer token123",\n "X-Tenant-ID": "acme-inc",\n "X-Environment": "staging"\n}'
className="font-mono text-sm min-h-[100px] resize-none"
/>
{jsonError && (
<p className="text-sm text-red-600">{jsonError}</p>
)}
<p className="text-xs text-muted-foreground">
Enter headers as a JSON object with string key-value pairs.
</p>
</div>
</div>
);
}
return (
<div className={`space-y-3 ${className}`}>
<div className="flex justify-between items-center gap-2">
<h4 className="text-sm font-semibold flex-shrink-0">Custom Headers</h4>
<div className="flex gap-1 flex-shrink-0">
<Button
type="button"
variant="outline"
size="sm"
onClick={switchToJsonMode}
className="text-xs px-2"
>
JSON
</Button>
<Button
type="button"
variant="outline"
size="sm"
onClick={addHeader}
className="text-xs px-2"
>
<Plus className="w-3 h-3 mr-1" />
Add
</Button>
</div>
</div>
{headers.length === 0 ? (
<div className="text-center py-4 text-muted-foreground">
<p className="text-sm">No custom headers configured</p>
<p className="text-xs mt-1">Click "Add" to get started</p>
</div>
) : (
<div className="space-y-2 max-h-[300px] overflow-y-auto">
{headers.map((header, index) => (
<div
key={index}
className="flex items-start gap-2 p-2 border rounded-md"
>
<Switch
checked={header.enabled}
onCheckedChange={(enabled) => updateHeader(index, "enabled", enabled)}
className="shrink-0 mt-2"
/>
<div className="flex-1 min-w-0 space-y-2">
<Input
placeholder="Header Name"
value={header.name}
onChange={(e) => updateHeader(index, "name", e.target.value)}
className="font-mono text-xs"
/>
<div className="relative">
<Input
placeholder="Header Value"
value={header.value}
onChange={(e) => updateHeader(index, "value", e.target.value)}
type={visibleValues.has(index) ? "text" : "password"}
className="font-mono text-xs pr-8"
/>
<Button
type="button"
variant="ghost"
size="sm"
onClick={() => toggleValueVisibility(index)}
className="absolute right-1 top-1/2 -translate-y-1/2 h-6 w-6 p-0"
>
{visibleValues.has(index) ? (
<EyeOff className="w-3 h-3" />
) : (
<Eye className="w-3 h-3" />
)}
</Button>
</div>
</div>
<Button
type="button"
variant="ghost"
size="sm"
onClick={() => removeHeader(index)}
className="shrink-0 text-red-600 hover:text-red-700 hover:bg-red-50 h-6 w-6 p-0 mt-2"
>
<Trash2 className="w-3 h-3" />
</Button>
</div>
))}
</div>
)}
{headers.length > 0 && (
<p className="text-xs text-muted-foreground">
Use the toggle to enable/disable headers. Only enabled headers with both name and value will be sent.
</p>
)}
</div>
);
};
export default CustomHeaders;

View file

@ -37,6 +37,8 @@ import {
TooltipTrigger,
TooltipContent,
} from "@/components/ui/tooltip";
import CustomHeaders from "./CustomHeaders";
import { CustomHeaders as CustomHeadersType } from "@/lib/types/customHeaders";
import { useToast } from "../lib/hooks/useToast";
interface SidebarProps {
@ -51,10 +53,14 @@ interface SidebarProps {
setSseUrl: (url: string) => void;
env: Record<string, string>;
setEnv: (env: Record<string, string>) => void;
// Legacy auth props (for backward compatibility)
bearerToken: string;
setBearerToken: (token: string) => void;
headerName?: string;
setHeaderName?: (name: string) => void;
// New custom headers support
customHeaders: CustomHeadersType;
setCustomHeaders: (headers: CustomHeadersType) => void;
oauthClientId: string;
setOauthClientId: (id: string) => void;
oauthScope: string;
@ -80,10 +86,12 @@ const Sidebar = ({
setSseUrl,
env,
setEnv,
bearerToken,
setBearerToken,
headerName,
setHeaderName,
bearerToken: _bearerToken,
setBearerToken: _setBearerToken,
headerName: _headerName,
setHeaderName: _setHeaderName,
customHeaders,
setCustomHeaders,
oauthClientId,
setOauthClientId,
oauthScope,
@ -497,38 +505,12 @@ const Sidebar = ({
</Button>
{showAuthConfig && (
<>
{/* Bearer Token Section */}
<div className="space-y-2 p-3 rounded border">
<h4 className="text-sm font-semibold flex items-center">
API Token Authentication
</h4>
<div className="space-y-2">
<label className="text-sm font-medium">Header Name</label>
<Input
placeholder="Authorization"
onChange={(e) =>
setHeaderName && setHeaderName(e.target.value)
}
data-testid="header-input"
className="font-mono"
value={headerName}
/>
<label
className="text-sm font-medium"
htmlFor="bearer-token-input"
>
Bearer Token
</label>
<Input
id="bearer-token-input"
placeholder="Bearer Token"
value={bearerToken}
onChange={(e) => setBearerToken(e.target.value)}
data-testid="bearer-token-input"
className="font-mono"
type="password"
/>
</div>
{/* Custom Headers Section */}
<div className="p-3 rounded border overflow-hidden">
<CustomHeaders
headers={customHeaders}
onChange={setCustomHeaders}
/>
</div>
{transportType !== "stdio" && (
// OAuth Configuration

View file

@ -50,6 +50,10 @@ describe("Sidebar", () => {
setEnv: jest.fn(),
bearerToken: "",
setBearerToken: jest.fn(),
headerName: "",
setHeaderName: jest.fn(),
customHeaders: [],
setCustomHeaders: jest.fn(),
onConnect: jest.fn(),
onDisconnect: jest.fn(),
stdErrNotifications: [],

View file

@ -0,0 +1,27 @@
import * as React from "react";
import * as SwitchPrimitives from "@radix-ui/react-switch";
import { cn } from "@/lib/utils";
const Switch = React.forwardRef<
React.ElementRef<typeof SwitchPrimitives.Root>,
React.ComponentPropsWithoutRef<typeof SwitchPrimitives.Root>
>(({ className, ...props }, ref) => (
<SwitchPrimitives.Root
className={cn(
"peer inline-flex h-5 w-9 shrink-0 cursor-pointer items-center rounded-full border-2 border-transparent shadow-sm transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background disabled:cursor-not-allowed disabled:opacity-50 data-[state=checked]:bg-primary data-[state=unchecked]:bg-input",
className
)}
{...props}
ref={ref}
>
<SwitchPrimitives.Thumb
className={cn(
"pointer-events-none block h-4 w-4 rounded-full bg-background shadow-lg ring-0 transition-transform data-[state=checked]:translate-x-4 data-[state=unchecked]:translate-x-0"
)}
/>
</SwitchPrimitives.Root>
));
Switch.displayName = SwitchPrimitives.Root.displayName;
export { Switch };

View file

@ -13,6 +13,7 @@ import {
} from "@modelcontextprotocol/sdk/types.js";
import { auth } from "@modelcontextprotocol/sdk/client/auth.js";
import { discoverScopes } from "../../auth";
import { CustomHeaders } from "../../types/customHeaders";
// Mock fetch
global.fetch = jest.fn().mockResolvedValue({
@ -740,6 +741,135 @@ describe("useConnection", () => {
});
});
describe("Custom Headers", () => {
beforeEach(() => {
jest.clearAllMocks();
// Reset the mock transport objects
mockSSETransport.url = undefined;
mockSSETransport.options = undefined;
mockStreamableHTTPTransport.url = undefined;
mockStreamableHTTPTransport.options = undefined;
});
test("sends multiple custom headers correctly", async () => {
const customHeaders: CustomHeaders = [
{ name: "Authorization", value: "Bearer token123", enabled: true },
{ name: "X-Tenant-ID", value: "acme-inc", enabled: true },
{ name: "X-Environment", value: "staging", enabled: true },
];
const propsWithCustomHeaders = {
...defaultProps,
customHeaders,
};
const { result } = renderHook(() =>
useConnection(propsWithCustomHeaders),
);
await act(async () => {
await result.current.connect();
});
// Check that the transport was created with the correct headers
expect(mockSSETransport.options).toBeDefined();
expect(mockSSETransport.options?.requestInit?.headers).toBeDefined();
const headers = mockSSETransport.options?.requestInit?.headers;
expect(headers).toHaveProperty("Authorization", "Bearer token123");
expect(headers).toHaveProperty("X-Tenant-ID", "acme-inc");
expect(headers).toHaveProperty("X-Environment", "staging");
expect(headers).toHaveProperty("x-custom-auth-headers", JSON.stringify(["X-Tenant-ID", "X-Environment"]));
});
test("ignores disabled custom headers", async () => {
const customHeaders: CustomHeaders = [
{ name: "Authorization", value: "Bearer token123", enabled: true },
{ name: "X-Disabled", value: "should-not-appear", enabled: false },
{ name: "X-Enabled", value: "should-appear", enabled: true },
];
const propsWithCustomHeaders = {
...defaultProps,
customHeaders,
};
const { result } = renderHook(() =>
useConnection(propsWithCustomHeaders),
);
await act(async () => {
await result.current.connect();
});
const headers = mockSSETransport.options?.requestInit?.headers;
expect(headers).toHaveProperty("Authorization", "Bearer token123");
expect(headers).toHaveProperty("X-Enabled", "should-appear");
expect(headers).not.toHaveProperty("X-Disabled");
});
test("migrates from legacy bearerToken and headerName", async () => {
const propsWithLegacyAuth = {
...defaultProps,
bearerToken: "legacy-token",
headerName: "X-Custom-Auth",
};
const { result } = renderHook(() =>
useConnection(propsWithLegacyAuth),
);
await act(async () => {
await result.current.connect();
});
const headers = mockSSETransport.options?.requestInit?.headers;
expect(headers).toHaveProperty("X-Custom-Auth", "legacy-token");
expect(headers).toHaveProperty("x-custom-auth-headers", JSON.stringify(["X-Custom-Auth"]));
});
test("uses OAuth token when no custom headers or legacy auth provided", async () => {
const propsWithoutAuth = {
...defaultProps,
};
const { result } = renderHook(() =>
useConnection(propsWithoutAuth),
);
await act(async () => {
await result.current.connect();
});
const headers = mockSSETransport.options?.requestInit?.headers;
expect(headers).toHaveProperty("Authorization", "Bearer mock-token");
});
test("prioritizes custom headers over legacy auth", async () => {
const customHeaders: CustomHeaders = [
{ name: "Authorization", value: "Bearer custom-token", enabled: true },
];
const propsWithBothAuth = {
...defaultProps,
customHeaders,
bearerToken: "legacy-token",
headerName: "Authorization",
};
const { result } = renderHook(() =>
useConnection(propsWithBothAuth),
);
await act(async () => {
await result.current.connect();
});
const headers = mockSSETransport.options?.requestInit?.headers;
expect(headers).toHaveProperty("Authorization", "Bearer custom-token");
});
});
describe("OAuth Error Handling with Scope Discovery", () => {
beforeEach(() => {
jest.clearAllMocks();

View file

@ -57,6 +57,7 @@ import {
import { getMCPServerRequestTimeout } from "@/utils/configUtils";
import { InspectorConfig } from "../configurationTypes";
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
import { CustomHeaders, migrateFromLegacyAuth } from "../types/customHeaders";
interface UseConnectionOptions {
transportType: "stdio" | "sse" | "streamable-http";
@ -64,8 +65,11 @@ interface UseConnectionOptions {
args: string;
sseUrl: string;
env: Record<string, string>;
// Legacy auth support (for backward compatibility)
bearerToken?: string;
headerName?: string;
// New custom headers support
customHeaders?: CustomHeaders;
oauthClientId?: string;
oauthScope?: string;
config: InspectorConfig;
@ -88,6 +92,7 @@ export function useConnection({
env,
bearerToken,
headerName,
customHeaders,
oauthClientId,
oauthScope,
config,
@ -379,21 +384,52 @@ export function useConnection({
// Create an auth provider with the current server URL
const serverAuthProvider = new InspectorOAuthClientProvider(sseUrl);
// Use manually provided bearer token if available, otherwise use OAuth tokens
const token =
bearerToken || (await serverAuthProvider.tokens())?.access_token;
if (token) {
const authHeaderName = headerName || "Authorization";
// Handle custom headers (new approach) or legacy auth (backward compatibility)
let finalHeaders: CustomHeaders = [];
// Add custom header name as a special request header to let the server know which header to pass through
if (authHeaderName.toLowerCase() !== "authorization") {
headers[authHeaderName] = token;
headers["x-custom-auth-header"] = authHeaderName;
} else {
headers[authHeaderName] = `Bearer ${token}`;
if (customHeaders && customHeaders.length > 0) {
// Use new custom headers approach
finalHeaders = customHeaders;
} else if (bearerToken || headerName) {
// Migrate from legacy auth approach
finalHeaders = migrateFromLegacyAuth(bearerToken, headerName);
}
// Add OAuth token if available and no custom headers are set
if (finalHeaders.length === 0) {
const oauthToken = (await serverAuthProvider.tokens())?.access_token;
if (oauthToken) {
finalHeaders = [
{
name: "Authorization",
value: `Bearer ${oauthToken}`,
enabled: true,
},
];
}
}
// Process all enabled custom headers
const customHeaderNames: string[] = [];
finalHeaders.forEach((header) => {
if (header.enabled && header.name.trim() && header.value.trim()) {
const headerName = header.name.trim();
const headerValue = header.value.trim();
headers[headerName] = headerValue;
// Track custom header names for server processing
if (headerName.toLowerCase() !== "authorization") {
customHeaderNames.push(headerName);
}
}
});
// Add custom header names as a special request header for server processing
if (customHeaderNames.length > 0) {
headers["x-custom-auth-headers"] = JSON.stringify(customHeaderNames);
}
// Add proxy authentication
const { token: proxyAuthToken, header: proxyAuthTokenHeader } =
getMCPProxyAuthToken(config);

View file

@ -0,0 +1,59 @@
export interface CustomHeader {
name: string;
value: string;
enabled: boolean;
}
export type CustomHeaders = CustomHeader[];
export const createEmptyHeader = (): CustomHeader => ({
name: "",
value: "",
enabled: true,
});
export const createHeaderFromBearerToken = (
bearerToken: string,
headerName?: string,
): CustomHeader => ({
name: headerName || "Authorization",
value: headerName?.toLowerCase() === "authorization" || !headerName
? `Bearer ${bearerToken}`
: bearerToken,
enabled: true,
});
export const getEnabledHeaders = (headers: CustomHeaders): CustomHeaders => {
return headers.filter(header => header.enabled && header.name.trim() && header.value.trim());
};
export const headersToRecord = (headers: CustomHeaders): Record<string, string> => {
const enabledHeaders = getEnabledHeaders(headers);
const record: Record<string, string> = {};
enabledHeaders.forEach(header => {
record[header.name.trim()] = header.value.trim();
});
return record;
};
export const recordToHeaders = (record: Record<string, string>): CustomHeaders => {
return Object.entries(record).map(([name, value]) => ({
name,
value,
enabled: true,
}));
};
// Migration helper for backward compatibility
export const migrateFromLegacyAuth = (
bearerToken?: string,
headerName?: string,
): CustomHeaders => {
if (!bearerToken) {
return [];
}
return [createHeaderFromBearerToken(bearerToken, headerName)];
};