feat(ui): add custom request headers to the API Playground

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-15 22:12:06 +00:00
parent a6127d2363
commit c05af7a12f
20 changed files with 364 additions and 38 deletions

View file

@ -35,10 +35,12 @@ beforeEach(() => {
Element.prototype.scrollIntoView = () => {};
});
const CHAT_REQUEST_ARG_COUNT = 26;
const CHAT_REQUEST_ARG_COUNT = 27;
const STREAMING_ENABLED_ARG_INDEX = 25;
const MESSAGES_REQUEST_ARG_COUNT = 19;
const CHAT_CUSTOM_HEADERS_ARG_INDEX = 26;
const MESSAGES_REQUEST_ARG_COUNT = 20;
const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18;
const MESSAGES_CUSTOM_HEADERS_ARG_INDEX = 19;
async function openComboboxByPlaceholder(placeholder: string) {
const user = userEvent.setup();
@ -447,6 +449,63 @@ describe("ChatUI", () => {
expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false);
});
it("should send custom headers entered in the sidebar with /v1/chat/completions and /v1/messages requests", async () => {
const user = userEvent.setup();
render(
<ChatUI
accessToken="1234567890"
token="1234567890"
userRole="user"
userID="1234567890"
disabledPersonalKeyCreation={false}
/>,
);
await waitFor(() => {
expect(screen.getByText("Test Key")).toBeInTheDocument();
});
await selectComboboxOption("Select a Model", "Model 1");
await user.click(screen.getByRole("button", { name: "Add Header" }));
await user.click(screen.getByRole("button", { name: "Add Header" }));
const [firstName] = screen.getAllByPlaceholderText("Header Name");
const [firstValue, secondValue] = screen.getAllByPlaceholderText("Header Value");
fireEvent.change(firstName, { target: { value: "anthropic-beta" } });
fireEvent.change(firstValue, { target: { value: "context-1m-2025-08-07" } });
fireEvent.change(secondValue, { target: { value: "ignored because the name is blank" } });
const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeOpenAIChatCompletionRequest).toHaveBeenCalledTimes(1);
});
const chatArgs = vi.mocked(makeOpenAIChatCompletionRequest).mock.calls[0];
expect(chatArgs).toHaveLength(CHAT_REQUEST_ARG_COUNT);
expect(chatArgs[CHAT_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
await selectComboboxOption("Select an endpoint", "/v1/messages");
await selectComboboxOption("Select a Model", "Model 1");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello again" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1);
});
const messagesArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0];
expect(messagesArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT);
expect(messagesArgs[MESSAGES_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
});
it("should force streaming in simplified mode even when the playground setting is off", async () => {
sessionStorage.setItem("streamingEnabled", "false");

View file

@ -9,6 +9,7 @@ import {
Info,
Key,
Link2,
ListPlus,
Loader2,
Settings,
Shield,
@ -40,6 +41,8 @@ import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages
import { makeOpenAIAudioSpeechRequest } from "../../llm_calls/audio_speech";
import { makeOpenAIAudioTranscriptionRequest } from "../../llm_calls/audio_transcriptions";
import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion";
import { customHeadersFromPairs, parseStoredHeaderPairs } from "@/components/llm_calls/request_headers";
import KeyValueInput, { type KeyValuePair } from "@/components/key_value_input";
import { makeOpenAIEmbeddingsRequest } from "../../llm_calls/embeddings_api";
import { Agent, fetchAvailableAgents } from "../../llm_calls/fetch_agents";
import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models";
@ -220,6 +223,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
return [];
}
});
const [customHeaderPairs, setCustomHeaderPairs] = useState<readonly KeyValuePair[]>(() =>
parseStoredHeaderPairs(getSecureItem("customHeaders")),
);
const customHeaders = useMemo(() => customHeadersFromPairs(customHeaderPairs), [customHeaderPairs]);
const [selectedVoice, setSelectedVoice] = useState<OpenAIVoice>(() => {
const saved = sessionStorage.getItem("selectedVoice");
if (!saved) return "alloy";
@ -346,6 +353,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedSdk,
selectedVoice,
proxySettings,
customHeaders,
});
setGeneratedCode(code);
}
@ -367,12 +375,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
endpointType,
selectedModel,
proxySettings,
customHeaders,
]);
useEffect(() => {
try {
setSecureItem("apiKeySource", JSON.stringify(apiKeySource));
setSecureItem("apiKey", apiKey);
setSecureItem("customHeaders", JSON.stringify(customHeaderPairs));
} catch {
// Storage full or unavailable — non-critical, skip persisting.
}
@ -410,6 +420,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpServerToolRestrictions,
selectedVoice,
streamingEnabled,
customHeaderPairs,
]);
useEffect(() => {
@ -921,6 +932,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mockTestFallbacks,
mcpToolsets,
streamingEnabled,
customHeaders,
);
} else if (endpointType === EndpointType.IMAGE) {
// For image generation
@ -932,6 +944,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
signal,
customProxyBaseUrl || undefined,
customHeaders,
);
} else if (endpointType === EndpointType.SPEECH) {
// For audio speech
@ -946,6 +959,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
undefined, // responseFormat
undefined, // speed
customProxyBaseUrl || undefined,
customHeaders,
);
} else if (endpointType === EndpointType.IMAGE_EDITS) {
// For image edits
@ -959,6 +973,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
signal,
customProxyBaseUrl || undefined,
customHeaders,
);
}
} else if (endpointType === EndpointType.RESPONSES) {
@ -1004,6 +1019,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpToolsets,
streamingEnabled,
updateTotalLatency,
customHeaders,
);
} else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) {
const apiChatHistory = [
@ -1033,6 +1049,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpServerToolRestrictions,
mcpToolsets,
streamingEnabled,
customHeaders,
);
} else if (endpointType === EndpointType.EMBEDDINGS) {
await makeOpenAIEmbeddingsRequest(
@ -1042,6 +1059,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
effectiveApiKey,
selectedTags,
customProxyBaseUrl || undefined,
customHeaders,
);
} else if (endpointType === EndpointType.TRANSCRIPTION) {
// For audio transcriptions
@ -1058,6 +1076,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
undefined, // responseFormat
undefined, // temperature
customProxyBaseUrl || undefined,
customHeaders,
);
}
} else if (endpointType === EndpointType.INTERACTIONS) {
@ -1069,6 +1088,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
signal,
customProxyBaseUrl || undefined,
undefined,
customHeaders,
);
}
}
@ -1086,13 +1107,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
resolvedServerId = toolEntry?.server_id ?? rawSelected;
}
if (resolvedServerId && !resolvedServerId.startsWith("toolset:") && selectedMCPDirectTool) {
const result = await callMCPTool(
effectiveApiKey,
resolvedServerId,
selectedMCPDirectTool,
mcpToolArguments,
selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined,
);
const result = await callMCPTool(effectiveApiKey, resolvedServerId, selectedMCPDirectTool, mcpToolArguments, {
...(selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : {}),
customHeaders,
});
const resultText =
result?.content?.length > 0
? JSON.stringify(
@ -1118,6 +1136,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
updateA2AMetadata,
customProxyBaseUrl || undefined,
selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
customHeaders,
);
}
} catch (error) {
@ -1485,6 +1504,18 @@ const ChatUI: React.FC<ChatUIProps> = ({
/>
</div>
{endpointType !== EndpointType.REALTIME && (
<div>
<label className="mb-2 flex items-center text-sm font-medium text-foreground">
<ListPlus className="mr-2 size-4" aria-hidden="true" /> Custom Headers
</label>
<KeyValueInput value={customHeaderPairs} onChange={setCustomHeaderPairs} />
<p className="mt-2 text-xs text-muted-foreground">
Sent with every playground request, e.g. provider-specific headers like anthropic-beta.
</p>
</div>
)}
<div>
<div className="mb-2 flex items-center gap-1 text-sm font-medium text-foreground">
<Wrench className="mr-1 size-4" aria-hidden="true" />

View file

@ -3,6 +3,7 @@
import { v4 as uuidv4 } from "uuid";
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
import type { CustomHeaders } from "@/components/llm_calls/request_headers";
import { A2ATaskMetadata } from "@/components/chat_ui/types";
interface A2AMessagePart {
@ -116,6 +117,7 @@ export const makeA2ASendMessageRequest = async (
onA2AMetadata?: (metadata: A2ATaskMetadata) => void,
customBaseUrl?: string,
guardrails?: string[],
customHeaders?: CustomHeaders,
): Promise<void> => {
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
const url = proxyBaseUrl ? `${proxyBaseUrl}/a2a/${agentId}/message/send` : `/a2a/${agentId}/message/send`;
@ -149,6 +151,7 @@ export const makeA2ASendMessageRequest = async (
headers: {
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
...customHeaders,
},
body: JSON.stringify(jsonRpcRequest),
signal,

View file

@ -1,4 +1,5 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import Anthropic from "@anthropic-ai/sdk";
import { makeAnthropicMessagesRequest } from "./anthropic_messages";
import type { TokenUsage } from "@/components/chat_ui/ResponseMetrics";
@ -122,4 +123,26 @@ describe("anthropic_messages non-streaming", () => {
expect(mockMessagesCreate).not.toHaveBeenCalled();
expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true });
});
it("sends custom headers alongside the tags header on the Anthropic client", async () => {
mockMessagesCreate.mockResolvedValue({ content: [{ type: "text", text: "OK" }], usage: {} });
await makeAnthropicMessagesRequest(
[{ role: "user", content: "Hello" }],
vi.fn(),
"claude-haiku-4-5",
"test-token",
["team-a"],
undefined,
undefined,
undefined,
undefined,
...NON_STREAMING_ARGS,
{ "anthropic-beta": "context-1m-2025-08-07" },
);
expect(vi.mocked(Anthropic).mock.calls[0][0]).toMatchObject({
defaultHeaders: { "x-litellm-tags": "team-a", "anthropic-beta": "context-1m-2025-08-07" },
});
});
});

View file

@ -2,6 +2,7 @@ import Anthropic from "@anthropic-ai/sdk";
import { MessageType } from "@/components/chat_ui/types";
import { TokenUsage } from "@/components/chat_ui/ResponseMetrics";
import { buildMcpToolBlocks } from "@/components/llm_calls/mcp_tool_blocks";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { MCPServer, MCPToolset } from "@/components/mcp_tools/types";
import { getProxyBaseUrl } from "@/components/networking";
import { toast } from "@/lib/toast";
@ -34,6 +35,7 @@ export async function makeAnthropicMessagesRequest(
mcpServerToolRestrictions?: Record<string, string[]>,
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
customHeaders?: CustomHeaders,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -46,11 +48,7 @@ export async function makeAnthropicMessagesRequest(
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
const client = new Anthropic({
apiKey: accessToken,

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
import type { OpenAIVoice } from "../components/chat_ui/chatConstants";
@ -14,6 +15,7 @@ export async function makeOpenAIAudioSpeechRequest(
responseFormat?: string,
speed?: number,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -25,7 +27,7 @@ export async function makeOpenAIAudioSpeechRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
export async function makeOpenAIAudioTranscriptionRequest(
@ -14,6 +15,7 @@ export async function makeOpenAIAudioTranscriptionRequest(
responseFormat?: string,
temperature?: number,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -26,7 +28,7 @@ export async function makeOpenAIAudioTranscriptionRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -76,4 +76,24 @@ describe("embeddings_api", () => {
input: "Sample text",
});
});
it("sends custom headers on the fetch request, letting them override the tags header", async () => {
await makeOpenAIEmbeddingsRequest(
"Sample text",
mockUpdateEmbeddingsUI,
"text-embedding-3-small",
"abcdef",
["team-a"],
undefined,
{ "x-litellm-tags": "team-b", "x-request-source": "playground" },
);
expect(mockFetch.mock.calls[0][1]).toMatchObject({
headers: {
Authorization: "Bearer abcdef",
"x-litellm-tags": "team-b",
"x-request-source": "playground",
},
});
});
});

View file

@ -1,5 +1,6 @@
import { toast } from "@/lib/toast";
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
export async function makeOpenAIEmbeddingsRequest(
input: string,
@ -8,6 +9,7 @@ export async function makeOpenAIEmbeddingsRequest(
accessToken: string,
tags?: string[],
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -20,11 +22,7 @@ export async function makeOpenAIEmbeddingsRequest(
}
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
try {
const normalizedBaseUrl = proxyBaseUrl.endsWith("/") ? proxyBaseUrl.slice(0, -1) : proxyBaseUrl;

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
export async function makeOpenAIImageEditsRequest(
@ -11,6 +12,7 @@ export async function makeOpenAIImageEditsRequest(
tags?: string[],
signal?: AbortSignal,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -23,7 +25,7 @@ export async function makeOpenAIImageEditsRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
export async function makeOpenAIImageGenerationRequest(
@ -10,6 +11,7 @@ export async function makeOpenAIImageGenerationRequest(
tags?: string[],
signal?: AbortSignal,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -21,7 +23,7 @@ export async function makeOpenAIImageGenerationRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -1,5 +1,6 @@
import { toast } from "@/lib/toast";
import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
export async function makeInteractionsRequest(
input: string,
@ -10,6 +11,7 @@ export async function makeInteractionsRequest(
signal?: AbortSignal,
customBaseUrl?: string,
previousInteractionId?: string,
customHeaders?: CustomHeaders,
): Promise<void> {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -27,10 +29,8 @@ export async function makeInteractionsRequest(
const headers: Record<string, string> = {
"Content-Type": "application/json",
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
...buildPlaygroundHeaders(tags, customHeaders),
};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const body: Record<string, unknown> = {
model: selectedModel,

View file

@ -48,6 +48,28 @@ describe("CodeSnippets", () => {
expect(code).toContain("print(response.data[0].embedding)");
});
describe("custom headers", () => {
const customHeaders = { "anthropic-beta": "context-1m-2025-08-07", "x-request-source": "playground" };
it("passes configured headers as default_headers on the OpenAI client", () => {
const code = generateCodeSnippet({ ...baseParams, endpointType: EndpointType.CHAT, customHeaders });
expect(code).toContain('base_url="http://localhost:4000",\n\tdefault_headers={');
expect(code).toContain('"anthropic-beta": "context-1m-2025-08-07"');
expect(code).toContain('"x-request-source": "playground"');
});
it("passes configured headers as default_headers on the Azure client", () => {
const code = generateCodeSnippet({ ...baseParams, selectedSdk: "azure", customHeaders });
expect(code).toContain('api_version="2024-02-01",\n\tdefault_headers={');
expect(code).toContain('"anthropic-beta": "context-1m-2025-08-07"');
});
it("omits default_headers when no custom headers are configured", () => {
expect(generateCodeSnippet(baseParams)).not.toContain("default_headers");
expect(generateCodeSnippet({ ...baseParams, customHeaders: {} })).not.toContain("default_headers");
});
});
describe("base URL selection", () => {
it("should use LITELLM_UI_API_DOC_BASE_URL when provided", () => {
const customBaseUrl = "https://custom-doc.example.com";

View file

@ -1,6 +1,7 @@
import { MessageType } from "./types";
import { EndpointType } from "./mode_endpoint_mapping";
import { MCPServer } from "@/components/mcp_tools/types";
import type { CustomHeaders } from "@/components/llm_calls/request_headers";
interface CodeGenMetadata {
tags?: string[];
@ -30,6 +31,7 @@ interface GenerateCodeParams {
PROXY_BASE_URL?: string;
LITELLM_UI_API_DOC_BASE_URL?: string | null;
};
customHeaders?: CustomHeaders;
}
export const generateCodeSnippet = (params: GenerateCodeParams): string => {
@ -48,6 +50,7 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
selectedModel,
selectedSdk,
proxySettings,
customHeaders,
} = params;
const effectiveApiKey = apiKeySource === "session" ? accessToken : apiKey;
@ -76,6 +79,11 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
const modelNameForCode = selectedModel || "your-model-name";
const defaultHeadersCode =
customHeaders && Object.keys(customHeaders).length > 0
? `,\n\tdefault_headers=${JSON.stringify(customHeaders, null, 2).replace(/\n/g, "\n\t")}`
: "";
const clientInitialization =
selectedSdk === "azure"
? `import openai
@ -83,13 +91,13 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
client = openai.AzureOpenAI(
api_key="${effectiveApiKey || "YOUR_LITELLM_API_KEY"}",
azure_endpoint="${apiBase}",
api_version="2024-02-01"
api_version="2024-02-01"${defaultHeadersCode}
)`
: `import openai
client = openai.OpenAI(
api_key="${effectiveApiKey || "YOUR_LITELLM_API_KEY"}",
base_url="${apiBase}"
base_url="${apiBase}"${defaultHeadersCode}
)`;
let endpointSpecificCode;

View file

@ -1,4 +1,5 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import openai from "openai";
import { makeOpenAIChatCompletionRequest } from "./chat_completion";
import type { TokenUsage } from "../chat_ui/ResponseMetrics";
@ -615,3 +616,47 @@ describe("chat_completion response cache", () => {
expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }));
});
});
describe("chat_completion custom headers", () => {
afterEach(() => {
vi.clearAllMocks();
});
it("sends custom headers alongside the tags header on the OpenAI client", async () => {
mockCreate.mockReturnValueOnce(nonStreamingResponse({ choices: [{ message: { content: "Hi" } }] }));
await makeOpenAIChatCompletionRequest(
[{ role: "user", content: "Hello" }],
vi.fn(),
"gpt-4",
"test-token",
["team-a"],
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
false,
{ "anthropic-beta": "context-1m-2025-08-07", "x-litellm-tags": "overridden" },
);
expect(vi.mocked(openai.OpenAI).mock.calls[0][0]).toMatchObject({
defaultHeaders: { "anthropic-beta": "context-1m-2025-08-07", "x-litellm-tags": "overridden" },
});
});
});

View file

@ -6,6 +6,7 @@ import { getProxyBaseUrl } from "@/components/networking";
import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types";
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
import { parseUsageCost } from "./usage_cost";
import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers";
const completionAsSingleChunk = (completion: ChatCompletion): ChatCompletionChunk =>
({
@ -50,6 +51,7 @@ export async function makeOpenAIChatCompletionRequest(
mockTestFallbacks?: boolean,
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -57,11 +59,7 @@ export async function makeOpenAIChatCompletionRequest(
console.log = function () {};
}
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
const client = new openai.OpenAI({
apiKey: accessToken,

View file

@ -0,0 +1,44 @@
import { describe, expect, it } from "vitest";
import { buildPlaygroundHeaders, customHeadersFromPairs, parseStoredHeaderPairs } from "./request_headers";
describe("customHeadersFromPairs", () => {
it("trims header names and drops rows without a name", () => {
expect(
customHeadersFromPairs([
[" anthropic-beta ", "context-1m-2025-08-07"],
["", "orphan value"],
[" ", "whitespace name"],
["x-empty", ""],
]),
).toEqual({ "anthropic-beta": "context-1m-2025-08-07", "x-empty": "" });
});
});
describe("parseStoredHeaderPairs", () => {
it("round-trips pairs persisted as JSON", () => {
const pairs = [["anthropic-beta", "context-1m-2025-08-07"]] as const;
expect(parseStoredHeaderPairs(JSON.stringify(pairs))).toEqual(pairs);
});
it("returns no pairs for missing, malformed, or wrongly shaped storage", () => {
expect(parseStoredHeaderPairs(null)).toEqual([]);
expect(parseStoredHeaderPairs("not json")).toEqual([]);
expect(parseStoredHeaderPairs(JSON.stringify({ "anthropic-beta": "x" }))).toEqual([]);
expect(parseStoredHeaderPairs(JSON.stringify([["ok", "pair"], ["one"], [1, 2], "str"]))).toEqual([["ok", "pair"]]);
});
});
describe("buildPlaygroundHeaders", () => {
it("joins tags into x-litellm-tags and lets custom headers override it", () => {
expect(buildPlaygroundHeaders(["a", "b"], { "x-custom": "1" })).toEqual({
"x-litellm-tags": "a,b",
"x-custom": "1",
});
expect(buildPlaygroundHeaders(["a"], { "x-litellm-tags": "b" })).toEqual({ "x-litellm-tags": "b" });
});
it("omits x-litellm-tags when there are no tags", () => {
expect(buildPlaygroundHeaders([], { "x-custom": "1" })).toEqual({ "x-custom": "1" });
expect(buildPlaygroundHeaders(undefined, undefined)).toEqual({});
});
});

View file

@ -0,0 +1,27 @@
import type { KeyValuePair } from "@/components/key_value_input";
export type CustomHeaders = Readonly<Record<string, string>>;
export const customHeadersFromPairs = (pairs: readonly KeyValuePair[]): CustomHeaders =>
Object.fromEntries(pairs.map(([name, value]) => [name.trim(), value]).filter(([name]) => name !== ""));
const isHeaderPair = (entry: unknown): entry is KeyValuePair =>
Array.isArray(entry) && entry.length === 2 && entry.every((part) => typeof part === "string");
export const parseStoredHeaderPairs = (raw: string | null): readonly KeyValuePair[] => {
if (!raw) return [];
try {
const parsed: unknown = JSON.parse(raw);
return Array.isArray(parsed) ? parsed.filter(isHeaderPair) : [];
} catch {
return [];
}
};
export const buildPlaygroundHeaders = (
tags?: readonly string[],
customHeaders?: CustomHeaders,
): Record<string, string> => ({
...(tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : {}),
...customHeaders,
});

View file

@ -1,4 +1,5 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import openai from "openai";
import { makeOpenAIResponsesRequest } from "./responses_api";
import { MessageType } from "../chat_ui/types";
import type { TokenUsage } from "../chat_ui/ResponseMetrics";
@ -611,3 +612,46 @@ describe("responses_api response cache", () => {
expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }), "");
});
});
describe("responses_api custom headers", () => {
afterEach(() => {
vi.clearAllMocks();
});
it("sends custom headers alongside the tags header on the OpenAI client", async () => {
mockResponsesCreate.mockReturnValueOnce(nonStreamingResponse({ id: "resp_1", output: [] }));
await makeOpenAIResponsesRequest(
[{ role: "user", content: "Hello" }],
vi.fn(),
"gpt-4",
"test-token",
["team-a"],
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
false,
undefined,
{ "anthropic-beta": "context-1m-2025-08-07" },
);
expect(vi.mocked(openai.OpenAI).mock.calls[0][0]).toMatchObject({
defaultHeaders: { "x-litellm-tags": "team-a", "anthropic-beta": "context-1m-2025-08-07" },
});
});
});

View file

@ -5,6 +5,7 @@ import { getProxyBaseUrl } from "@/components/networking";
import { toast } from "@/lib/toast";
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
import { parseUsageCost } from "./usage_cost";
import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers";
import type { MCPEvent } from "@/components/mcp_tools/types";
import { MCPServer, MCPToolset } from "@/components/mcp_tools/types";
import {
@ -85,6 +86,7 @@ export async function makeOpenAIResponsesRequest(
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
onTotalLatency?: (latency: number) => void,
customHeaders?: CustomHeaders,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -101,11 +103,7 @@ export async function makeOpenAIResponsesRequest(
}
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
const client = new openai.OpenAI({
apiKey: accessToken,