From 56a0e5a262049b092f03839799dd65fea6c95719 Mon Sep 17 00:00:00 2001 From: halfaipg Date: Wed, 29 Jul 2026 16:57:57 -0400 Subject: [PATCH] Fix OAuth refresh for long-lived ACP sessions --- package-lock.json | 4 +-- package.json | 2 +- src/acp/server.test.ts | 47 ++++++++++++++++++++++++++++++++++- src/acp/server.ts | 3 +++ src/agent/config.test.ts | 27 ++++++++++++++++++++ src/agent/config.ts | 9 ++++++- src/auth/ensure-fresh.test.ts | 46 ++++++++++++++++++++++++++++++++++ src/auth/ensure-fresh.ts | 13 +++++----- 8 files changed, 139 insertions(+), 12 deletions(-) create mode 100644 src/auth/ensure-fresh.test.ts diff --git a/package-lock.json b/package-lock.json index 2442261..9466e1c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "codebase-cli", - "version": "2.0.0-pre.82", + "version": "2.0.0-pre.83", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "codebase-cli", - "version": "2.0.0-pre.82", + "version": "2.0.0-pre.83", "license": "MIT", "dependencies": { "@agentclientprotocol/sdk": "^1.3.0", diff --git a/package.json b/package.json index 1aaed78..82b2fef 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "codebase-cli", - "version": "2.0.0-pre.82", + "version": "2.0.0-pre.83", "description": "Codebase CLI — a TypeScript coding agent on the pi-mono runtime. OAuth-aware, any LLM provider, single install.", "keywords": [ "ai", diff --git a/src/acp/server.test.ts b/src/acp/server.test.ts index 926d599..04d63a5 100644 --- a/src/acp/server.test.ts +++ b/src/acp/server.test.ts @@ -6,7 +6,8 @@ import { fileURLToPath } from "node:url"; import * as acp from "@agentclientprotocol/sdk"; import type { Model } from "@earendil-works/pi-ai"; import { fauxAssistantMessage, fauxToolCall, registerFauxProvider } from "@earendil-works/pi-ai"; -import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { CredentialsStore } from "../auth/credentials.js"; import { runAcpServer } from "./server.js"; const MOCK_MCP = fileURLToPath(new URL("../mcp/__test__/mock-server.mjs", import.meta.url)); @@ -136,6 +137,50 @@ describe("runAcpServer", () => { await serverDone; }); + it("refreshes OAuth credentials before a long-lived server creates a new session", async () => { + const store = new CredentialsStore(); + store.save({ + accessToken: "access-current", + refreshToken: "refresh-current", + expiresAt: Date.now() + 4 * 60_000, + scopes: ["inference"], + source: "codebase", + }); + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValueOnce( + new Response(JSON.stringify({ access_token: "access-new", expires_in: 3600 }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); + + const toAgent = new PassThrough(); + const fromAgent = new PassThrough(); + const serverDone = runAcpServer({ + stdin: toAgent, + stdout: fromAgent, + configOverride: { model, apiKey: "faux-key", source: "byok" }, + }); + const stream = acp.ndJsonStream(Writable.toWeb(toAgent), Readable.toWeb(fromAgent) as ReadableStream); + + await acp.client({ name: "buzz-refresh-test" }).connectWith(stream, async (client) => { + await client.request(acp.methods.agent.initialize, { + protocolVersion: acp.PROTOCOL_VERSION, + clientInfo: { name: "buzz-acp", version: "test" }, + }); + await client.request(acp.methods.agent.session.new, { + cwd, + mcpServers: [], + }); + }); + + const refreshCalls = fetchSpy.mock.calls.filter(([input]) => String(input).endsWith("/api/oauth/token")); + expect(refreshCalls).toHaveLength(1); + expect(store.load()?.accessToken).toBe("access-new"); + + toAgent.end(); + await serverDone; + }); + it("streams provider failures as visible ACP messages", async () => { faux.setResponses([ fauxAssistantMessage([], { diff --git a/src/acp/server.ts b/src/acp/server.ts index 889a678..5eea108 100644 --- a/src/acp/server.ts +++ b/src/acp/server.ts @@ -7,6 +7,7 @@ import type { AgentEvent } from "@earendil-works/pi-agent-core"; import type { ImageContent } from "@earendil-works/pi-ai"; import { type AgentBundle, type CreateAgentOptions, createAgent } from "../agent/agent.js"; import type { ModelOption } from "../agent/model-list.js"; +import { ensureFreshCredentials } from "../auth/ensure-fresh.js"; import { latestAssistantError, userFacingErrorMessage } from "../errors/user-facing.js"; import type { NamedServer } from "../mcp/config.js"; import type { PermissionRequest, ResponseChoice } from "../permissions/store.js"; @@ -72,6 +73,7 @@ export async function runAcpServer(options: AcpServerOptions = {}): Promise { if (!isAbsolute(params.cwd)) throw new Error("ACP session cwd must be an absolute path"); + await ensureFreshCredentials(); const bundle = createAgent({ cwd: params.cwd, autoApprove: false, @@ -115,6 +117,7 @@ export async function runAcpServer(options: AcpServerOptions = {}): Promise { expect(config.model.contextWindow).toBe(131_072); }); + it("constructs a proxy agent from expired OAuth credentials that can refresh", () => { + credentials.save({ + accessToken: "expired-oauth-token", + refreshToken: "refresh-token", + expiresAt: Date.now() - 60_000, + scopes: ["inference"], + source: "codebase", + }); + + const config = resolveConfig({ env: {}, credentials }); + + expect(config.source).toBe("proxy"); + expect(config.apiKey).toBe("expired-oauth-token"); + expect(config.model.provider).toBe("codebase"); + }); + + it("does not use an expired manual token that cannot refresh", () => { + credentials.save({ + accessToken: "expired-manual-token", + expiresAt: Date.now() - 60_000, + scopes: ["inference"], + source: "manual", + }); + + expect(() => resolveConfig({ env: {}, credentials })).toThrow(); + }); + it("treats provider=codebase model overrides as proxy-routed model ids", () => { credentials.save({ accessToken: "oauth-token", diff --git a/src/agent/config.ts b/src/agent/config.ts index cf76e3b..a716b60 100644 --- a/src/agent/config.ts +++ b/src/agent/config.ts @@ -78,7 +78,14 @@ export function resolveConfig(envOrOpts: NodeJS.ProcessEnv | ResolveConfigOption // - byok → direct call against the provider's own API const useProxy = env.CODEBASE_DISABLE_PROXY !== "1"; const creds = credentials.load(); - if (creds && !credentials.isExpired(creds)) { + // Refreshable OAuth credentials remain sufficient to construct a proxy + // agent even when the current access token is wire-dead. The agent's + // TokenManager refreshes before the first provider request. Rejecting the + // credentials here creates a chicken-and-egg for long-lived ACP servers: + // they cannot create the next session that would perform the refresh. + const usableCredentials = + creds && (!credentials.isExpired(creds) || (creds.source === "codebase" && Boolean(creds.refreshToken))); + if (usableCredentials) { if (creds.source === "byok" && creds.provider === "openai-compat" && creds.baseUrl) { // Custom Chat Completions endpoint saved by the wizard (Ollama, // LM Studio, vLLM, …). Same synthesis as the OPENAI_BASE_URL env diff --git a/src/auth/ensure-fresh.test.ts b/src/auth/ensure-fresh.test.ts new file mode 100644 index 0000000..e04a89a --- /dev/null +++ b/src/auth/ensure-fresh.test.ts @@ -0,0 +1,46 @@ +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { CredentialsStore } from "./credentials.js"; +import { ensureFreshCredentials } from "./ensure-fresh.js"; + +describe("ensureFreshCredentials", () => { + let home: string; + let previousHome: string | undefined; + + beforeEach(() => { + previousHome = process.env.HOME; + home = mkdtempSync(join(tmpdir(), "ensure-fresh-")); + process.env.HOME = home; + }); + + afterEach(() => { + if (previousHome === undefined) delete process.env.HOME; + else process.env.HOME = previousHome; + rmSync(home, { recursive: true, force: true }); + vi.restoreAllMocks(); + }); + + it("refreshes within the five-minute preflight window", async () => { + const store = new CredentialsStore(); + store.save({ + accessToken: "access-current", + refreshToken: "refresh-current", + expiresAt: Date.now() + 4 * 60_000, + scopes: ["inference"], + source: "codebase", + }); + const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValueOnce( + new Response(JSON.stringify({ access_token: "access-new", expires_in: 3600 }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); + + await ensureFreshCredentials(); + + expect(fetchSpy).toHaveBeenCalledTimes(1); + expect(store.load()?.accessToken).toBe("access-new"); + }); +}); diff --git a/src/auth/ensure-fresh.ts b/src/auth/ensure-fresh.ts index 07e6847..cd308fc 100644 --- a/src/auth/ensure-fresh.ts +++ b/src/auth/ensure-fresh.ts @@ -14,15 +14,15 @@ import { TokenManager } from "./token-manager.js"; * succeeds. That's a chicken/egg the previous code didn't notice * because a fresh login + fresh process never tripped it. * - * Safe to call on every launch. No-ops when: + * Safe to call before every launch or long-lived server session. No-ops when: * - no credentials exist (wizard handles it) * - source isn't "codebase" (BYOK / env-key sessions don't refresh) - * - access token is still good (TokenManager's preemptive window) + * - access token is outside TokenManager's preemptive refresh window * - * Errors are NOT thrown: a network blip during cold start should fall - * through to `resolveConfig`, which will produce the existing - * "please run `codebase auth login`" path. The user sees one error, - * not two layered ones. + * Errors are NOT thrown: a network blip during preflight should not prevent + * a refreshable proxy session from being created. The agent's TokenManager + * retries at the provider-call boundary, where callers can render the + * canonical sign-in or network error instead of an opaque startup failure. */ export async function ensureFreshCredentials(): Promise { const store = new CredentialsStore(); @@ -30,7 +30,6 @@ export async function ensureFreshCredentials(): Promise { if (!creds) return; if (creds.source !== "codebase") return; if (!creds.refreshToken) return; - if (!store.isExpired(creds)) return; const manager = new TokenManager({ store, oauthConfig: defaultOAuthConfig() }); try {