From c2be3ab2354bcff277071010719e481ca2745bfc Mon Sep 17 00:00:00 2001 From: halfaipg Date: Wed, 29 Jul 2026 17:32:03 -0400 Subject: [PATCH 1/4] feat(search): use hosted search for signed-in sessions --- README.md | 5 +++ docs/MIGRATION_v1_to_v2.md | 2 +- src/agent/agent.ts | 6 ++++ src/agent/system-prompt.test.ts | 6 ++++ src/agent/system-prompt.ts | 3 ++ src/diagnostics/doctor.test.ts | 4 +-- src/diagnostics/doctor.ts | 13 ++++--- src/tools/types.ts | 8 +++++ src/tools/web-search.test.ts | 58 ++++++++++++++++++++++++++++++- src/tools/web-search.ts | 61 +++++++++++++++++++++++++++++---- 10 files changed, 151 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index e45e8fb..f760098 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,11 @@ codebase project build --wait "build a launch waitlist page" OAuth uses Codebase Auto by default (`codebase/d4f`, in-house DeepSeek V4 Flash). Swap models live with `/model`. Set reasoning depth with `/effort`. `project build` hands a prompt to the web builder and prints the session, status, event stream, and preview URL when you pass `--wait`. +Signed-in sessions also get hosted web search with no extra key. BYOK users can +set `TAVILY_API_KEY`, `BRAVE_API_KEY`, or `SEARXNG_URL`; local search providers +take priority over the hosted service. Run `codebase doctor` to verify which +path is active. + ## Use Codebase from an ACP client Codebase speaks the [Agent Client Protocol](https://agentclientprotocol.com/) over stdio: diff --git a/docs/MIGRATION_v1_to_v2.md b/docs/MIGRATION_v1_to_v2.md index 8cdc8d5..a65af36 100644 --- a/docs/MIGRATION_v1_to_v2.md +++ b/docs/MIGRATION_v1_to_v2.md @@ -29,7 +29,7 @@ and project memory all keep working. | `~/.codebase/config.json` | User config (theme, MCP servers, hooks) | Compatible. New fields are added; existing ones are read as before. | | `CLAUDE.md`, `AGENTS.md`, `CODEX.md`, `.cursorrules` | Project instructions | Identical pickup logic. | | `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `OPENAI_BASE_URL`, `OPENAI_MODEL` | Provider env vars | Identical. | -| `GLUE_*`, `TAVILY_API_KEY`, `BRAVE_API_KEY`, `SEARXNG_URL` | Sidecar + search keys | Identical. | +| `GLUE_*`, `TAVILY_API_KEY`, `BRAVE_API_KEY`, `SEARXNG_URL` | Sidecar + search keys | Still supported. Signed-in Codebase sessions can use hosted search without a separate search key. | | `CODEBASE_NOBOOT`, `CODEBASE_NOSOUND` | Behavior toggles | Identical. | If you sign in via codebase.foundation in v1, you stay signed in diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 28ef5ab..23447d2 100644 --- a/src/agent/agent.ts +++ b/src/agent/agent.ts @@ -324,6 +324,12 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle { const toolContext: ToolContext = { cwd, + ...(tokenManager && { + platform: { + baseUrl: (process.env.CODEBASE_AUTH_BASE_URL ?? "https://codebase.design").replace(/\/+$/, ""), + getAccessToken: () => tokenManager.getAccessToken(), + }, + }), fileStateCache: new FileStateCache(), tasks: new TaskStore({ cwd, taskListId: opts.taskListId ?? sessions.id }), userQueries, diff --git a/src/agent/system-prompt.test.ts b/src/agent/system-prompt.test.ts index 71c7821..5be074f 100644 --- a/src/agent/system-prompt.test.ts +++ b/src/agent/system-prompt.test.ts @@ -28,6 +28,12 @@ describe("buildSystemPrompt", () => { expect(out).toMatch(/Issue independent tool calls together/); }); + it("forbids presenting invented results after a failed web search", () => { + const out = buildSystemPrompt(); + expect(out).toMatch(/When web_search fails, no search occurred/); + expect(out).toMatch(/never invent.*links as search results/i); + }); + it("teaches subagent dispatch for fan-out work", () => { const out = buildSystemPrompt(); expect(out).toMatch(/dispatch_agent/); diff --git a/src/agent/system-prompt.ts b/src/agent/system-prompt.ts index c96b51e..8e3f088 100644 --- a/src/agent/system-prompt.ts +++ b/src/agent/system-prompt.ts @@ -60,6 +60,9 @@ export function buildSystemPrompt(opts: BuildSystemPromptOptions = {}): string { lines.push( "When a task fans out cleanly — multi-file audits, security reviews, broad codebase exploration — prefer dispatching subagents via dispatch_agent so each stream runs in parallel and their context stays out of your main loop. Don't also do the same searches yourself; that wastes turns and doubles the noise.", ); + lines.push( + "When web_search fails, no search occurred. Say that it failed and why; never invent, reconstruct from memory, or present plausible-looking links as search results.", + ); lines.push( 'Subagents come in types: "explore" (read-only, the default) for investigation, "general" for work that edits files or runs commands, plus any project-defined types. For parallel write work, give each general subagent isolation: "worktree" so their edits can\'t collide.', ); diff --git a/src/diagnostics/doctor.test.ts b/src/diagnostics/doctor.test.ts index 06e3009..9d5f272 100644 --- a/src/diagnostics/doctor.test.ts +++ b/src/diagnostics/doctor.test.ts @@ -52,7 +52,7 @@ describe("buildDoctorReport", () => { const out = buildDoctorReport({ cwd, dataRoot, - env: { TAVILY_API_KEY: "configured" } as NodeJS.ProcessEnv, + env: {}, model: { provider: "codebase", id: "d4f", name: "Codebase Auto" }, source: "proxy", mcpStatuses: [{ name: "db", connected: false, toolCount: 0, error: "spawn failed" }], @@ -66,7 +66,7 @@ describe("buildDoctorReport", () => { expect(out).toContain("config "); expect(out).toContain("is not valid JSON"); expect(out).toContain("mcp db: spawn failed"); - expect(out).toContain("✓ web_search configured"); + expect(out).toContain("✓ web_search configured (Codebase hosted)"); expect(out).toContain("sessions for this directory: 3"); expect(out).toContain("subagent types: general"); }); diff --git a/src/diagnostics/doctor.ts b/src/diagnostics/doctor.ts index d04cab8..6e6ec8a 100644 --- a/src/diagnostics/doctor.ts +++ b/src/diagnostics/doctor.ts @@ -87,11 +87,16 @@ export function buildDoctorReport(options: DoctorReportOptions): string[] { lines.push(check(s.connected, `mcp ${s.name}: ${s.toolCount} tools`, `mcp ${s.name}: ${s.error ?? "failed"}`)); } - const hasSearch = Boolean(env.TAVILY_API_KEY || env.BRAVE_API_KEY || env.SEARXNG_URL); + const localSearch = Boolean(env.TAVILY_API_KEY || env.BRAVE_API_KEY || env.SEARXNG_URL); + const hostedSearch = Boolean(creds && creds.source !== "byok"); lines.push( - hasSearch - ? check(true, "web_search configured", "") - : info("web_search unconfigured — set TAVILY_API_KEY, BRAVE_API_KEY, or SEARXNG_URL to enable"), + localSearch + ? check(true, "web_search configured (local provider)", "") + : hostedSearch + ? check(true, "web_search configured (Codebase hosted)", "") + : info( + "web_search unconfigured — sign in with `codebase auth login` or set TAVILY_API_KEY, BRAVE_API_KEY, or SEARXNG_URL", + ), ); lines.push( diff --git a/src/tools/types.ts b/src/tools/types.ts index f41793d..d699c3e 100644 --- a/src/tools/types.ts +++ b/src/tools/types.ts @@ -18,6 +18,14 @@ import type { TaskStore } from "./task-store.js"; */ export interface ToolContext { cwd: string; + /** + * Authenticated Codebase platform access. Present for OAuth/manual proxy + * sessions so tools can use hosted services without exposing service keys. + */ + platform?: { + baseUrl: string; + getAccessToken: () => Promise; + }; fileStateCache: FileStateCache; tasks: TaskStore; userQueries: UserQueryStore; diff --git a/src/tools/web-search.test.ts b/src/tools/web-search.test.ts index 90a3a1f..0edfca6 100644 --- a/src/tools/web-search.test.ts +++ b/src/tools/web-search.test.ts @@ -61,8 +61,14 @@ describe("web_search provider selection", () => { expect(pickProvider({ TAVILY_API_KEY: "x", BRAVE_API_KEY: "y" }).name).toBe("tavily"); }); + it("uses Codebase hosted search only when no local provider is configured", () => { + const hosted = { baseUrl: "https://codebase.design", getAccessToken: async () => "token" }; + expect(pickProvider({}, hosted).name).toBe("codebase"); + expect(pickProvider({ BRAVE_API_KEY: "x" }, hosted).name).toBe("brave"); + }); + it("errors with onboarding when nothing is configured", () => { - expect(() => pickProvider({})).toThrow(/TAVILY_API_KEY/); + expect(() => pickProvider({})).toThrow(/SEARCH_FAILED_NO_RESULTS/); }); }); @@ -140,6 +146,56 @@ describe("web_search end-to-end via mock servers", () => { expect(result.details.results[0].snippet).toBe("s1"); }); + it("uses authenticated Codebase hosted search when signed in", async () => { + routes["/api/v1/search"] = (req, res) => { + expect(req.method).toBe("POST"); + expect(req.headers.authorization).toBe("Bearer oauth-token"); + const chunks: Buffer[] = []; + req.on("data", (chunk) => chunks.push(chunk)); + req.on("end", () => { + expect(JSON.parse(Buffer.concat(chunks).toString("utf8"))).toEqual({ + query: "AIPG", + max_results: 4, + }); + res.setHeader("Content-Type", "application/json"); + res.end( + JSON.stringify({ + query: "AIPG", + provider: "tavily", + results: [ + { + title: "AI Power Grid", + url: "https://example.com/aipg", + snippet: "A sourced result.", + }, + ], + duration_ms: 12, + }), + ); + }); + }; + + 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, + getAccessToken: async () => "oauth-token", + }, + }).execute("call", { query: "AIPG", max_results: 4 }); + + expect(result.details.provider).toBe("codebase"); + expect(result.details.results[0]).toMatchObject({ + title: "AI Power Grid", + url: "https://example.com/aipg", + }); + } 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 index 85286a3..73f83e0 100644 --- a/src/tools/web-search.ts +++ b/src/tools/web-search.ts @@ -48,10 +48,11 @@ Provider auto-detected from env (highest-priority match wins): - TAVILY_API_KEY → Tavily (recommended; free tier available) - BRAVE_API_KEY → Brave Search - SEARXNG_URL → self-hosted SearXNG instance +- Codebase OAuth → hosted search (when signed in) -If none are set, this tool errors with onboarding instructions instead of attempting an unauthenticated fallback. Default 5 results, max 20.`; +Local providers take priority over hosted search. If none are available, this tool errors with onboarding instructions instead of attempting an unauthenticated fallback. Default 5 results, max 20.`; -export type ProviderName = "tavily" | "brave" | "searxng"; +export type ProviderName = "tavily" | "brave" | "searxng" | "codebase"; interface ProviderConfig { name: ProviderName; @@ -66,7 +67,12 @@ export interface ProviderEnv { SEARXNG_URL?: string; } -export function createWebSearch(_ctx: ToolContext): AgentTool { +export interface HostedSearchConfig { + baseUrl: string; + getAccessToken: () => Promise; +} + +export function createWebSearch(ctx: ToolContext): AgentTool { return { name: "web_search", label: "Search", @@ -74,7 +80,7 @@ export function createWebSearch(_ctx: ToolContext): AgentTool { - const provider = pickProvider(process.env); + const provider = pickProvider(process.env, ctx.platform); const max = params.max_results ?? DEFAULT_RESULTS; const timeout = params.timeout_ms ?? DEFAULT_TIMEOUT_MS; @@ -107,7 +113,7 @@ export function createWebSearch(_ctx: ToolContext): AgentTool { + const token = await config.getAccessToken(); + const res = await fetch(url, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + Accept: "application/json", + }, + body: JSON.stringify({ query, max_results: max }), + signal, + }); + const json = (await res.json().catch(() => ({}))) as { + error?: string; + error_description?: string; + results?: SearchResult[]; + }; + if (!res.ok) { + const reason = json.error_description ?? json.error ?? res.statusText; + throw new Error(`codebase ${res.status} ${reason}`); + } + return Array.isArray(json.results) + ? json.results.filter( + (result) => + typeof result?.title === "string" && + typeof result?.url === "string" && + typeof result?.snippet === "string", + ) + : []; + }, + }; +} + function tavilyProvider(apiKey: string, baseUrl?: string): ProviderConfig { const url = `${baseUrl ?? "https://api.tavily.com"}/search`; return { From 4c75628cbf05e0886b7be971ec761ee04efb0279 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Wed, 29 Jul 2026 18:20:24 -0400 Subject: [PATCH 2/4] chore(release): prepare 2.0.0-pre.84 --- package-lock.json | 4 ++-- package.json | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/package-lock.json b/package-lock.json index 9466e1c..865313b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "codebase-cli", - "version": "2.0.0-pre.83", + "version": "2.0.0-pre.84", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "codebase-cli", - "version": "2.0.0-pre.83", + "version": "2.0.0-pre.84", "license": "MIT", "dependencies": { "@agentclientprotocol/sdk": "^1.3.0", diff --git a/package.json b/package.json index 82b2fef..9fcdd01 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "codebase-cli", - "version": "2.0.0-pre.83", + "version": "2.0.0-pre.84", "description": "Codebase CLI — a TypeScript coding agent on the pi-mono runtime. OAuth-aware, any LLM provider, single install.", "keywords": [ "ai", From b7174350facb9171daca4da482874099ded3290b Mon Sep 17 00:00:00 2001 From: halfaipg Date: Wed, 29 Jul 2026 18:51:50 -0400 Subject: [PATCH 3/4] fix(auth): route manual keys with x-api-key --- src/agent/agent.ts | 26 +++++++-- src/agent/config.test.ts | 14 +++++ src/agent/config.ts | 105 +++++++++++++++++++++++------------ src/agent/model-list.test.ts | 17 ++++++ src/agent/model-list.ts | 11 ++-- src/tools/types.ts | 2 +- src/tools/web-search.test.ts | 32 ++++++++++- src/tools/web-search.ts | 6 +- 8 files changed, 160 insertions(+), 53 deletions(-) diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 23447d2..72dd2e5 100644 --- 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 = + 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 index 3beaca5..b87006a 100644 --- 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 index a716b60..b3c5b81 100644 --- a/src/agent/config.ts +++ b/src/agent/config.ts @@ -51,6 +51,8 @@ export interface ResolvedConfig { model: Model; 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 | undefined; if (!template) return null; const isDefault = !explicitModel; - const model: Model = { - ...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["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["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 | undefined; if (baseModel) { - const proxiedModel: Model = { ...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 | undefined; if (!template) return null; - const synthesized: Model = { - ...template, - id: modelId, - name: modelId, - baseUrl: proxyBase, - provider: explicitProvider as Model["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["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, + accessToken: string, + authMode: "bearer" | "api-key", +): Model { + 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 index a8a6786..fce8940 100644 --- 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).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 index 7911c95..fee6e50 100644 --- a/src/agent/model-list.ts +++ b/src/agent/model-list.ts @@ -24,7 +24,7 @@ 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. */ @@ -32,12 +32,15 @@ async function fetchOpenAiCompat( baseUrl: string, apiKey: string | undefined, provider: string, + modelHeaders: Record | undefined, signal?: AbortSignal, ): Promise { - const headers: Record = { Accept: "application/json" }; + const headers: Record = { ...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 index d699c3e..ec592e5 100644 --- a/src/tools/types.ts +++ b/src/tools/types.ts @@ -24,7 +24,7 @@ export interface ToolContext { */ platform?: { baseUrl: string; - getAccessToken: () => Promise; + getAuthHeaders: () => Promise>; }; fileStateCache: FileStateCache; tasks: TaskStore; diff --git a/src/tools/web-search.test.ts b/src/tools/web-search.test.ts index 0edfca6..27468e8 100644 --- 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 index 73f83e0..716b673 100644 --- 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; + getAuthHeaders: () => Promise>; } export function createWebSearch(ctx: ToolContext): AgentTool { @@ -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", }, From 5b8c0c96d5d631ea5554075f10453d73d5211bd4 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Wed, 29 Jul 2026 18:54:53 -0400 Subject: [PATCH 4/4] test(auth): prove manual key transport on the wire --- src/agent/proxy-auth-wire.test.ts | 77 +++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 src/agent/proxy-auth-wire.test.ts diff --git a/src/agent/proxy-auth-wire.test.ts b/src/agent/proxy-auth-wire.test.ts new file mode 100644 index 0000000..6d1ab09 --- /dev/null +++ b/src/agent/proxy-auth-wire.test.ts @@ -0,0 +1,77 @@ +import { mkdtempSync, rmSync } from "node:fs"; +import { type AddressInfo, createServer } from "node:http"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { completeSimple } from "@earendil-works/pi-ai"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { CredentialsStore } from "../auth/credentials.js"; +import { resolveConfig } from "./config.js"; + +describe("manual proxy authentication on the wire", () => { + const requests: Array> = []; + let server: ReturnType; + let baseUrl: string; + + beforeAll(async () => { + server = createServer((req, res) => { + requests.push(req.headers); + res.writeHead(200, { "Content-Type": "text/event-stream" }); + res.write( + `data: ${JSON.stringify({ + id: "chatcmpl-auth", + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + model: "d4f", + choices: [{ index: 0, delta: { role: "assistant", content: "authenticated" }, finish_reason: null }], + })}\n\n`, + ); + res.write( + `data: ${JSON.stringify({ + id: "chatcmpl-auth", + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + model: "d4f", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + })}\n\n`, + ); + res.end("data: [DONE]\n\n"); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address() as AddressInfo; + baseUrl = `http://127.0.0.1:${address.port}/inference`; + }); + + afterAll(async () => { + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())); + }); + }); + + it("sends X-API-Key through the real OpenAI-compatible transport", async () => { + const dataRoot = mkdtempSync(join(tmpdir(), "codebase-proxy-wire-")); + try { + const credentials = new CredentialsStore({ dataRoot }); + credentials.save({ + accessToken: "cb_test_wire", + scopes: ["inference"], + source: "manual", + }); + const config = resolveConfig({ + env: { CODEBASE_PROXY_BASE_URL: baseUrl }, + credentials, + }); + + const response = await completeSimple( + config.model, + { messages: [{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() }] }, + { apiKey: config.apiKey }, + ); + + expect(response.content).toContainEqual({ type: "text", text: "authenticated" }); + expect(requests).toHaveLength(1); + expect(requests[0]["x-api-key"]).toBe("cb_test_wire"); + } finally { + rmSync(dataRoot, { recursive: true, force: true }); + } + }); +});