fix(auth): refresh expired OAuth access tokens on demand

halfaipg · halfaipg · commit 8ade465f02b2 · 2026-05-13T13:12:33.000-04:00
The agent and glue clients both captured the access token by value at
startup, so once the OAuth token expired (~1h after launch) every
subsequent request kept sending the stale token and the proxy returned
401 invalid_token. Manual \`codebase auth refresh\` was the only way out.

Introduce TokenManager: a read-through getter that loads credentials
from disk on each call, checks expiry with a 60s skew, and refreshes
via the existing refresh-token flow before returning. Single-flight
lock prevents concurrent requests from firing parallel refreshes.

Wire the token manager into agent.ts (when source === "proxy") and
make GlueClient take a getApiKey callback instead of a captured key.
BYOK / explicit-env paths keep the static key — those don't expire.
diff --git a/src/agent/agent.ts b/src/agent/agent.ts
@@ -2,6 +2,9 @@ import { homedir } from "node:os";
 import { isAbsolute, join, resolve } from "node:path";
 import { Agent, type AgentEvent } from "@earendil-works/pi-agent-core";
 import type { Model } from "@earendil-works/pi-ai";
+import { defaultOAuthConfig } from "../auth/cli.js";
+import { CredentialsStore } from "../auth/credentials.js";
+import { TokenManager } from "../auth/token-manager.js";
 import { CompactionEngine } from "../compaction/engine.js";
 import { CompactionMonitor } from "../compaction/monitor.js";
 import { ConfigStore } from "../config/store.js";
@@ -92,6 +95,15 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle {
 	const cwd = opts.cwd ?? process.cwd();
 	const systemPrompt = opts.systemPrompt ?? buildSystemPrompt(cwd);
 
+	// 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;
+	const getApiKey = tokenManager ? () => tokenManager.getAccessToken() : () => apiKey;
+
 	const config = new ConfigStore({ cwd });
 	const permissions = new PermissionStore({
 		allowPatterns: config.allowPatterns(),
@@ -109,7 +121,7 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle {
 	const glue = new GlueClient({
 		fastModel: glueModels.fast,
 		smartModel: glueModels.smart,
-		apiKey: glueModels.apiKey,
+		getApiKey,
 	});
 	const compaction = new CompactionEngine({ glue, modelId: model.id });
 	const compactionMonitor = new CompactionMonitor();
@@ -127,7 +139,7 @@ export function createAgent(opts: CreateAgentOptions = {}): AgentBundle {
 		spawnSubagent: ({ systemPrompt: subPrompt, tools: subTools }) =>
 			new Agent({
 				initialState: { model, systemPrompt: subPrompt, tools: subTools },
-				getApiKey: () => apiKey,
+				getApiKey,
 			}),
 	};
 
diff --git a/src/auth/cli.ts b/src/auth/cli.ts
@@ -20,7 +20,7 @@ const DEFAULT_AUTH_BASE = "https://codebase.design";
  * codebase.design and codebase.foundation both alias to the same Next
  * app; the user-facing brand is .design so it's the default.
  */
-function defaultOAuthConfig(env: NodeJS.ProcessEnv = process.env): OAuthConfig {
+export function defaultOAuthConfig(env: NodeJS.ProcessEnv = process.env): OAuthConfig {
 	const base = (env.CODEBASE_AUTH_BASE_URL ?? DEFAULT_AUTH_BASE).replace(/\/+$/, "");
 	return {
 		authorizationUrl: `${base}/login`,
diff --git a/src/auth/token-manager.test.ts b/src/auth/token-manager.test.ts
@@ -0,0 +1,140 @@
+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 { type Credentials, CredentialsStore } from "./credentials.js";
+import type { OAuthConfig } from "./flow.js";
+import { TokenManager } from "./token-manager.js";
+
+const OAUTH: OAuthConfig = {
+	authorizationUrl: "https://example.test/login",
+	tokenUrl: "https://example.test/token",
+	refreshUrl: "https://example.test/refresh",
+	revokeUrl: "https://example.test/revoke",
+	clientId: "test-client",
+	scopes: ["read", "write"],
+};
+
+function freshCreds(overrides: Partial<Credentials> = {}): Omit<Credentials, "version"> {
+	return {
+		accessToken: "access-current",
+		refreshToken: "refresh-current",
+		expiresAt: Date.now() + 60 * 60 * 1000,
+		scopes: ["read", "write"],
+		source: "codebase",
+		userId: "u-1",
+		email: "u@example.test",
+		...overrides,
+	};
+}
+
+describe("TokenManager", () => {
+	let dataRoot: string;
+	let store: CredentialsStore;
+
+	beforeEach(() => {
+		dataRoot = mkdtempSync(join(tmpdir(), "tm-test-"));
+		store = new CredentialsStore({ dataRoot });
+	});
+
+	afterEach(() => {
+		rmSync(dataRoot, { recursive: true, force: true });
+		vi.restoreAllMocks();
+	});
+
+	it("returns the cached access token when not near expiry", async () => {
+		store.save(freshCreds());
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		const fetchSpy = vi.spyOn(globalThis, "fetch");
+		await expect(tm.getAccessToken()).resolves.toBe("access-current");
+		expect(fetchSpy).not.toHaveBeenCalled();
+	});
+
+	it("refreshes and persists when the stored token is within the skew window", async () => {
+		store.save(freshCreds({ expiresAt: Date.now() + 10_000 })); // 10s left, under 60s skew
+		vi.spyOn(globalThis, "fetch").mockResolvedValueOnce(
+			new Response(JSON.stringify({ access_token: "access-NEW", refresh_token: "refresh-NEW", expires_in: 3600 }), {
+				status: 200,
+				headers: { "Content-Type": "application/json" },
+			}),
+		);
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		await expect(tm.getAccessToken()).resolves.toBe("access-NEW");
+		const persisted = store.load();
+		expect(persisted?.accessToken).toBe("access-NEW");
+		expect(persisted?.refreshToken).toBe("refresh-NEW");
+		// Metadata that the refresh endpoint doesn't echo back is preserved.
+		expect(persisted?.email).toBe("u@example.test");
+		expect(persisted?.source).toBe("codebase");
+	});
+
+	it("preserves the existing refresh token when the refresh response omits it", async () => {
+		store.save(freshCreds({ expiresAt: Date.now() + 5_000 }));
+		vi.spyOn(globalThis, "fetch").mockResolvedValueOnce(
+			new Response(JSON.stringify({ access_token: "access-NEW", expires_in: 3600 }), {
+				status: 200,
+				headers: { "Content-Type": "application/json" },
+			}),
+		);
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		await tm.getAccessToken();
+		expect(store.load()?.refreshToken).toBe("refresh-current");
+	});
+
+	it("single-flights concurrent refresh calls into one network round-trip", async () => {
+		store.save(freshCreds({ expiresAt: Date.now() + 5_000 }));
+		const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation(
+			() =>
+				new Promise((resolve) => {
+					setTimeout(
+						() =>
+							resolve(
+								new Response(JSON.stringify({ access_token: "access-NEW", expires_in: 3600 }), {
+									status: 200,
+									headers: { "Content-Type": "application/json" },
+								}),
+							),
+						20,
+					);
+				}),
+		);
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		const results = await Promise.all([tm.getAccessToken(), tm.getAccessToken(), tm.getAccessToken()]);
+		expect(results).toEqual(["access-NEW", "access-NEW", "access-NEW"]);
+		expect(fetchSpy).toHaveBeenCalledTimes(1);
+	});
+
+	it("subsequent calls AFTER a refresh resolve hit the cached path again", async () => {
+		store.save(freshCreds({ expiresAt: Date.now() + 5_000 }));
+		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 tm = new TokenManager({ store, oauthConfig: OAUTH });
+		await tm.getAccessToken();
+		// Now the stored token has a fresh 1h lifetime — should not re-refresh.
+		await tm.getAccessToken();
+		expect(fetchSpy).toHaveBeenCalledTimes(1);
+	});
+
+	it("throws a sign-in-needed error when no credentials exist", async () => {
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		await expect(tm.getAccessToken()).rejects.toThrow(/codebase auth login/);
+	});
+
+	it("throws when the token expired but no refresh token was saved", async () => {
+		store.save(freshCreds({ expiresAt: Date.now() + 5_000, refreshToken: undefined }));
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		await expect(tm.getAccessToken()).rejects.toThrow(/codebase auth login/);
+	});
+
+	it("treats undefined expiresAt as never-expires (no refresh attempted)", async () => {
+		store.save(freshCreds({ expiresAt: undefined }));
+		const fetchSpy = vi.spyOn(globalThis, "fetch");
+		const tm = new TokenManager({ store, oauthConfig: OAUTH });
+		await expect(tm.getAccessToken()).resolves.toBe("access-current");
+		expect(fetchSpy).not.toHaveBeenCalled();
+	});
+});
diff --git a/src/auth/token-manager.ts b/src/auth/token-manager.ts
@@ -0,0 +1,85 @@
+import type { Credentials, CredentialsStore } from "./credentials.js";
+import { type OAuthConfig, refreshAccessToken } from "./flow.js";
+
+export interface TokenManagerOptions {
+	store: CredentialsStore;
+	oauthConfig: OAuthConfig;
+	/**
+	 * Refresh when the access token's remaining lifetime falls below this
+	 * many milliseconds. Default 60s — generous enough to ride out clock
+	 * skew + a slow refresh round-trip without ever sending an expired
+	 * token over the wire.
+	 */
+	refreshSkewMs?: number;
+}
+
+/**
+ * Read-through, refresh-aware accessor for the OAuth access token.
+ *
+ * Pi-mono's `getApiKey` runs on every API call, so we have to be fast in
+ * the common case (token still valid). Slow path: refresh + persist before
+ * returning the new token. A single in-flight promise (`pending`) prevents
+ * a burst of concurrent calls from firing parallel refreshes at the same
+ * moment — they all await the one refresh that's already running.
+ */
+export class TokenManager {
+	private readonly store: CredentialsStore;
+	private readonly oauthConfig: OAuthConfig;
+	private readonly refreshSkewMs: number;
+	private pending: Promise<string> | null = null;
+
+	constructor(options: TokenManagerOptions) {
+		this.store = options.store;
+		this.oauthConfig = options.oauthConfig;
+		this.refreshSkewMs = options.refreshSkewMs ?? 60_000;
+	}
+
+	/**
+	 * Return a valid access token, refreshing if the stored one is within
+	 * `refreshSkewMs` of expiry. Throws when no credentials exist or when
+	 * a refresh fails (network or expired refresh token); the caller turns
+	 * that into a user-facing "please run `codebase auth login`" message.
+	 */
+	async getAccessToken(): Promise<string> {
+		const creds = this.store.load();
+		if (!creds) throw new Error("not signed in — run `codebase auth login`");
+		if (!this.needsRefresh(creds)) return creds.accessToken;
+		if (!creds.refreshToken) {
+			throw new Error("access token expired and no refresh token saved — run `codebase auth login`");
+		}
+		return this.refresh(creds.refreshToken);
+	}
+
+	private needsRefresh(creds: Credentials): boolean {
+		if (!creds.expiresAt) return false;
+		return creds.expiresAt - this.refreshSkewMs <= Date.now();
+	}
+
+	private refresh(refreshToken: string): Promise<string> {
+		if (this.pending) return this.pending;
+		this.pending = (async () => {
+			try {
+				const next = await refreshAccessToken(this.oauthConfig, refreshToken);
+				// Preserve fields the refresh response doesn't echo back (source,
+				// email, userId) so they survive every rotation. The refresh
+				// response is authoritative for tokens + expiry; we layer the
+				// stable metadata back on top.
+				const existing = this.store.load();
+				this.store.save({
+					accessToken: next.accessToken,
+					refreshToken: next.refreshToken ?? existing?.refreshToken ?? refreshToken,
+					expiresAt: next.expiresAt,
+					scopes: next.scopes,
+					source: existing?.source ?? next.source,
+					userId: existing?.userId ?? next.userId,
+					email: existing?.email ?? next.email,
+					provider: existing?.provider ?? next.provider,
+				});
+				return next.accessToken;
+			} finally {
+				this.pending = null;
+			}
+		})();
+		return this.pending;
+	}
+}
diff --git a/src/glue/client.ts b/src/glue/client.ts
@@ -18,7 +18,12 @@ import { completeSimple, getModel, type KnownProvider, type Model } from "@earen
 export interface GlueOptions {
 	fastModel: Model<string>;
 	smartModel: Model<string>;
-	apiKey: string;
+	/**
+	 * Resolves the API key for each call. Async so callers can layer
+	 * OAuth refresh-on-demand: when the underlying access token is near
+	 * expiry, the getter swaps in a new one transparently.
+	 */
+	getApiKey: () => Promise<string> | string;
 	/** Default 8000 — glue calls cap their reply length so a runaway summary doesn't hang the UI. */
 	maxTokens?: number;
 }
@@ -40,14 +45,15 @@ export class GlueClient {
 		system: string | undefined,
 		signal: AbortSignal | undefined,
 	): Promise<string> {
+		const apiKey = await this.options.getApiKey();
 		const message = await completeSimple(
 			model,
 			{
 				systemPrompt: system,
 				messages: [{ role: "user", content: prompt, timestamp: Date.now() }],
 			},
 			{
-				apiKey: this.options.apiKey,
+				apiKey,
 				signal,
 				maxTokens: this.options.maxTokens ?? 8000,
 			},
