diff --git a/src/codex/account-store.ts b/src/codex/account-store.ts index 5932a614b7..85eaab3ef8 100644 --- a/src/codex/account-store.ts +++ b/src/codex/account-store.ts @@ -386,6 +386,19 @@ function findFreshCredentialForGrant( return null; } +async function notePlanFromRefreshedAccessToken( + id: string, + accessToken: string, + generation: number, +): Promise { + try { + const { noteCodexAccountAccessToken } = await import("./plan-from-token"); + noteCodexAccountAccessToken(id, accessToken, generation); + } catch { + // Derived plan metadata must not fail credential refresh. + } +} + export async function getValidCodexToken(id: string): Promise { const record = readCodexAccountRecord(id); const cred = record?.deletedAt == null ? record?.credential : undefined; @@ -415,10 +428,12 @@ export async function getValidCodexToken(id: string): Promise if (!saveCodexAccountCredentialIfGeneration(id, current.generation, refreshed.credential)) { throw new CodexCredentialGenerationConflictError(); } + const generation = current.generation + 1; + await notePlanFromRefreshedAccessToken(id, refreshed.credential.accessToken, generation); return { accessToken: refreshed.credential.accessToken, chatgptAccountId: refreshed.credential.chatgptAccountId, - generation: current.generation + 1, + generation, }; } return getValidCodexToken(id); @@ -520,6 +535,7 @@ export async function getValidCodexToken(id: string): Promise flight = { promise: refreshPromise, startedAt: Date.now(), abort }; refreshLocks.set(refreshGrantFingerprint, flight); const result = await refreshPromise; + await notePlanFromRefreshedAccessToken(id, result.accessToken, result.generation); return { accessToken: result.accessToken, chatgptAccountId: result.chatgptAccountId, diff --git a/src/codex/main-account.ts b/src/codex/main-account.ts index 495ad39300..2bc7457d5d 100644 --- a/src/codex/main-account.ts +++ b/src/codex/main-account.ts @@ -1,5 +1,6 @@ import { readCodexTokens } from "./auth-collision"; import { decodeJwtPayload } from "../oauth/chatgpt"; +import { extractChatgptPlanType } from "./plan"; import { MAIN_CODEX_ACCOUNT_ID } from "./account-id"; export { MAIN_CODEX_ACCOUNT_ID } from "./account-id"; @@ -10,13 +11,23 @@ export { MAIN_CODEX_ACCOUNT_ID } from "./account-id"; * percent, matching pool-account behavior. */ let mainAccountPlan: string | null = null; +let jwtPlanAttempted = false; export function setMainAccountPlan(plan: string | null): void { mainAccountPlan = plan; + if (plan === null) jwtPlanAttempted = false; } export function getMainAccountPlan(): string | undefined { - return mainAccountPlan ?? undefined; + if (mainAccountPlan) return mainAccountPlan; + if (jwtPlanAttempted) return undefined; + jwtPlanAttempted = true; + const tokens = readCodexTokens(); + const jwtPlan = tokens + ? extractChatgptPlanType(tokens.id_token, tokens.access_token) + : undefined; + if (jwtPlan) mainAccountPlan = jwtPlan; + return jwtPlan; } /** Read-only main account token from ~/.codex/auth.json, or null when not logged in. */ diff --git a/src/codex/plan-from-token.ts b/src/codex/plan-from-token.ts new file mode 100644 index 0000000000..40549364aa --- /dev/null +++ b/src/codex/plan-from-token.ts @@ -0,0 +1,115 @@ +import { + ConfigMutationLockError, + loadConfig, + mutatePersistedConfig, +} from "../config"; +import type { CodexAccount, OcxConfig } from "../types"; +import { isSelectableCodexPoolAccount, isValidCodexAccountId } from "./account-id"; +import { + getCodexAccountCredential, + isCodexAccountGenerationLive, + readCodexAccountRecord, +} from "./account-store"; +import { extractChatgptPlanType, codexPlanValue } from "./plan"; + +interface FreshPoolPlanUpdate { + accountId: string; + plan: string; + credentialGeneration: number; +} + +function configuredPoolAccount(config: OcxConfig, accountId: string): CodexAccount | null { + if (!isValidCodexAccountId(accountId)) return null; + return (config.codexAccounts ?? []) + .find(account => account.id === accountId && isSelectableCodexPoolAccount(account)) ?? null; +} + +function jwtPlanFromPoolCredential(accountId: string): string | undefined { + const cred = getCodexAccountCredential(accountId); + return cred ? extractChatgptPlanType(undefined, cred.accessToken) : undefined; +} + +function collectJwtPoolPlanUpdates(runtimeConfig: OcxConfig): FreshPoolPlanUpdate[] { + const updates: FreshPoolPlanUpdate[] = []; + for (const account of (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount)) { + const jwtPlan = jwtPlanFromPoolCredential(account.id); + if (!jwtPlan || codexPlanValue(account.plan) === jwtPlan) continue; + const generation = readCodexAccountRecord(account.id)?.generation; + if (generation === undefined) continue; + updates.push({ accountId: account.id, plan: jwtPlan, credentialGeneration: generation }); + } + return updates; +} + +const appliedJwtPlans = new Map(); + +function persistJwtPlanUpdates(runtimeConfig: OcxConfig, updates: FreshPoolPlanUpdate[]): void { + if (updates.length === 0) return; + let outcome: ReturnType>; + try { + outcome = mutatePersistedConfig(persistedConfig => { + const accepted: FreshPoolPlanUpdate[] = []; + let changed = false; + for (const update of updates) { + if (!isCodexAccountGenerationLive(update.accountId, update.credentialGeneration)) continue; + const liveAccount = configuredPoolAccount(runtimeConfig, update.accountId); + const persistedAccount = configuredPoolAccount(persistedConfig, update.accountId); + if (!liveAccount || !persistedAccount) continue; + accepted.push(update); + if (persistedAccount.plan !== update.plan) { + persistedAccount.plan = update.plan; + changed = true; + } + } + return { changed, value: accepted }; + }); + } catch (error) { + if (error instanceof ConfigMutationLockError) return; + throw error; + } + if (outcome.status === "unavailable") return; + for (const update of outcome.value) { + if (!isCodexAccountGenerationLive(update.accountId, update.credentialGeneration)) continue; + const liveAccount = configuredPoolAccount(runtimeConfig, update.accountId); + if (liveAccount) { + liveAccount.plan = update.plan; + appliedJwtPlans.set(update.accountId, update.plan); + } + } +} + +/** + * Persist JWT `chatgpt_plan_type` onto `codexAccounts[].plan` when it contradicts the stored + * label. Generation-gated, same fail-closed lock policy as the WHAM plan patch. Does not + * overwrite a plan that already matches the token. + */ +export function reconcileCodexPlansFromTokens(runtimeConfig: OcxConfig = loadConfig()): void { + persistJwtPlanUpdates(runtimeConfig, collectJwtPoolPlanUpdates(runtimeConfig)); +} + +export function resetJwtPlanNotesForTests(): void { + appliedJwtPlans.clear(); +} + +/** Apply one account's live token claim without blocking the credential read path. */ +export function noteCodexAccountAccessToken( + accountId: string, + accessToken: string, + credentialGeneration: number, +): void { + const jwtPlan = extractChatgptPlanType(undefined, accessToken); + if (!jwtPlan) return; + if (appliedJwtPlans.get(accountId) === jwtPlan) return; + try { + const runtimeConfig = loadConfig(); + const live = configuredPoolAccount(runtimeConfig, accountId); + if (!live || codexPlanValue(live.plan) === jwtPlan) { + appliedJwtPlans.set(accountId, jwtPlan); + return; + } + persistJwtPlanUpdates(runtimeConfig, [{ accountId, plan: jwtPlan, credentialGeneration }]); + if (codexPlanValue(live.plan) === jwtPlan) appliedJwtPlans.set(accountId, jwtPlan); + } catch { + appliedJwtPlans.delete(accountId); + } +} diff --git a/src/codex/plan.ts b/src/codex/plan.ts index 6acb2d1b46..e3a9f70f80 100644 --- a/src/codex/plan.ts +++ b/src/codex/plan.ts @@ -1,3 +1,5 @@ +import { decodeJwtPayload } from "../oauth/chatgpt"; + /** Preserve user/provider plan labels only when they are usable strings. */ export function codexPlanValue(value: unknown): string | undefined { if (typeof value !== "string") return undefined; @@ -13,3 +15,26 @@ export function isThirtyDayOnlyCodexPlan(value: unknown): boolean { const key = codexPlanKey(value); return key === "go" || key === "free"; } + +/** + * ChatGPT plan label from a live access/id token (`chatgpt_plan_type`). + * Used when WHAM has not run yet so a stale stored `free` cannot outrank the token (#1989). + */ +export function extractChatgptPlanType(idToken?: string, accessToken?: string): string | undefined { + for (const token of [idToken, accessToken]) { + if (!token) continue; + const payload = decodeJwtPayload(token); + if (!payload) continue; + if (typeof payload.chatgpt_plan_type === "string") { + const plan = codexPlanValue(payload.chatgpt_plan_type); + if (plan) return plan; + } + const ns = payload["https://api.openai.com/auth"]; + if (ns && typeof ns === "object") { + const nested = (ns as Record).chatgpt_plan_type; + const plan = codexPlanValue(nested); + if (plan) return plan; + } + } + return undefined; +} diff --git a/src/server/index.ts b/src/server/index.ts index 87ece913e1..341d4bbdd4 100644 --- a/src/server/index.ts +++ b/src/server/index.ts @@ -1728,7 +1728,15 @@ export function startServer(port?: number, deps: StartServerDeps = {}): Server { + try { + reconcileCodexPlansFromTokens(config); + } catch { + // Derived plan metadata must not block WHAM priming. + } + return import("../codex/auth-api"); + }) .then(({ primeCodexPoolQuotas }) => primeCodexPoolQuotas(config, "startup")) .catch(() => {}); } diff --git a/tests/codex-auth-api.test.ts b/tests/codex-auth-api.test.ts index bcd7eb4399..027c28301e 100644 --- a/tests/codex-auth-api.test.ts +++ b/tests/codex-auth-api.test.ts @@ -17,6 +17,7 @@ import { clearCodexQuotaPrimeState, primeCodexPoolQuotas, seedCodexAuthAdmissionForTests, type CodexAuthAccountDto, listCodexAuthAccounts, + setAccountQuotaFromParsed, } from "../src/codex/auth-api"; import { getCodexAccountCredential, @@ -44,6 +45,7 @@ import type { WsData } from "../src/server/ws-bridge"; import { handleNativeProfileAPI } from "../src/codex/native-profile-api"; import type { NativeProfileManager } from "../src/codex/native-profile-manager"; import { MAIN_CODEX_ACCOUNT_ID, setMainAccountPlan } from "../src/codex/main-account"; +import { reconcileCodexPlansFromTokens, resetJwtPlanNotesForTests } from "../src/codex/plan-from-token"; import { deleteCodexAccount, reconcileMainCodexAccountRuntimeState, @@ -210,6 +212,15 @@ async function completeMockCodexOAuth(options: { } } +function chatgptPlanJwt(plan: string, accountId = "acct"): string { + const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ + chatgpt_account_id: accountId, + "https://api.openai.com/auth": { chatgpt_account_id: accountId, chatgpt_plan_type: plan }, + })).toString("base64url"); + return `${header}.${body}.sig`; +} + function seedPoolAccount( config: OcxConfig, account: { @@ -255,6 +266,7 @@ beforeEach(() => { clearPoolRotationState(); clearCodexWebSocketRegistry(); resetMainCodexAccountIdentityTrackingForTests(); + resetJwtPlanNotesForTests(); }); afterEach(() => { @@ -1325,6 +1337,77 @@ describe("codex-auth API", () => { expect(configCommits).toBe(0); }); + test("quota cache hit still corrects a stale stored pool plan from the access-token JWT (#1989)", async () => { + const config = makeConfig(); + const accountId = "pool-jwt-plan"; + seedPoolAccount(config, { + id: accountId, + email: "pool-jwt-plan@example.com", + plan: "free", + accessToken: chatgptPlanJwt("pro", `acct-${accountId}`), + chatgptAccountId: `acct-${accountId}`, + }); + saveConfig(structuredClone(config)); + setAccountQuotaFromParsed(accountId, { weeklyPercent: 4 }, captureConfigGeneration()); + reconcileCodexPlansFromTokens(config); + let whamCalls = 0; + globalThis.fetch = (async () => { + whamCalls += 1; + return Response.json({ plan_type: "free" }); + }) as typeof fetch; + + const req = new Request("http://localhost/api/codex-auth/accounts", { method: "GET" }); + const resp = await handleCodexAuthAPI(req, new URL(req.url), config); + const data = await resp!.json() as { accounts: Array<{ id: string; plan?: string }> }; + + expect(whamCalls).toBe(0); + expect(data.accounts.find(account => account.id === accountId)?.plan).toBe("pro"); + expect(config.codexAccounts?.find(account => account.id === accountId)?.plan).toBe("pro"); + expect(loadConfig().codexAccounts?.find(account => account.id === accountId)?.plan).toBe("pro"); + }); + + test("a live WHAM plan_type still outranks a contradicting access-token JWT", async () => { + const config = makeConfig(); + seedPoolAccount(config, { + id: "pool-wham-wins", + email: "pool-wham-wins@example.com", + plan: "free", + accessToken: chatgptPlanJwt("plus", "acct-pool-wham-wins"), + chatgptAccountId: "acct-pool-wham-wins", + }); + saveConfig(structuredClone(config)); + globalThis.fetch = (async () => Response.json({ + plan_type: "prolite", + rate_limit: { primary_window: { used_percent: 11, reset_at: 1782628379 } }, + })) as typeof fetch; + + const req = new Request("http://localhost/api/codex-auth/accounts?refresh=1", { method: "GET" }); + const resp = await handleCodexAuthAPI(req, new URL(req.url), config); + const data = await resp!.json() as { accounts: Array<{ id: string; plan?: string }> }; + + expect(data.accounts.find(account => account.id === "pool-wham-wins")?.plan).toBe("prolite"); + expect(loadConfig().codexAccounts?.find(account => account.id === "pool-wham-wins")?.plan).toBe("prolite"); + }); + + test("main account list uses chatgpt_plan_type when WHAM omits plan_type (#1989)", async () => { + writeFileSync(join(TEST_CODEX_HOME, "auth.json"), JSON.stringify({ + tokens: { + access_token: chatgptPlanJwt("pro", "acct-main-jwt"), + account_id: "acct-main-jwt", + }, + })); + globalThis.fetch = (async () => Response.json({ + email: "main-jwt@example.test", + rate_limit: { primary_window: { used_percent: 2, reset_at: 1782628379 } }, + })) as typeof fetch; + + const req = new Request("http://localhost/api/codex-auth/accounts", { method: "GET" }); + const resp = await handleCodexAuthAPI(req, new URL(req.url), makeConfig()); + const data = await resp!.json() as { accounts: Array<{ id: string; plan?: string | null }> }; + + expect(data.accounts.find(account => account.id === MAIN_CODEX_ACCOUNT_ID)?.plan).toBe("pro"); + }); + test("pool plan refresh does not recreate a config file deleted while the server is running", async () => { const config = makeConfig(); seedPoolAccount(config, { diff --git a/tests/codex-plan.test.ts b/tests/codex-plan.test.ts new file mode 100644 index 0000000000..7654a5f780 --- /dev/null +++ b/tests/codex-plan.test.ts @@ -0,0 +1,126 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; +import { saveCodexAccountCredential } from "../src/codex/account-store"; +import { getMainAccountPlan, setMainAccountPlan } from "../src/codex/main-account"; +import { extractChatgptPlanType } from "../src/codex/plan"; +import { + reconcileCodexPlansFromTokens, + resetJwtPlanNotesForTests, +} from "../src/codex/plan-from-token"; +import { loadConfig, saveConfig } from "../src/config"; +import type { OcxConfig } from "../src/types"; + +const TEST_DIR = join(import.meta.dir, ".tmp-codex-plan-test"); +const TEST_CODEX_HOME = join(TEST_DIR, "codex"); +let previousOpencodexHome: string | undefined; +let previousCodexHome: string | undefined; + +function chatgptPlanJwt(plan: string, accountId = "acct"): string { + const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ + chatgpt_account_id: accountId, + chatgpt_plan_type: plan, + "https://api.openai.com/auth": { chatgpt_account_id: accountId, chatgpt_plan_type: plan }, + })).toString("base64url"); + return `${header}.${body}.sig`; +} + +beforeEach(() => { + previousOpencodexHome = process.env.OPENCODEX_HOME; + previousCodexHome = process.env.CODEX_HOME; + if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true }); + mkdirSync(TEST_CODEX_HOME, { recursive: true }); + process.env.OPENCODEX_HOME = TEST_DIR; + process.env.CODEX_HOME = TEST_CODEX_HOME; + setMainAccountPlan(null); + resetJwtPlanNotesForTests(); +}); + +afterEach(() => { + setMainAccountPlan(null); + resetJwtPlanNotesForTests(); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; + else process.env.CODEX_HOME = previousCodexHome; + if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true }); +}); + +describe("extractChatgptPlanType", () => { + test("reads the namespaced chatgpt_plan_type claim", () => { + const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ + "https://api.openai.com/auth": { chatgpt_plan_type: "pro" }, + })).toString("base64url"); + expect(extractChatgptPlanType(undefined, `${header}.${body}.sig`)).toBe("pro"); + }); + + test("reads a top-level chatgpt_plan_type claim", () => { + const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ chatgpt_plan_type: "plus" })).toString("base64url"); + expect(extractChatgptPlanType(`${header}.${body}.sig`)).toBe("plus"); + }); + + test("ignores non-JWT access tokens", () => { + expect(extractChatgptPlanType(undefined, "access-pool-1")).toBeUndefined(); + }); +}); + +describe("reconcileCodexPlansFromTokens", () => { + test("persists a stale stored free plan from the live access-token JWT (#1989)", () => { + const config: OcxConfig = { + port: 10100, + providers: {}, + defaultProvider: "openai", + codexAccounts: [{ id: "pool-jwt-plan", email: "pool@example.test", plan: "free", isMain: false }], + }; + saveConfig(config); + saveCodexAccountCredential("pool-jwt-plan", { + accessToken: chatgptPlanJwt("pro", "acct-pool-jwt-plan"), + refreshToken: "refresh-pool-jwt-plan", + expiresAt: Date.now() + 5 * 60_000, + chatgptAccountId: "acct-pool-jwt-plan", + }); + + reconcileCodexPlansFromTokens(config); + + expect(config.codexAccounts?.[0]?.plan).toBe("pro"); + expect(loadConfig().codexAccounts?.[0]?.plan).toBe("pro"); + }); + + test("leaves a non-JWT pool credential's stored plan alone", () => { + const config: OcxConfig = { + port: 10100, + providers: {}, + defaultProvider: "openai", + codexAccounts: [{ id: "pool-plain", email: "plain@example.test", plan: "free", isMain: false }], + }; + saveConfig(config); + saveCodexAccountCredential("pool-plain", { + accessToken: "access-pool-plain", + refreshToken: "refresh-pool-plain", + expiresAt: Date.now() + 5 * 60_000, + chatgptAccountId: "acct-pool-plain", + }); + + reconcileCodexPlansFromTokens(config); + + expect(config.codexAccounts?.[0]?.plan).toBe("free"); + expect(loadConfig().codexAccounts?.[0]?.plan).toBe("free"); + }); +}); + +describe("getMainAccountPlan JWT fallback", () => { + test("reads chatgpt_plan_type from auth.json when WHAM has not cached a plan (#1989)", () => { + writeFileSync(join(TEST_CODEX_HOME, "auth.json"), JSON.stringify({ + tokens: { + access_token: chatgptPlanJwt("pro", "acct-main-jwt"), + account_id: "acct-main-jwt", + }, + })); + + expect(getMainAccountPlan()).toBe("pro"); + expect(getMainAccountPlan()).toBe("pro"); + }); +});