-
Notifications
You must be signed in to change notification settings - Fork 872
fix(codex): re-derive pool plan from JWT chatgpt_plan_type (#1989) #1998
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
|
Comment on lines
+26
to
+28
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When Useful? React with 👍 / 👎. |
||
| if (jwtPlan) mainAccountPlan = jwtPlan; | ||
| return jwtPlan; | ||
|
Comment on lines
21
to
+30
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift Invalidate the JWT-derived cache when token material changes. Line 22 returns the prior JWT-derived plan indefinitely. Line 23 returns Track the token value or a token fingerprint with the derived cache. Preserve a fresh WHAM plan, but re-derive fallback data after the token changes. Add a regression test that replaces the token between two calls. The PR objective requires current JWT evidence when WHAM has not supplied a fresh plan. 🤖 Prompt for AI Agents |
||
| } | ||
|
|
||
| /** Read-only main account token from ~/.codex/auth.json, or null when not logged in. */ | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
|
Comment on lines
+35
to
+37
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When another OpenCodex process replaces this credential between the two file reads, Useful? React with 👍 / 👎. |
||
| if (generation === undefined) continue; | ||
| updates.push({ accountId: account.id, plan: jwtPlan, credentialGeneration: generation }); | ||
| } | ||
| return updates; | ||
| } | ||
|
|
||
| const appliedJwtPlans = new Map<string, string>(); | ||
|
|
||
| function persistJwtPlanUpdates(runtimeConfig: OcxConfig, updates: FreshPoolPlanUpdate[]): void { | ||
| if (updates.length === 0) return; | ||
| let outcome: ReturnType<typeof mutatePersistedConfig<FreshPoolPlanUpdate[]>>; | ||
| 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; | ||
|
Comment on lines
+53
to
+61
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift Preserve a fresh WHAM plan before writing a JWT fallback. Line 59 overwrites every differing Track plan provenance or fresh-WHAM state. Apply the JWT value only when WHAM has not supplied a current 🤖 Prompt for AI Agents |
||
| } | ||
| } | ||
| 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; | ||
|
Comment on lines
+100
to
+102
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
If startup records plan Useful? React with 👍 / 👎. |
||
| 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); | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1728,7 +1728,15 @@ export function startServer(port?: number, deps: StartServerDeps = {}): Server<W | |
| && isCanonicalOpenAiForwardProvider(openAiProvider) | ||
| && providerCodexAccountMode("openai", openAiProvider) === "pool" | ||
| ) { | ||
| import("../codex/auth-api") | ||
| import("../codex/plan-from-token") | ||
| .then(({ reconcileCodexPlansFromTokens }) => { | ||
| try { | ||
| reconcileCodexPlansFromTokens(config); | ||
| } catch { | ||
| // Derived plan metadata must not block WHAM priming. | ||
| } | ||
| return import("../codex/auth-api"); | ||
| }) | ||
| .then(({ primeCodexPoolQuotas }) => primeCodexPoolQuotas(config, "startup")) | ||
| .catch(() => {}); | ||
|
Comment on lines
+1731
to
1741
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win Add a startup-order regression test. This changes startup behavior. The changed tests call Add a flat Bun test that starts a pool-mode server with a stale stored plan and a JWT-derived plan. Assert that quota priming observes the reconciled plan before its first WHAM request. As per path instructions, “A behavior change in src/ should come with a focused regression test near the existing tests for that subsystem.” 🤖 Prompt for AI AgentsSource: Path instructions |
||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
After the first lookup, replacing
auth.jsonwith a refreshed token for the same ChatGPT account does not reset eithermainAccountPlanorjwtPlanAttempted, because the identity reconciliation only purges state when the account ID changes. Consequently an upgrade or downgrade reflected in a new access token keeps returning the old plan until WHAM supplies an explicit plan or some unrelated reset occurs. Bind this cache to the observed token/account snapshot or invalidate it when the token file changes.Useful? React with 👍 / 👎.