diff --git a/.changeset/steady-codex-refresh.md b/.changeset/steady-codex-refresh.md new file mode 100644 index 00000000000..2d8438541e1 --- /dev/null +++ b/.changeset/steady-codex-refresh.md @@ -0,0 +1,5 @@ +--- +"@kilocode/cli": patch +--- + +Prevent concurrent Kilo processes from reusing a ChatGPT Codex refresh token. diff --git a/packages/opencode/src/kilocode/provider/codex-refresh.ts b/packages/opencode/src/kilocode/provider/codex-refresh.ts index 81a160a3411..0d20e059e30 100644 --- a/packages/opencode/src/kilocode/provider/codex-refresh.ts +++ b/packages/opencode/src/kilocode/provider/codex-refresh.ts @@ -1,4 +1,4 @@ -import type { PluginInput } from "@kilocode/plugin" +import { Flock } from "@opencode-ai/core/util/flock" export class CodexAuthExpiredError extends Error { constructor( @@ -25,14 +25,24 @@ type Tokens = { } type Input = { - input: PluginInput + input: { + client: { + auth: { + set: (input: { path: { id: string }; body: Auth }) => Promise + } + } + } getAuth: () => Promise auth: Auth - refresh: (refresh: string) => Promise + refresh: (refresh: string, signal: AbortSignal) => Promise account: (tokens: Tokens) => string | undefined + lock?: Flock.Options + timeout?: number } const pending = new Map>() +const lock = "codex-auth-refresh:openai" +const timeout = 30_000 function valid(auth: Auth) { return auth.access && auth.expires > Date.now() @@ -59,43 +69,53 @@ function recoverable(err: unknown) { } export async function refreshCodexAuth(input: Input) { - const inflight = pending.get(input.auth.refresh) + const token = input.auth.refresh + const inflight = pending.get(token) if (inflight) { const next = await inflight assign(input.auth, next) return next } - const promise = (async () => { - const fresh = await input.getAuth() - const current = oauth(fresh) - if (current && valid(current)) return current - - try { - const base = current && current.refresh !== input.auth.refresh ? current : input.auth - const tokens = await input.refresh(base.refresh) - const id = input.account(tokens) || base.accountId - const next = { - type: "oauth" as const, - refresh: tokens.refresh_token, - access: tokens.access_token, - expires: Date.now() + (tokens.expires_in ?? 3600) * 1000, - ...(id && { accountId: id }), + const promise = Flock.withLock( + lock, + async () => { + const fresh = await input.getAuth() + const current = oauth(fresh) + if (current && valid(current)) return current + + try { + const base = current && current.refresh !== token ? current : input.auth + const controller = new AbortController() + const timer = setTimeout( + () => controller.abort(new DOMException("The operation timed out.", "TimeoutError")), + input.timeout ?? timeout, + ) + const tokens = await input.refresh(base.refresh, controller.signal).finally(() => clearTimeout(timer)) + const id = input.account(tokens) || base.accountId + const next = { + type: "oauth" as const, + refresh: tokens.refresh_token, + access: tokens.access_token, + expires: Date.now() + (tokens.expires_in ?? 3600) * 1000, + ...(id && { accountId: id }), + } + await input.input.client.auth.set({ path: { id: "openai" }, body: next }) + return next + } catch (err) { + if (!recoverable(err)) throw err + + const latest = await input.getAuth() + const next = oauth(latest) + if (next && usable(next, token)) return next + + throw new CodexAuthExpiredError() } - await input.input.client.auth.set({ path: { id: "openai" }, body: next }) - return next - } catch (err) { - if (!recoverable(err)) throw err - - const latest = await input.getAuth() - const next = oauth(latest) - if (next && usable(next, input.auth.refresh)) return next - - throw new CodexAuthExpiredError() - } - })().finally(() => pending.delete(input.auth.refresh)) + }, + input.lock, + ).finally(() => pending.delete(token)) - pending.set(input.auth.refresh, promise) + pending.set(token, promise) const next = await promise assign(input.auth, next) return next diff --git a/packages/opencode/src/plugin/codex.ts b/packages/opencode/src/plugin/codex.ts index 65892fa7111..887a1443385 100644 --- a/packages/opencode/src/plugin/codex.ts +++ b/packages/opencode/src/plugin/codex.ts @@ -142,10 +142,13 @@ async function exchangeCodeForTokens(code: string, redirectUri: string, pkce: Pk return response.json() } -async function refreshAccessToken(refreshToken: string): Promise { +// kilocode_change start +async function refreshAccessToken(refreshToken: string, signal?: AbortSignal): Promise { const response = await fetch(`${ISSUER}/oauth/token`, { method: "POST", - headers: { "Content-Type": "application/x-www-form-urlencoded" }, + signal, + headers: { "Content-Type": "application/x-www-form-urlencoded", "User-Agent": `kilo/${InstallationVersion}` }, + // kilocode_change end body: new URLSearchParams({ grant_type: "refresh_token", refresh_token: refreshToken, diff --git a/packages/opencode/test/kilocode/agent-permission-overrides.test.ts b/packages/opencode/test/kilocode/agent-permission-overrides.test.ts index 202af60df0c..d55f04629e9 100644 --- a/packages/opencode/test/kilocode/agent-permission-overrides.test.ts +++ b/packages/opencode/test/kilocode/agent-permission-overrides.test.ts @@ -99,7 +99,7 @@ test("system utility agents ignore per-agent permission allows", async () => { }, }) - await WithInstance.provide({ + await provideTestInstance({ directory: tmp.path, fn: async () => { const title = await load(tmp.path, (svc) => svc.get("title")) @@ -131,7 +131,7 @@ test("system utility agents deny tools after configured name override", async () }, }) - await WithInstance.provide({ + await provideTestInstance({ directory: tmp.path, fn: async () => { const title = await load(tmp.path, (svc) => svc.get("title")) diff --git a/packages/opencode/test/kilocode/codex-auth-refresh.test.ts b/packages/opencode/test/kilocode/codex-auth-refresh.test.ts index 28ea37c41ae..1388149615d 100644 --- a/packages/opencode/test/kilocode/codex-auth-refresh.test.ts +++ b/packages/opencode/test/kilocode/codex-auth-refresh.test.ts @@ -1,8 +1,11 @@ import { describe, expect, test } from "bun:test" import { CodexAuthExpiredError, refreshCodexAuth } from "../../src/kilocode/provider/codex-refresh" -import type { PluginInput } from "@kilocode/plugin" import { MessageV2 } from "../../src/session/message-v2" import { ProviderID } from "../../src/provider/schema" +import { spawn } from "child_process" +import fs from "fs/promises" +import os from "os" +import path from "path" type Auth = { type: "oauth" @@ -12,6 +15,13 @@ type Auth = { accountId?: string } +type Lock = { + staleMs: number + timeoutMs: number + baseDelayMs: number + maxDelayMs: number +} + const expired = (): Auth => ({ type: "oauth", access: "old-access", @@ -19,7 +29,10 @@ const expired = (): Auth => ({ expires: 0, }) -function plugin(persist: (auth: Auth) => void): PluginInput { +const root = path.join(import.meta.dir, "../..") +const worker = path.join(import.meta.dir, "fixture/codex-auth-refresh-worker.ts") + +function plugin(persist: (auth: Auth) => void) { const set = async (req: { body: Auth }) => { persist(req.body) } @@ -27,7 +40,101 @@ function plugin(persist: (auth: Auth) => void): PluginInput { client: { auth: { set }, }, - } as unknown as PluginInput + } +} + +function sleep(ms: number) { + return new Promise((resolve) => setTimeout(resolve, ms)) +} + +async function wait(file: string) { + const stop = Date.now() + 10_000 + while (Date.now() < stop) { + if ( + await fs + .stat(file) + .then(() => true) + .catch(() => false) + ) + return + await sleep(10) + } + throw new Error(`Timed out waiting for file: ${file}`) +} + +function run(input: { root: string; url: string; ready: string; start: string; lock?: Lock }) { + const proc = spawn(process.execPath, [worker, JSON.stringify(input)], { + cwd: root, + windowsHide: true, + }) + const stdout: Buffer[] = [] + const stderr: Buffer[] = [] + proc.stdout?.on("data", (data) => stdout.push(Buffer.from(data))) + proc.stderr?.on("data", (data) => stderr.push(Buffer.from(data))) + return { + proc, + done: new Promise<{ code: number; stdout: string; stderr: string }>((resolve) => { + proc.on("close", (code) => { + resolve({ + code: code ?? 1, + stdout: Buffer.concat(stdout).toString(), + stderr: Buffer.concat(stderr).toString(), + }) + }) + }), + } +} + +async function race(input: { reuse: "early" | "late"; delay: number; lock?: Lock }) { + const dir = await fs.mkdtemp(path.join(os.tmpdir(), "codex-auth-refresh-")) + const calls: string[] = [] + const used = new Set() + const server = Bun.serve({ + hostname: "127.0.0.1", + port: 0, + async fetch(req) { + const body = new URLSearchParams(await req.text()) + const token = body.get("refresh_token") + if (!token) return new Response("missing refresh token", { status: 400 }) + calls.push(token) + if (used.has(token)) { + if (input.reuse === "late") await sleep(input.delay + 50) + return new Response("refresh token reused", { status: 401 }) + } + used.add(token) + await sleep(input.delay) + return Response.json({ + id_token: "", + access_token: "next-access", + refresh_token: "next-refresh", + expires_in: 60, + }) + }, + }) + + try { + const data = path.join(dir, "share", "kilo") + const start = path.join(dir, "start") + const first = path.join(dir, "first") + const second = path.join(dir, "second") + await fs.mkdir(data, { recursive: true }) + await fs.writeFile(path.join(data, "auth.json"), JSON.stringify({ openai: expired() })) + const url = `http://127.0.0.1:${server.port}/oauth/token` + const a = run({ root: dir, url, ready: first, start, lock: input.lock }) + const b = run({ root: dir, url, ready: second, start, lock: input.lock }) + try { + await Promise.all([wait(first), wait(second)]) + await fs.writeFile(start, "") + const out = await Promise.all([a.done, b.done]) + return { calls, out } + } finally { + if (a.proc.exitCode === null) a.proc.kill() + if (b.proc.exitCode === null) b.proc.kill() + } + } finally { + await server.stop(true) + await fs.rm(dir, { recursive: true, force: true }) + } } describe("Codex auth refresh", () => { @@ -121,4 +228,113 @@ describe("Codex auth refresh", () => { }), ).rejects.toBeInstanceOf(CodexAuthExpiredError) }) + + test("refreshes a newer stored token instead of the stale caller token", async () => { + const auth = expired() + const fresh = { ...expired(), refresh: "fresh-refresh", accountId: "account-1" } + const calls: string[] = [] + + const result = await refreshCodexAuth({ + input: plugin(() => {}), + getAuth: async () => fresh, + auth, + refresh: async (token) => { + calls.push(token) + return { id_token: "", access_token: "next-access", refresh_token: "next-refresh", expires_in: 60 } + }, + account: () => undefined, + }) + + expect(calls).toEqual(["fresh-refresh"]) + expect(result.accountId).toBe("account-1") + }) + + test("releases the lock and pending entry after a transient failure", async () => { + const auth = expired() + const calls: string[] = [] + const failed = await refreshCodexAuth({ + input: plugin(() => {}), + getAuth: async () => auth, + auth, + refresh: async (token) => { + calls.push(token) + throw new Error("offline") + }, + account: () => undefined, + }).catch((err) => err) + + expect(failed).toEqual(new Error("offline")) + + await refreshCodexAuth({ + input: plugin(() => {}), + getAuth: async () => auth, + auth, + refresh: async (token) => { + calls.push(token) + return { id_token: "", access_token: "next-access", refresh_token: "next-refresh", expires_in: 60 } + }, + account: () => undefined, + }) + + expect(calls).toEqual(["old-refresh", "old-refresh"]) + }) + + test("aborts a stalled refresh and releases the lock", async () => { + const auth = expired() + const calls: string[] = [] + const failed = await refreshCodexAuth({ + input: plugin(() => {}), + getAuth: async () => auth, + auth, + refresh: async (token, signal) => { + calls.push(token) + return new Promise((_, reject) => { + signal.addEventListener("abort", () => reject(signal.reason), { once: true }) + }) + }, + account: () => undefined, + timeout: 20, + }).catch((err) => err) + + expect(failed).toBeInstanceOf(DOMException) + expect(failed.name).toBe("TimeoutError") + + await refreshCodexAuth({ + input: plugin(() => {}), + getAuth: async () => auth, + auth, + refresh: async (token) => { + calls.push(token) + return { id_token: "", access_token: "next-access", refresh_token: "next-refresh", expires_in: 60 } + }, + account: () => undefined, + }) + + expect(calls).toEqual(["old-refresh", "old-refresh"]) + }) + + test("serializes refreshes across processes before token reuse", async () => { + for (const reuse of ["early", "late"] as const) { + for (let trial = 0; trial < 5; trial++) { + const result = await race({ reuse, delay: 100 }) + expect(result.calls).toEqual(["old-refresh"]) + expect(result.out.map((x) => x.code)).toEqual([0, 0]) + expect(result.out.map((x) => x.stderr)).toEqual(["", ""]) + } + } + }, 30_000) + + test("keeps the process lock alive during a delayed token response", async () => { + const lock = { + staleMs: 300, + timeoutMs: 10_000, + baseDelayMs: 20, + maxDelayMs: 30, + } + const result = await race({ reuse: "early", delay: 1_000, lock }) + + expect(result.calls).toEqual(["old-refresh"]) + expect(result.out.map((x) => x.code)).toEqual([0, 0]) + expect(result.out.map((x) => x.stderr)).toEqual(["", ""]) + }, 15_000) }) diff --git a/packages/opencode/test/kilocode/codex-refresh-user-agent.test.ts b/packages/opencode/test/kilocode/codex-refresh-user-agent.test.ts new file mode 100644 index 00000000000..2e4ad682286 --- /dev/null +++ b/packages/opencode/test/kilocode/codex-refresh-user-agent.test.ts @@ -0,0 +1,55 @@ +import { expect, test } from "bun:test" +import type { PluginInput } from "@kilocode/plugin" +import { CodexAuthPlugin } from "../../src/plugin/codex" + +test("identifies Codex refresh requests as Kilo", async () => { + const original = globalThis.fetch + const seen: Request[] = [] + const signals: (AbortSignal | null | undefined)[] = [] + let auth = { + type: "oauth" as const, + access: "old-access", + refresh: "old-refresh", + expires: 0, + } + const input = { + client: { + auth: { + set: async (req: { body: typeof auth }) => { + auth = req.body + }, + }, + }, + } as unknown as PluginInput + globalThis.fetch = Object.assign( + async (...args: Parameters) => { + const req = new Request(...args) + seen.push(req) + signals.push(args[1]?.signal) + if (req.url === "https://auth.openai.com/oauth/token") { + return Response.json({ + id_token: "", + access_token: "next-access", + refresh_token: "next-refresh", + expires_in: 60, + }) + } + return new Response("", { status: 200 }) + }, + { preconnect: original.preconnect }, + ) + + try { + const plugin = await CodexAuthPlugin(input) + const loaded = await plugin.auth!.loader!(async () => auth, {} as never) + await loaded.fetch("https://api.openai.com/v1/responses") + } finally { + globalThis.fetch = original + } + + const refresh = seen[0] + expect(refresh.url).toBe("https://auth.openai.com/oauth/token") + expect(refresh.headers.get("user-agent")).toMatch(/^kilo\//) + expect(signals[0]).toBeInstanceOf(AbortSignal) + expect(await refresh.text()).toContain("refresh_token=old-refresh") +}) diff --git a/packages/opencode/test/kilocode/fixture/codex-auth-refresh-worker.ts b/packages/opencode/test/kilocode/fixture/codex-auth-refresh-worker.ts new file mode 100644 index 00000000000..845b25fde88 --- /dev/null +++ b/packages/opencode/test/kilocode/fixture/codex-auth-refresh-worker.ts @@ -0,0 +1,113 @@ +import fs from "fs/promises" +import path from "path" +import z from "zod" + +const Auth = z.object({ + type: z.literal("oauth"), + refresh: z.string(), + access: z.string(), + expires: z.number(), + accountId: z.string().optional(), +}) +type Auth = z.infer + +const Tokens = z.object({ + id_token: z.string(), + access_token: z.string(), + refresh_token: z.string(), + expires_in: z.number().optional(), +}) + +const Msg = z.object({ + root: z.string(), + url: z.string(), + ready: z.string(), + start: z.string(), + lock: z + .object({ + staleMs: z.number(), + timeoutMs: z.number(), + baseDelayMs: z.number(), + maxDelayMs: z.number(), + }) + .optional(), +}) + +function input() { + const raw = process.argv[2] + if (!raw) throw new Error("Missing Codex auth refresh worker input") + return Msg.parse(JSON.parse(raw)) +} + +function sleep(ms: number) { + return new Promise((resolve) => setTimeout(resolve, ms)) +} + +async function wait(file: string) { + const stop = Date.now() + 10_000 + while (Date.now() < stop) { + if ( + await fs + .stat(file) + .then(() => true) + .catch(() => false) + ) + return + await sleep(10) + } + throw new Error(`Timed out waiting for file: ${file}`) +} + +async function main() { + const msg = input() + process.env.XDG_DATA_HOME = path.join(msg.root, "share") + process.env.XDG_CACHE_HOME = path.join(msg.root, "cache") + process.env.XDG_CONFIG_HOME = path.join(msg.root, "config") + process.env.XDG_STATE_HOME = path.join(msg.root, "state") + process.env.KILO_TEST_HOME = path.join(msg.root, "home") + + const { Path } = await import("@opencode-ai/core/global") + const { refreshCodexAuth } = await import("../../../src/kilocode/provider/codex-refresh") + const file = path.join(Path.data, "auth.json") + const read = async () => { + const data = z.object({ openai: Auth }).parse(JSON.parse(await fs.readFile(file, "utf8"))) + return data.openai + } + const plugin = { + client: { + auth: { + set: async (req: { body: Auth }) => { + await fs.writeFile(file, JSON.stringify({ openai: req.body })) + }, + }, + }, + } + + await fs.writeFile(msg.ready, String(process.pid)) + await wait(msg.start) + const auth = await read() + const next = await refreshCodexAuth({ + input: plugin, + getAuth: read, + auth, + refresh: async (token) => { + const response = await fetch(msg.url, { + method: "POST", + headers: { "Content-Type": "application/x-www-form-urlencoded" }, + body: new URLSearchParams({ refresh_token: token }).toString(), + }) + if (!response.ok) throw new Error(`Token refresh failed: ${response.status}`) + return Tokens.parse(await response.json()) + }, + account: () => undefined, + lock: msg.lock, + }) + + process.stdout.write(JSON.stringify(next)) +} + +await main().catch((err) => { + const text = err instanceof Error ? (err.stack ?? err.message) : String(err) + process.stderr.write(text) + process.exit(1) +}) diff --git a/packages/opencode/test/kilocode/oauth-branding.test.ts b/packages/opencode/test/kilocode/oauth-branding.test.ts index 1a8556a5b5c..71cc234b3db 100644 --- a/packages/opencode/test/kilocode/oauth-branding.test.ts +++ b/packages/opencode/test/kilocode/oauth-branding.test.ts @@ -8,6 +8,7 @@ describe("Kilo OAuth branding", () => { const src = await Bun.file(path.join(root, "src", "plugin", "codex.ts")).text() expect(src).toContain('originator: "kilo"') + expect(src).toContain('"User-Agent": `kilo/${InstallationVersion}`') expect(src).toContain("return to Kilo") expect(src).not.toContain('originator: "opencode"') expect(src).not.toContain("return to OpenCode")