Merge pull request #13 from codebase/codex/acp-session-refresh

Fix OAuth refresh for long-lived ACP sessions

Author: halfaipg

package-lock.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",

package.json

@@ -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<Uint8Array>);
+
+		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([], {

src/acp/server.test.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<numb
 	async function newSession(params: acp.NewSessionRequest): Promise<acp.NewSessionResponse> {
 		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<numb
 		if (!selected) throw new Error(`unknown model: ${params.value}`);
 		if (session.bundle.model.id === selected.id) return { configOptions: [modelConfig(session)] };
 
+		await ensureFreshCredentials();
 		const previous = session.bundle;
 		const next = createAgent({
 			cwd: session.cwd,

src/acp/server.ts

@@ -73,6 +73,33 @@ describe("resolveConfig", () => {
 		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",

src/agent/config.test.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

src/agent/config.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");
+	});
+});

src/auth/ensure-fresh.test.ts

@@ -14,23 +14,22 @@ 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<void> {
 	const store = new CredentialsStore();
 	const creds = store.load();
 	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 {