Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -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",
Expand Down
47 changes: 46 additions & 1 deletion src/acp/server.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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([], {
Expand Down
3 changes: 3 additions & 0 deletions src/acp/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
27 changes: 27 additions & 0 deletions src/agent/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
9 changes: 8 additions & 1 deletion src/agent/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
46 changes: 46 additions & 0 deletions src/auth/ensure-fresh.test.ts
Original file line number Diff line number Diff line change
@@ -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");
});
});
13 changes: 6 additions & 7 deletions src/auth/ensure-fresh.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading