fix(auth): route manual keys with x-api-key

halfaipg · halfaipg · commit b7174350facb · 2026-07-29T18:51:50.000-04:00
diff --git a/src/agent/agent.ts b/src/agent/agent.ts
@@ -82,7 +82,12 @@ export interface CreateAgentOptions {
 	 * faux provider so tests can run without env vars or credentials.
 	 * Production code never sets this.
 	 */
-	configOverride?: { model: ResolvedConfig["model"]; apiKey: string; source: ResolvedConfig["source"] };
+	configOverride?: {
+		model: ResolvedConfig["model"];
+		apiKey: string;
+		source: ResolvedConfig["source"];
+		proxyAuth?: ResolvedConfig["proxyAuth"];
+	};
 	/**
 	 * Runtime model override for proxy/OAuth sessions. Lets the user swap
 	 * models via /model without restarting. Format: `{ provider?, modelId }`.
@@ -184,15 +189,19 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle {
 	const effectiveOverride =
 		opts.modelOverride ??
 		(persistedModel?.modelId ? { provider: persistedModel.provider, modelId: persistedModel.modelId } : undefined);
-	const { model, apiKey, source } = opts.configOverride ?? resolveConfig({ modelOverride: effectiveOverride });
+	const credentials = new CredentialsStore();
+	const {
+		model,
+		apiKey,
+		source,
+		proxyAuth = "bearer",
+	} = opts.configOverride ?? resolveConfig({ credentials, modelOverride: effectiveOverride });
 
 	// OAuth-sourced credentials rotate ~hourly; build a refresh-aware getter
 	// so the agent never sends a stale access token after the first refresh
 	// window. BYOK / explicit / auto sources use the static key passed in.
 	const tokenManager =
-		source === "proxy"
-			? new TokenManager({ store: new CredentialsStore(), oauthConfig: defaultOAuthConfig() })
-			: null;
+		source === "proxy" ? new TokenManager({ store: credentials, oauthConfig: defaultOAuthConfig() }) : null;
 	const getApiKey = tokenManager ? () => tokenManager.getAccessToken() : () => apiKey;
 
 	const config = persistedConfig;
@@ -327,7 +336,12 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle {
 		...(tokenManager && {
 			platform: {
 				baseUrl: (process.env.CODEBASE_AUTH_BASE_URL ?? "https://codebase.design").replace(/\/+$/, ""),
-				getAccessToken: () => tokenManager.getAccessToken(),
+				getAuthHeaders: async () => {
+					const token = await tokenManager.getAccessToken();
+					const headers: Record<string, string> =
+						proxyAuth === "api-key" ? { "X-API-Key": token } : { Authorization: `Bearer ${token}` };
+					return headers;
+				},
 			},
 		}),
 		fileStateCache: new FileStateCache(),
diff --git a/src/agent/config.test.ts b/src/agent/config.test.ts
@@ -89,6 +89,20 @@ describe("resolveConfig", () => {
 		expect(config.model.provider).toBe("codebase");
 	});
 
+	it("uses X-API-Key transport for a manually entered Codebase API key", () => {
+		credentials.save({
+			accessToken: "cb_test_manual",
+			scopes: ["inference"],
+			source: "manual",
+		});
+
+		const config = resolveConfig({ env: {}, credentials });
+
+		expect(config.source).toBe("proxy");
+		expect(config.proxyAuth).toBe("api-key");
+		expect(config.model.headers).toMatchObject({ "X-API-Key": "cb_test_manual" });
+	});
+
 	it("does not use an expired manual token that cannot refresh", () => {
 		credentials.save({
 			accessToken: "expired-manual-token",
diff --git a/src/agent/config.ts b/src/agent/config.ts
@@ -51,6 +51,8 @@ export interface ResolvedConfig {
 	model: Model<string>;
 	apiKey: string;
 	source: "explicit" | "auto" | "proxy" | "byok";
+	/** Header contract for Codebase proxy credentials. OAuth uses Bearer; issued keys use X-API-Key. */
+	proxyAuth?: "bearer" | "api-key";
 }
 
 export class ConfigError extends Error {}
@@ -102,7 +104,12 @@ export function resolveConfig(envOrOpts: NodeJS.ProcessEnv | ResolveConfigOption
 			const byok = buildByokConfig(creds.provider as KnownProvider, creds.accessToken, override);
 			if (byok) return byok;
 		} else if (useProxy) {
-			const proxied = buildProxiedConfig(env, creds.accessToken, override);
+			const proxied = buildProxiedConfig(
+				env,
+				creds.accessToken,
+				creds.source === "manual" ? "api-key" : "bearer",
+				override,
+			);
 			if (proxied) return proxied;
 		}
 	}
@@ -202,6 +209,7 @@ export function resolveConfig(envOrOpts: NodeJS.ProcessEnv | ResolveConfigOption
 function buildProxiedConfig(
 	env: NodeJS.ProcessEnv,
 	accessToken: string,
+	authMode: "bearer" | "api-key",
 	override?: { provider?: string; modelId: string },
 ): ResolvedConfig | null {
 	// Runtime override (set via /model) wins over env vars wins over the default.
@@ -220,35 +228,39 @@ function buildProxiedConfig(
 		const template = getModel("groq", "llama-3.3-70b-versatile") as Model<string> | undefined;
 		if (!template) return null;
 		const isDefault = !explicitModel;
-		const model: Model<string> = {
-			...template,
-			// "Codebase Auto" — backend renamed this slot from MiniMax-M2.7
-			// to d4f (DeepSeek V4 Flash) when the underlying SGLang server
-			// was repointed. Keep this in sync with web/backend/providers/
-			// registry.js DEFAULT_MODEL.
-			id: explicitModel ?? "d4f",
-			name: isDefault ? "Codebase Auto" : (explicitModel ?? "Codebase Auto"),
-			baseUrl: proxyBase,
-			// Override provider so the status bar and /model don't lie about
-			// where this is served from. pi-ai uses `provider` mainly for
-			// display + a few baseUrl heuristics; the request body sends
-			// `model.id` only, so this cast is safe.
-			provider: "codebase" as Model<string>["provider"],
-			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-			// Keep this in sync with web/backend/providers/registry.js. The
-			// Groq llama template's context window is only a synthesis
-			// scaffold; compaction must budget against the actual d4f backend.
-			contextWindow: 131_072,
-		};
-		return { model, apiKey: accessToken, source: "proxy" };
+		const model = withProxyAuthHeaders(
+			{
+				...template,
+				// "Codebase Auto" — backend renamed this slot from MiniMax-M2.7
+				// to d4f (DeepSeek V4 Flash) when the underlying SGLang server
+				// was repointed. Keep this in sync with web/backend/providers/
+				// registry.js DEFAULT_MODEL.
+				id: explicitModel ?? "d4f",
+				name: isDefault ? "Codebase Auto" : (explicitModel ?? "Codebase Auto"),
+				baseUrl: proxyBase,
+				// Override provider so the status bar and /model don't lie about
+				// where this is served from. pi-ai uses `provider` mainly for
+				// display + a few baseUrl heuristics; the request body sends
+				// `model.id` only, so this cast is safe.
+				provider: "codebase" as Model<string>["provider"],
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+				// Keep this in sync with web/backend/providers/registry.js. The
+				// Groq llama template's context window is only a synthesis
+				// scaffold; compaction must budget against the actual d4f backend.
+				contextWindow: 131_072,
+			},
+			accessToken,
+			authMode,
+		);
+		return { model, apiKey: accessToken, source: "proxy", proxyAuth: authMode };
 	}
 
 	const modelId = explicitModel ?? DEFAULT_MODELS[explicitProvider];
 	if (!modelId) return null;
 	const baseModel = getModel(explicitProvider, modelId as never) as Model<string> | undefined;
 	if (baseModel) {
-		const proxiedModel: Model<string> = { ...baseModel, baseUrl: proxyBase };
-		return { model: proxiedModel, apiKey: accessToken, source: "proxy" };
+		const proxiedModel = withProxyAuthHeaders({ ...baseModel, baseUrl: proxyBase }, accessToken, authMode);
+		return { model: proxiedModel, apiKey: accessToken, source: "proxy", proxyAuth: authMode };
 	}
 	// pi-ai's registry doesn't know this provider+modelId combo (common for
 	// backend-only models the proxy exposes, e.g. "codebase:MiniMax-M2.7"
@@ -260,20 +272,39 @@ function buildProxiedConfig(
 	// didn't have native cost / context-window data for.
 	const template = getModel("groq", "llama-3.3-70b-versatile") as Model<string> | undefined;
 	if (!template) return null;
-	const synthesized: Model<string> = {
-		...template,
-		id: modelId,
-		name: modelId,
-		baseUrl: proxyBase,
-		provider: explicitProvider as Model<string>["provider"],
-		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
-		// Same reasoning as Codebase Auto: synthesized proxy models point
-		// at large-context backends. Without overriding contextWindow the
-		// Groq llama template's 128k leaks through and the compaction
-		// engine triggers ~30% earlier than it should.
-		contextWindow: guessContextWindow(modelId, 200_000),
+	const synthesized = withProxyAuthHeaders(
+		{
+			...template,
+			id: modelId,
+			name: modelId,
+			baseUrl: proxyBase,
+			provider: explicitProvider as Model<string>["provider"],
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+			// Same reasoning as Codebase Auto: synthesized proxy models point
+			// at large-context backends. Without overriding contextWindow the
+			// Groq llama template's 128k leaks through and the compaction
+			// engine triggers ~30% earlier than it should.
+			contextWindow: guessContextWindow(modelId, 200_000),
+		},
+		accessToken,
+		authMode,
+	);
+	return { model: synthesized, apiKey: accessToken, source: "proxy", proxyAuth: authMode };
+}
+
+function withProxyAuthHeaders(
+	model: Model<string>,
+	accessToken: string,
+	authMode: "bearer" | "api-key",
+): Model<string> {
+	if (authMode === "bearer") return model;
+	return {
+		...model,
+		headers: {
+			...model.headers,
+			"X-API-Key": accessToken,
+		},
 	};
-	return { model: synthesized, apiKey: accessToken, source: "proxy" };
 }
 
 /**
diff --git a/src/agent/model-list.test.ts b/src/agent/model-list.test.ts
@@ -32,6 +32,23 @@ describe("fetchAvailableModels", () => {
 		expect(out).toEqual([{ id: "d4f", name: "Codebase Auto", provider: "codebase" }]);
 	});
 
+	it("preserves model X-API-Key headers for Codebase-issued credentials", async () => {
+		const f = mockFetch(200, { models: [] });
+		vi.stubGlobal("fetch", f);
+		const manualModel = {
+			...model("codebase", "https://codebase.design/api/inference"),
+			headers: { "X-API-Key": "cb_test_manual" },
+		};
+
+		await fetchAvailableModels(manualModel, "cb_test_manual");
+
+		const headers = (f as unknown as ReturnType<typeof vi.fn>).mock.calls[0][1]?.headers;
+		expect(headers).toMatchObject({
+			"X-API-Key": "cb_test_manual",
+		});
+		expect(headers.Authorization).toBeUndefined();
+	});
+
 	it("hits Anthropic's endpoint and maps display_name", async () => {
 		const f = mockFetch(200, { data: [{ id: "claude-opus-4-1", display_name: "Claude Opus 4.1" }] });
 		vi.stubGlobal("fetch", f);
diff --git a/src/agent/model-list.ts b/src/agent/model-list.ts
@@ -24,20 +24,23 @@ export async function fetchAvailableModels(
 	if (provider === "google") return fetchGoogle(apiKey, signal);
 	const baseUrl = (model.baseUrl ?? "").replace(/\/+$/, "");
 	if (!baseUrl) return [];
-	return fetchOpenAiCompat(baseUrl, apiKey, provider, signal);
+	return fetchOpenAiCompat(baseUrl, apiKey, provider, model.headers, signal);
 }
 
 /** OpenAI-compatible `GET {baseUrl}/models`. Also parses our proxy's `{models:[…]}` shape. */
 async function fetchOpenAiCompat(
 	baseUrl: string,
 	apiKey: string | undefined,
 	provider: string,
+	modelHeaders: Record<string, string> | undefined,
 	signal?: AbortSignal,
 ): Promise<ModelOption[]> {
-	const headers: Record<string, string> = { Accept: "application/json" };
+	const headers: Record<string, string> = { ...modelHeaders, Accept: "application/json" };
 	// Local servers (Ollama/LM Studio) usually need no auth; only send a key
-	// when we have a real one.
-	if (apiKey) headers.Authorization = `Bearer ${apiKey}`;
+	// when we have a real one. A model-level X-API-Key means this is a
+	// manually entered Codebase key, so do not duplicate it as Bearer.
+	const hasApiKeyHeader = Object.keys(headers).some((name) => name.toLowerCase() === "x-api-key");
+	if (apiKey && !hasApiKeyHeader) headers.Authorization = `Bearer ${apiKey}`;
 	const res = await fetch(`${baseUrl}/models`, { headers, signal });
 	if (!res.ok) throw new Error(`${res.status} ${res.statusText}`);
 	const json = (await res.json()) as { data?: unknown; models?: unknown };
diff --git a/src/tools/types.ts b/src/tools/types.ts
@@ -24,7 +24,7 @@ export interface ToolContext {
 	 */
 	platform?: {
 		baseUrl: string;
-		getAccessToken: () => Promise<string>;
+		getAuthHeaders: () => Promise<Record<string, string>>;
 	};
 	fileStateCache: FileStateCache;
 	tasks: TaskStore;
diff --git a/src/tools/web-search.test.ts b/src/tools/web-search.test.ts
@@ -62,7 +62,10 @@ describe("web_search provider selection", () => {
 	});
 
 	it("uses Codebase hosted search only when no local provider is configured", () => {
-		const hosted = { baseUrl: "https://codebase.design", getAccessToken: async () => "token" };
+		const hosted = {
+			baseUrl: "https://codebase.design",
+			getAuthHeaders: async () => ({ Authorization: "Bearer token" }),
+		};
 		expect(pickProvider({}, hosted).name).toBe("codebase");
 		expect(pickProvider({ BRAVE_API_KEY: "x" }, hosted).name).toBe("brave");
 	});
@@ -182,7 +185,7 @@ describe("web_search end-to-end via mock servers", () => {
 				...makeCtx(),
 				platform: {
 					baseUrl,
-					getAccessToken: async () => "oauth-token",
+					getAuthHeaders: async () => ({ Authorization: "Bearer oauth-token" }),
 				},
 			}).execute("call", { query: "AIPG", max_results: 4 });
 
@@ -196,6 +199,31 @@ describe("web_search end-to-end via mock servers", () => {
 		}
 	});
 
+	it("uses X-API-Key for hosted search with a manually entered Codebase key", async () => {
+		routes["/api/v1/search"] = (req, res) => {
+			expect(req.headers["x-api-key"]).toBe("cb_test_manual");
+			expect(req.headers.authorization).toBeUndefined();
+			res.setHeader("Content-Type", "application/json");
+			res.end(JSON.stringify({ results: [] }));
+		};
+
+		const original = { ...process.env };
+		for (const key of ["TAVILY_API_KEY", "BRAVE_API_KEY", "SEARXNG_URL"]) delete process.env[key];
+		try {
+			const result = await createWebSearch({
+				...makeCtx(),
+				platform: {
+					baseUrl,
+					getAuthHeaders: async () => ({ "X-API-Key": "cb_test_manual" }),
+				},
+			}).execute("call", { query: "manual auth" });
+
+			expect(result.details.provider).toBe("codebase");
+		} finally {
+			Object.assign(process.env, original);
+		}
+	});
+
 	it("surfaces provider HTTP errors with the status code", async () => {
 		routes["/search"] = (_req, res) => {
 			res.statusCode = 401;
diff --git a/src/tools/web-search.ts b/src/tools/web-search.ts
@@ -69,7 +69,7 @@ export interface ProviderEnv {
 
 export interface HostedSearchConfig {
 	baseUrl: string;
-	getAccessToken: () => Promise<string>;
+	getAuthHeaders: () => Promise<Record<string, string>>;
 }
 
 export function createWebSearch(ctx: ToolContext): AgentTool<typeof Params, WebSearchDetails> {
@@ -138,11 +138,11 @@ function codebaseProvider(config: HostedSearchConfig): ProviderConfig {
 	return {
 		name: "codebase",
 		search: async (query, max, signal) => {
-			const token = await config.getAccessToken();
+			const authHeaders = await config.getAuthHeaders();
 			const res = await fetch(url, {
 				method: "POST",
 				headers: {
-					Authorization: `Bearer ${token}`,
+					...authHeaders,
 					"Content-Type": "application/json",
 					Accept: "application/json",
 				},
