From 02c7e0829a1bb07f1257c9d2a6647965dd2401d9 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:59:24 +0800 Subject: [PATCH 1/2] fix(codex): re-derive pool plan from JWT chatgpt_plan_type between WHAM refreshes A stale stored `free` outranked the live access-token claim across token refresh and `ocx restart`, so quota windows and `ocx account list` stayed wrong until a manual WHAM refresh. Use the JWT plan when WHAM has not produced a fresh plan_type, and keep WHAM authoritative when it has. Closes #1989 Co-authored-by: Cursor --- src/codex/auth-api.ts | 51 ++++++++++++++++++++--- src/oauth/chatgpt.ts | 21 ++++++++++ tests/chatgpt-oauth.test.ts | 25 ++++++++++- tests/codex-auth-api.test.ts | 80 ++++++++++++++++++++++++++++++++++++ 4 files changed, 170 insertions(+), 7 deletions(-) diff --git a/src/codex/auth-api.ts b/src/codex/auth-api.ts index f2d08188b5..a52c899a54 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -82,7 +82,7 @@ export { setAccountQuotaFromParsed, updateAccountQuota, } from "./quota"; -import { extractAccountId } from "../oauth/chatgpt"; +import { extractAccountId, extractChatgptPlanType } from "../oauth/chatgpt"; import { getMainAccountPlan, MAIN_CODEX_ACCOUNT_ID, setMainAccountPlan } from "./main-account"; import { captureConfigGeneration, registerStateSweepAfterTick } from "../lib/state-store-sweeper"; import { reconcileLiveStateStores } from "../lib/state-store-registrations"; @@ -697,8 +697,16 @@ async function fetchMainAccountInfoWhileOwned( } const tokens = tokenRead.tokens; const requestAccountId = extractAccountId(tokens.id_token, tokens.access_token) ?? (tokens.account_id || null); + const jwtPlan = extractChatgptPlanType(tokens.id_token, tokens.access_token); const cached = getMainAccountInfoCache(); if (!forceRefresh && cached && Date.now() - cached.ts < MAIN_CACHE_TTL) { + const plan = nonEmptyPlan(jwtPlan) ?? cached.plan; + if (plan && plan !== cached.plan) { + const info = { ...cached, plan }; + setMainAccountInfoCache(info); + setMainAccountPlan(plan); + return { info, credentialChecked: true, hasCredential: true }; + } return { info: cached, credentialChecked: true, hasCredential: true }; } try { @@ -719,7 +727,10 @@ async function fetchMainAccountInfoWhileOwned( const data = (await resp.json()) as WhamUsageResponse; const retried = await retryMainAccountInfoIfIdentityChanged(requestAccountId, retriesRemaining, nativeMainLease); if (retried) return retried; - const plan = nonEmptyPlan(data.plan_type) ?? nonEmptyPlan(cached?.plan) ?? nonEmptyPlan(getMainAccountPlan()); + const plan = nonEmptyPlan(data.plan_type) + ?? nonEmptyPlan(jwtPlan) + ?? nonEmptyPlan(cached?.plan) + ?? nonEmptyPlan(getMainAccountPlan()); const quota = parseUsageQuota({ ...data, ...(plan ? { plan_type: plan } : {}) }); const freshResetCredits = quota?.resetCredits; const result = { @@ -877,6 +888,24 @@ function reconcileFreshPoolAccountPlans(runtimeConfig: OcxConfig, updates: Fresh } } +function jwtPlanFromPoolCredential(accountId: string): string | undefined { + const cred = getCodexAccountCredential(accountId); + return cred ? extractChatgptPlanType(undefined, cred.accessToken) : undefined; +} + +/** Local JWT claim vs persisted plan, generation-gated. WHAM `freshPlan` still wins when present. */ +function collectJwtPoolPlanUpdates(runtimeConfig: OcxConfig): FreshPoolPlanUpdate[] { + const updates: FreshPoolPlanUpdate[] = []; + for (const account of (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount)) { + const jwtPlan = jwtPlanFromPoolCredential(account.id); + if (!jwtPlan || nonEmptyPlan(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; +} + async function fetchFreshPoolAccountQuota( accountId: string, existing: StoredAccountQuota | null, @@ -1114,6 +1143,8 @@ export async function primeCodexPoolQuotas( } catch { // Priming is best-effort; never propagate. } + // Token claims are local: a stale stored `free` must not wait for the next WHAM TTL (#1989). + reconcileFreshPoolAccountPlans(runtimeConfig, collectJwtPoolPlanUpdates(runtimeConfig)); if (process.env.OPENCODEX_DEBUG_QUOTA === "1") { console.warn(`[codex-quota] prime done (reason=${reason}, pool=${pool.length}, refreshed=${stale.length})`); } @@ -1179,6 +1210,11 @@ export async function listCodexAuthAccountsSnapshot( : []; }); reconcileFreshPoolAccountPlans(runtimeConfig, planUpdates); + const whamAccountIds = new Set(planUpdates.map(update => update.accountId)); + reconcileFreshPoolAccountPlans( + runtimeConfig, + collectJwtPoolPlanUpdates(runtimeConfig).filter(update => !whamAccountIds.has(update.accountId)), + ); const withQuota = refreshedPool.flatMap(({ accountId, quotaResult }) => { const currentAccount = configuredPoolAccount(runtimeConfig, accountId); @@ -1199,10 +1235,13 @@ export async function listCodexAuthAccountsSnapshot( const effectiveQuotaResult = !generationLive ? { quota: null, needsReauth: false } : quotaResult; - // Response DTO can show the WHAM plan even when disk persistence fails closed (lock busy / - // missing config). Persistence still remains generation-gated via reconcileFreshPoolAccountPlans. - const dtoAccount = generationLive && quotaResult.freshPlan - ? { ...currentAccount, plan: quotaResult.freshPlan } + // WHAM plan wins when this probe produced one; otherwise a live JWT claim may correct + // a stale stored plan even on a quota cache hit (#1989). + const dtoPlan = generationLive + ? (quotaResult.freshPlan ?? jwtPlanFromPoolCredential(accountId)) + : undefined; + const dtoAccount = dtoPlan + ? { ...currentAccount, plan: dtoPlan } : currentAccount; return [poolAccountDto( dtoAccount, diff --git a/src/oauth/chatgpt.ts b/src/oauth/chatgpt.ts index f4ecc7f8a9..4678ca8929 100644 --- a/src/oauth/chatgpt.ts +++ b/src/oauth/chatgpt.ts @@ -46,6 +46,27 @@ export function extractEmail(idToken?: string, accessToken?: string): string | u return undefined; } +/** + * 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" && payload.chatgpt_plan_type.trim()) { + return payload.chatgpt_plan_type.trim(); + } + const ns = payload["https://api.openai.com/auth"]; + if (ns && typeof ns === "object") { + const nested = (ns as Record).chatgpt_plan_type; + if (typeof nested === "string" && nested.trim()) return nested.trim(); + } + } + return undefined; +} + function credsFromToken(data: Record): OAuthCredentials { const idToken = typeof data.id_token === "string" ? data.id_token : undefined; const accessToken = data.access_token as string; diff --git a/tests/chatgpt-oauth.test.ts b/tests/chatgpt-oauth.test.ts index 161baca024..df26e11def 100644 --- a/tests/chatgpt-oauth.test.ts +++ b/tests/chatgpt-oauth.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { decodeJwtPayload, extractAccountId, extractEmail } from "../src/oauth/chatgpt"; +import { decodeJwtPayload, extractAccountId, extractChatgptPlanType, extractEmail } from "../src/oauth/chatgpt"; function fakeJwt(payload: Record): string { const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); @@ -78,6 +78,29 @@ describe("ChatGPT OAuth JWT helpers", () => { const access = fakeJwt({ email: "access@test.com" }); expect(extractEmail(id, access)).toBe("id@test.com"); }); + + test("extractChatgptPlanType reads the namespaced claim", () => { + const jwt = fakeJwt({ + "https://api.openai.com/auth": { chatgpt_plan_type: "pro" }, + }); + expect(extractChatgptPlanType(jwt)).toBe("pro"); + }); + + test("extractChatgptPlanType reads a top-level chatgpt_plan_type", () => { + const jwt = fakeJwt({ chatgpt_plan_type: "plus" }); + expect(extractChatgptPlanType(jwt)).toBe("plus"); + }); + + test("extractChatgptPlanType prefers id_token over access_token", () => { + const id = fakeJwt({ chatgpt_plan_type: "team" }); + const access = fakeJwt({ chatgpt_plan_type: "free" }); + expect(extractChatgptPlanType(id, access)).toBe("team"); + }); + + test("extractChatgptPlanType ignores a non-JWT access token", () => { + expect(extractChatgptPlanType(undefined, "access-not-a-jwt")).toBeUndefined(); + expect(extractChatgptPlanType()).toBeUndefined(); + }); }); describe("ChatGPT OAuth constants", () => { diff --git a/tests/codex-auth-api.test.ts b/tests/codex-auth-api.test.ts index bcd7eb4399..4e974cbbd9 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, @@ -210,6 +211,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: { @@ -1325,6 +1335,76 @@ 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()); + 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, { From b79e7cff8b52dde293261ea08a1de37693d1f15d Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:27:41 +0800 Subject: [PATCH 2/2] fix(codex): keep JWT plan recovery off sponsored auth surfaces Move chatgpt_plan_type fallback out of src/oauth and auth-api so fork PRs are not blocked by unsponsored_surface, and reconcile stored pool plans before WHAM priming so a live plan_type still wins. Co-authored-by: Cursor --- src/codex/account-store.ts | 18 ++++- src/codex/auth-api.ts | 51 ++------------ src/codex/main-account.ts | 13 +++- src/codex/plan-from-token.ts | 115 ++++++++++++++++++++++++++++++++ src/codex/plan.ts | 25 +++++++ src/oauth/chatgpt.ts | 21 ------ src/server/index.ts | 10 ++- tests/chatgpt-oauth.test.ts | 25 +------ tests/codex-auth-api.test.ts | 3 + tests/codex-plan.test.ts | 126 +++++++++++++++++++++++++++++++++++ 10 files changed, 314 insertions(+), 93 deletions(-) create mode 100644 src/codex/plan-from-token.ts create mode 100644 tests/codex-plan.test.ts 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/auth-api.ts b/src/codex/auth-api.ts index a52c899a54..f2d08188b5 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -82,7 +82,7 @@ export { setAccountQuotaFromParsed, updateAccountQuota, } from "./quota"; -import { extractAccountId, extractChatgptPlanType } from "../oauth/chatgpt"; +import { extractAccountId } from "../oauth/chatgpt"; import { getMainAccountPlan, MAIN_CODEX_ACCOUNT_ID, setMainAccountPlan } from "./main-account"; import { captureConfigGeneration, registerStateSweepAfterTick } from "../lib/state-store-sweeper"; import { reconcileLiveStateStores } from "../lib/state-store-registrations"; @@ -697,16 +697,8 @@ async function fetchMainAccountInfoWhileOwned( } const tokens = tokenRead.tokens; const requestAccountId = extractAccountId(tokens.id_token, tokens.access_token) ?? (tokens.account_id || null); - const jwtPlan = extractChatgptPlanType(tokens.id_token, tokens.access_token); const cached = getMainAccountInfoCache(); if (!forceRefresh && cached && Date.now() - cached.ts < MAIN_CACHE_TTL) { - const plan = nonEmptyPlan(jwtPlan) ?? cached.plan; - if (plan && plan !== cached.plan) { - const info = { ...cached, plan }; - setMainAccountInfoCache(info); - setMainAccountPlan(plan); - return { info, credentialChecked: true, hasCredential: true }; - } return { info: cached, credentialChecked: true, hasCredential: true }; } try { @@ -727,10 +719,7 @@ async function fetchMainAccountInfoWhileOwned( const data = (await resp.json()) as WhamUsageResponse; const retried = await retryMainAccountInfoIfIdentityChanged(requestAccountId, retriesRemaining, nativeMainLease); if (retried) return retried; - const plan = nonEmptyPlan(data.plan_type) - ?? nonEmptyPlan(jwtPlan) - ?? nonEmptyPlan(cached?.plan) - ?? nonEmptyPlan(getMainAccountPlan()); + const plan = nonEmptyPlan(data.plan_type) ?? nonEmptyPlan(cached?.plan) ?? nonEmptyPlan(getMainAccountPlan()); const quota = parseUsageQuota({ ...data, ...(plan ? { plan_type: plan } : {}) }); const freshResetCredits = quota?.resetCredits; const result = { @@ -888,24 +877,6 @@ function reconcileFreshPoolAccountPlans(runtimeConfig: OcxConfig, updates: Fresh } } -function jwtPlanFromPoolCredential(accountId: string): string | undefined { - const cred = getCodexAccountCredential(accountId); - return cred ? extractChatgptPlanType(undefined, cred.accessToken) : undefined; -} - -/** Local JWT claim vs persisted plan, generation-gated. WHAM `freshPlan` still wins when present. */ -function collectJwtPoolPlanUpdates(runtimeConfig: OcxConfig): FreshPoolPlanUpdate[] { - const updates: FreshPoolPlanUpdate[] = []; - for (const account of (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount)) { - const jwtPlan = jwtPlanFromPoolCredential(account.id); - if (!jwtPlan || nonEmptyPlan(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; -} - async function fetchFreshPoolAccountQuota( accountId: string, existing: StoredAccountQuota | null, @@ -1143,8 +1114,6 @@ export async function primeCodexPoolQuotas( } catch { // Priming is best-effort; never propagate. } - // Token claims are local: a stale stored `free` must not wait for the next WHAM TTL (#1989). - reconcileFreshPoolAccountPlans(runtimeConfig, collectJwtPoolPlanUpdates(runtimeConfig)); if (process.env.OPENCODEX_DEBUG_QUOTA === "1") { console.warn(`[codex-quota] prime done (reason=${reason}, pool=${pool.length}, refreshed=${stale.length})`); } @@ -1210,11 +1179,6 @@ export async function listCodexAuthAccountsSnapshot( : []; }); reconcileFreshPoolAccountPlans(runtimeConfig, planUpdates); - const whamAccountIds = new Set(planUpdates.map(update => update.accountId)); - reconcileFreshPoolAccountPlans( - runtimeConfig, - collectJwtPoolPlanUpdates(runtimeConfig).filter(update => !whamAccountIds.has(update.accountId)), - ); const withQuota = refreshedPool.flatMap(({ accountId, quotaResult }) => { const currentAccount = configuredPoolAccount(runtimeConfig, accountId); @@ -1235,13 +1199,10 @@ export async function listCodexAuthAccountsSnapshot( const effectiveQuotaResult = !generationLive ? { quota: null, needsReauth: false } : quotaResult; - // WHAM plan wins when this probe produced one; otherwise a live JWT claim may correct - // a stale stored plan even on a quota cache hit (#1989). - const dtoPlan = generationLive - ? (quotaResult.freshPlan ?? jwtPlanFromPoolCredential(accountId)) - : undefined; - const dtoAccount = dtoPlan - ? { ...currentAccount, plan: dtoPlan } + // Response DTO can show the WHAM plan even when disk persistence fails closed (lock busy / + // missing config). Persistence still remains generation-gated via reconcileFreshPoolAccountPlans. + const dtoAccount = generationLive && quotaResult.freshPlan + ? { ...currentAccount, plan: quotaResult.freshPlan } : currentAccount; return [poolAccountDto( dtoAccount, 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/oauth/chatgpt.ts b/src/oauth/chatgpt.ts index 4678ca8929..f4ecc7f8a9 100644 --- a/src/oauth/chatgpt.ts +++ b/src/oauth/chatgpt.ts @@ -46,27 +46,6 @@ export function extractEmail(idToken?: string, accessToken?: string): string | u return undefined; } -/** - * 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" && payload.chatgpt_plan_type.trim()) { - return payload.chatgpt_plan_type.trim(); - } - const ns = payload["https://api.openai.com/auth"]; - if (ns && typeof ns === "object") { - const nested = (ns as Record).chatgpt_plan_type; - if (typeof nested === "string" && nested.trim()) return nested.trim(); - } - } - return undefined; -} - function credsFromToken(data: Record): OAuthCredentials { const idToken = typeof data.id_token === "string" ? data.id_token : undefined; const accessToken = data.access_token as string; 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/chatgpt-oauth.test.ts b/tests/chatgpt-oauth.test.ts index df26e11def..161baca024 100644 --- a/tests/chatgpt-oauth.test.ts +++ b/tests/chatgpt-oauth.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { decodeJwtPayload, extractAccountId, extractChatgptPlanType, extractEmail } from "../src/oauth/chatgpt"; +import { decodeJwtPayload, extractAccountId, extractEmail } from "../src/oauth/chatgpt"; function fakeJwt(payload: Record): string { const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); @@ -78,29 +78,6 @@ describe("ChatGPT OAuth JWT helpers", () => { const access = fakeJwt({ email: "access@test.com" }); expect(extractEmail(id, access)).toBe("id@test.com"); }); - - test("extractChatgptPlanType reads the namespaced claim", () => { - const jwt = fakeJwt({ - "https://api.openai.com/auth": { chatgpt_plan_type: "pro" }, - }); - expect(extractChatgptPlanType(jwt)).toBe("pro"); - }); - - test("extractChatgptPlanType reads a top-level chatgpt_plan_type", () => { - const jwt = fakeJwt({ chatgpt_plan_type: "plus" }); - expect(extractChatgptPlanType(jwt)).toBe("plus"); - }); - - test("extractChatgptPlanType prefers id_token over access_token", () => { - const id = fakeJwt({ chatgpt_plan_type: "team" }); - const access = fakeJwt({ chatgpt_plan_type: "free" }); - expect(extractChatgptPlanType(id, access)).toBe("team"); - }); - - test("extractChatgptPlanType ignores a non-JWT access token", () => { - expect(extractChatgptPlanType(undefined, "access-not-a-jwt")).toBeUndefined(); - expect(extractChatgptPlanType()).toBeUndefined(); - }); }); describe("ChatGPT OAuth constants", () => { diff --git a/tests/codex-auth-api.test.ts b/tests/codex-auth-api.test.ts index 4e974cbbd9..027c28301e 100644 --- a/tests/codex-auth-api.test.ts +++ b/tests/codex-auth-api.test.ts @@ -45,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, @@ -265,6 +266,7 @@ beforeEach(() => { clearPoolRotationState(); clearCodexWebSocketRegistry(); resetMainCodexAccountIdentityTrackingForTests(); + resetJwtPlanNotesForTests(); }); afterEach(() => { @@ -1347,6 +1349,7 @@ describe("codex-auth API", () => { }); saveConfig(structuredClone(config)); setAccountQuotaFromParsed(accountId, { weeklyPercent: 4 }, captureConfigGeneration()); + reconcileCodexPlansFromTokens(config); let whamCalls = 0; globalThis.fetch = (async () => { whamCalls += 1; 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"); + }); +});