Skip to content

Commit 3a9c7e0

Browse files
GregBrimbleGregBrimble
committed
feat: add provider-specific custom headers support
- Support provider-specific environment variables (OPENAI_CUSTOM_HEADERS, GROQ_CUSTOM_HEADERS, etc.) - Environment variables take precedence over provider config headers - Implemented in both TypeScript and Rust implementations - Added comprehensive test coverage for all supported providers - Header format: "key: value\nkey2: value2"
1 parent 4cb3c76 commit 3a9c7e0

File tree

13 files changed

+416
-11
lines changed

13 files changed

+416
-11
lines changed

‎codex-cli/src/utils/agent/agent-loop.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import {
2222
} from "../config.js";
2323
import { log } from "../logger/log.js";
2424
import { parseToolCallArguments } from "../parsers.js";
25+
import { providers, getCustomHeaders } from "../providers.js";
2526
import { responsesCreateViaChatCompletions } from "../responses.js";
2627
import {
2728
ORIGIN,
@@ -309,6 +310,12 @@ export class AgentLoop {
309310
const apiKey = this.config.apiKey ?? process.env["OPENAI_API_KEY"] ?? "";
310311
const baseURL = getBaseUrl(this.provider);
311312

313+
// Get custom headers from provider config and environment variables
314+
const providerInfo = providers[this.provider.toLowerCase()];
315+
const customHeaders = providerInfo
316+
? getCustomHeaders(providerInfo, this.provider.toLowerCase())
317+
: {};
318+
312319
this.oai = new OpenAI({
313320
// The OpenAI JS SDK only requires `apiKey` when making requests against
314321
// the official API. When running unit‑tests we stub out all network
@@ -326,6 +333,7 @@ export class AgentLoop {
326333
? { "OpenAI-Organization": OPENAI_ORGANIZATION }
327334
: {}),
328335
...(OPENAI_PROJECT ? { "OpenAI-Project": OPENAI_PROJECT } : {}),
336+
...customHeaders,
329337
},
330338
httpAgent: PROXY_URL ? new HttpsProxyAgent(PROXY_URL) : undefined,
331339
...(timeoutMs !== undefined ? { timeout: timeoutMs } : {}),
@@ -344,6 +352,7 @@ export class AgentLoop {
344352
? { "OpenAI-Organization": OPENAI_ORGANIZATION }
345353
: {}),
346354
...(OPENAI_PROJECT ? { "OpenAI-Project": OPENAI_PROJECT } : {}),
355+
...customHeaders,
347356
},
348357
httpAgent: PROXY_URL ? new HttpsProxyAgent(PROXY_URL) : undefined,
349358
...(timeoutMs !== undefined ? { timeout: timeoutMs } : {}),

‎codex-cli/src/utils/openai-client.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
OPENAI_ORGANIZATION,
99
OPENAI_PROJECT,
1010
} from "./config.js";
11+
import { providers, getCustomHeaders } from "./providers.js";
1112
import OpenAI, { AzureOpenAI } from "openai";
1213

1314
type OpenAIClientConfig = {
@@ -32,6 +33,14 @@ export function createOpenAIClient(
3233
headers["OpenAI-Project"] = OPENAI_PROJECT;
3334
}
3435

36+
// Add custom headers from provider config and environment variables.
37+
const providerName = config.provider?.toLowerCase() || "openai";
38+
const providerInfo = providers[providerName];
39+
if (providerInfo) {
40+
const customHeaders = getCustomHeaders(providerInfo, providerName);
41+
Object.assign(headers, customHeaders);
42+
}
43+
3544
if (config.provider?.toLowerCase() === "azure") {
3645
return new AzureOpenAI({
3746
apiKey: getApiKey(config.provider),

‎codex-cli/src/utils/providers.ts

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
1-
export const providers: Record<
2-
string,
3-
{ name: string; baseURL: string; envKey: string }
4-
> = {
1+
export interface ProviderInfo {
2+
name: string;
3+
baseURL: string;
4+
envKey: string;
5+
customHeaders?: Record<string, string>;
6+
}
7+
8+
export const providers: Record<string, ProviderInfo> = {
59
openai: {
610
name: "OpenAI",
711
baseURL: "https://api.openai.com/v1",
@@ -53,3 +57,62 @@ export const providers: Record<
5357
envKey: "ARCEEAI_API_KEY",
5458
},
5559
};
60+
61+
/**
62+
* Parse custom headers from provider-specific environment variable.
63+
* Format: "key: value\nkey2: value2"
64+
*/
65+
export function parseCustomHeadersFromEnv(
66+
providerName: string,
67+
): Record<string, string> {
68+
// Check for a PROVIDER-specific custom headers: e.g. OPENAI_CUSTOM_HEADERS or GROQ_CUSTOM_HEADERS.
69+
const envVarName = `${providerName.toUpperCase()}_CUSTOM_HEADERS`;
70+
const envHeaders = process.env[envVarName];
71+
if (!envHeaders) {
72+
return {};
73+
}
74+
75+
const headers: Record<string, string> = {};
76+
77+
for (const line of envHeaders.split("\n")) {
78+
const trimmedLine = line.trim();
79+
if (!trimmedLine) {
80+
continue;
81+
}
82+
83+
const colonIndex = trimmedLine.indexOf(":");
84+
if (colonIndex === -1) {
85+
continue;
86+
}
87+
88+
const key = trimmedLine.substring(0, colonIndex).trim();
89+
const value = trimmedLine.substring(colonIndex + 1).trim();
90+
91+
if (key) {
92+
headers[key] = value;
93+
}
94+
}
95+
96+
return headers;
97+
}
98+
99+
/**
100+
* Get merged custom headers from provider config and provider-specific environment variable.
101+
* Environment variable headers take precedence over provider config headers.
102+
*/
103+
export function getCustomHeaders(
104+
provider: ProviderInfo,
105+
providerName: string = "openai",
106+
): Record<string, string> {
107+
const headers: Record<string, string> = {};
108+
109+
// Add headers from provider config.
110+
if (provider.customHeaders) {
111+
Object.assign(headers, provider.customHeaders);
112+
}
113+
114+
// Add headers from provider-specific environment variable (takes precedence).
115+
Object.assign(headers, parseCustomHeadersFromEnv(providerName));
116+
117+
return headers;
118+
}

‎codex-cli/tests/custom-headers.test.ts

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
import { describe, it, expect, beforeEach, afterEach } from "vitest";
2+
import {
3+
parseCustomHeadersFromEnv,
4+
getCustomHeaders,
5+
type ProviderInfo,
6+
} from "../src/utils/providers.js";
7+
8+
describe("Custom Headers", () => {
9+
let originalEnvVars: Record<string, string | undefined> = {};
10+
11+
beforeEach(() => {
12+
// Store original environment variables
13+
const envVarNames = [
14+
"OPENAI_CUSTOM_HEADERS",
15+
"AZURE_CUSTOM_HEADERS",
16+
"TEST_CUSTOM_HEADERS",
17+
];
18+
for (const envVar of envVarNames) {
19+
originalEnvVars[envVar] = process.env[envVar];
20+
delete process.env[envVar];
21+
}
22+
});
23+
24+
afterEach(() => {
25+
// Restore original environment variables
26+
for (const [envVar, value] of Object.entries(originalEnvVars)) {
27+
if (value !== undefined) {
28+
process.env[envVar] = value;
29+
} else {
30+
delete process.env[envVar];
31+
}
32+
}
33+
originalEnvVars = {};
34+
});
35+
36+
describe("parseCustomHeadersFromEnv", () => {
37+
it("returns empty object when provider-specific env var is not set", () => {
38+
const headers = parseCustomHeadersFromEnv("openai");
39+
expect(headers).toEqual({});
40+
});
41+
42+
it("parses single header correctly", () => {
43+
process.env["OPENAI_CUSTOM_HEADERS"] = "X-Test-Header: test-value";
44+
const headers = parseCustomHeadersFromEnv("openai");
45+
expect(headers).toEqual({
46+
"X-Test-Header": "test-value",
47+
});
48+
});
49+
50+
it("parses multiple headers correctly", () => {
51+
process.env["AZURE_CUSTOM_HEADERS"] =
52+
"X-Custom-Auth: Bearer token123\nX-App-Version: 2.0.0\nX-Extra: some value";
53+
const headers = parseCustomHeadersFromEnv("azure");
54+
expect(headers).toEqual({
55+
"X-Custom-Auth": "Bearer token123",
56+
"X-App-Version": "2.0.0",
57+
"X-Extra": "some value",
58+
});
59+
});
60+
61+
it("handles headers with extra whitespace", () => {
62+
process.env["TEST_CUSTOM_HEADERS"] =
63+
" X-Header1 : value1 \n X-Header2:value2\n";
64+
const headers = parseCustomHeadersFromEnv("test");
65+
expect(headers).toEqual({
66+
"X-Header1": "value1",
67+
"X-Header2": "value2",
68+
});
69+
});
70+
71+
it("ignores empty lines", () => {
72+
process.env["TEST_CUSTOM_HEADERS"] =
73+
"X-Header1: value1\n\n\nX-Header2: value2\n";
74+
const headers = parseCustomHeadersFromEnv("test");
75+
expect(headers).toEqual({
76+
"X-Header1": "value1",
77+
"X-Header2": "value2",
78+
});
79+
});
80+
81+
it("ignores lines without colons", () => {
82+
process.env["TEST_CUSTOM_HEADERS"] =
83+
"X-Header1: value1\ninvalid-line-without-colon\nX-Header2: value2";
84+
const headers = parseCustomHeadersFromEnv("test");
85+
expect(headers).toEqual({
86+
"X-Header1": "value1",
87+
"X-Header2": "value2",
88+
});
89+
});
90+
91+
it("ignores lines with empty keys", () => {
92+
process.env["TEST_CUSTOM_HEADERS"] =
93+
"X-Header1: value1\n: empty-key\nX-Header2: value2";
94+
const headers = parseCustomHeadersFromEnv("test");
95+
expect(headers).toEqual({
96+
"X-Header1": "value1",
97+
"X-Header2": "value2",
98+
});
99+
});
100+
101+
it("handles values with colons", () => {
102+
process.env["TEST_CUSTOM_HEADERS"] =
103+
"X-Auth: Bearer token:with:colons\nX-URL: https://example.com:8080";
104+
const headers = parseCustomHeadersFromEnv("test");
105+
expect(headers).toEqual({
106+
"X-Auth": "Bearer token:with:colons",
107+
"X-URL": "https://example.com:8080",
108+
});
109+
});
110+
});
111+
112+
describe("getCustomHeaders", () => {
113+
it("returns empty object when provider has no custom headers and no env var", () => {
114+
const provider: ProviderInfo = {
115+
name: "Test Provider",
116+
baseURL: "https://api.example.com",
117+
envKey: "TEST_KEY",
118+
};
119+
const headers = getCustomHeaders(provider, "test");
120+
expect(headers).toEqual({});
121+
});
122+
123+
it("returns provider headers when no env var is set", () => {
124+
const provider: ProviderInfo = {
125+
name: "Test Provider",
126+
baseURL: "https://api.example.com",
127+
envKey: "TEST_KEY",
128+
customHeaders: {
129+
"X-Provider-Header": "provider-value",
130+
"X-Provider-Version": "1.0.0",
131+
},
132+
};
133+
const headers = getCustomHeaders(provider, "test");
134+
expect(headers).toEqual({
135+
"X-Provider-Header": "provider-value",
136+
"X-Provider-Version": "1.0.0",
137+
});
138+
});
139+
140+
it("returns env headers when provider has no custom headers", () => {
141+
process.env["OPENAI_CUSTOM_HEADERS"] =
142+
"X-Env-Header: env-value\nX-Env-Version: 2.0.0";
143+
const provider: ProviderInfo = {
144+
name: "OpenAI",
145+
baseURL: "https://api.openai.com/v1",
146+
envKey: "OPENAI_API_KEY",
147+
};
148+
const headers = getCustomHeaders(provider, "openai");
149+
expect(headers).toEqual({
150+
"X-Env-Header": "env-value",
151+
"X-Env-Version": "2.0.0",
152+
});
153+
});
154+
155+
it("merges provider and env headers with env taking precedence", () => {
156+
process.env["AZURE_CUSTOM_HEADERS"] =
157+
"X-Shared-Header: env-value\nX-Env-Only: env-only";
158+
const provider: ProviderInfo = {
159+
name: "AzureOpenAI",
160+
baseURL: "https://test.openai.azure.com/openai",
161+
envKey: "AZURE_OPENAI_API_KEY",
162+
customHeaders: {
163+
"X-Shared-Header": "provider-value",
164+
"X-Provider-Only": "provider-only",
165+
},
166+
};
167+
const headers = getCustomHeaders(provider, "azure");
168+
expect(headers).toEqual({
169+
"X-Shared-Header": "env-value", // env takes precedence
170+
"X-Provider-Only": "provider-only",
171+
"X-Env-Only": "env-only",
172+
});
173+
});
174+
});
175+
});

‎codex-rs/core/src/chat_completions.rs

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ pub(crate) async fn stream_chat_completions(
3434
model: &str,
3535
client: &reqwest::Client,
3636
provider: &ModelProviderInfo,
37+
provider_name: &str,
3738
) -> Result<ResponseStream> {
3839
// Build messages array
3940
let mut messages = Vec::<serde_json::Value>::new();
@@ -130,8 +131,14 @@ pub(crate) async fn stream_chat_completions(
130131
if let Some(api_key) = &api_key {
131132
req_builder = req_builder.bearer_auth(api_key.clone());
132133
}
134+
req_builder = req_builder.header(reqwest::header::ACCEPT, "text/event-stream");
135+
136+
// Add custom headers
137+
for (key, value) in provider.get_custom_headers(provider_name) {
138+
req_builder = req_builder.header(key, value);
139+
}
140+
133141
let res = req_builder
134-
.header(reqwest::header::ACCEPT, "text/event-stream")
135142
.json(&payload)
136143
.send()
137144
.await;

‎codex-rs/core/src/client.rs

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ pub struct ModelClient {
4343
model: String,
4444
client: reqwest::Client,
4545
provider: ModelProviderInfo,
46+
provider_name: String,
4647
effort: ReasoningEffortConfig,
4748
summary: ReasoningSummaryConfig,
4849
}
@@ -51,13 +52,15 @@ impl ModelClient {
5152
pub fn new(
5253
model: impl ToString,
5354
provider: ModelProviderInfo,
55+
provider_name: impl ToString,
5456
effort: ReasoningEffortConfig,
5557
summary: ReasoningSummaryConfig,
5658
) -> Self {
5759
Self {
5860
model: model.to_string(),
5961
client: reqwest::Client::new(),
6062
provider,
63+
provider_name: provider_name.to_string(),
6164
effort,
6265
summary,
6366
}
@@ -72,7 +75,7 @@ impl ModelClient {
7275
WireApi::Chat => {
7376
// Create the raw streaming connection first.
7477
let response_stream =
75-
stream_chat_completions(prompt, &self.model, &self.client, &self.provider)
78+
stream_chat_completions(prompt, &self.model, &self.client, &self.provider, &self.provider_name)
7679
.await?;
7780

7881
// Wrap it with the aggregation adapter so callers see *only*
@@ -136,12 +139,19 @@ impl ModelClient {
136139
instructions: None,
137140
})
138141
})?;
139-
let res = self
142+
let mut req_builder = self
140143
.client
141144
.post(&url)
142145
.bearer_auth(api_key)
143146
.header("OpenAI-Beta", "responses=experimental")
144-
.header(reqwest::header::ACCEPT, "text/event-stream")
147+
.header(reqwest::header::ACCEPT, "text/event-stream");
148+
149+
// Add custom headers
150+
for (key, value) in self.provider.get_custom_headers(&self.provider_name) {
151+
req_builder = req_builder.header(key, value);
152+
}
153+
154+
let res = req_builder
145155
.json(&payload)
146156
.send()
147157
.await;

0 commit comments

Comments
 (0)