diff --git a/src/codex/account-store.ts b/src/codex/account-store.ts index 85eaab3ef8..bf4bcced74 100644 --- a/src/codex/account-store.ts +++ b/src/codex/account-store.ts @@ -11,13 +11,15 @@ import { withConfigMutationLockSync, } from "../config"; import { assertNotRealHomeUnderTest } from "../lib/test-home-guard"; +import { CHATGPT_CLIENT_ID, CHATGPT_TOKEN_URL } from "../oauth/chatgpt"; import type { CodexAccountCredentialRecord, CodexAccountCredentials } from "../types"; type LegacyCodexAccountStore = Record; type CodexAccountStore = Record; type RawCodexAccountStore = Record; -const REFRESH_SKEW_MS = 60_000; +export const CODEX_REFRESH_SKEW_MS = 60_000; +const REFRESH_SKEW_MS = CODEX_REFRESH_SKEW_MS; const REFRESH_LOCK_STALE_MS = 60_000; const REFRESH_LOCK_WAIT_MS = REFRESH_LOCK_STALE_MS + 5_000; const REFRESH_LOCK_POLL_MS = 50; @@ -47,6 +49,23 @@ function isCredential(value: unknown): value is CodexAccountCredentials { && typeof value.chatgptAccountId === "string"; } +function requiredTokenResponseString(data: Record, field: string): string { + const value = data[field]; + if (typeof value !== "string" || value.trim().length === 0) { + throw new TokenRefreshError("unknown", `Codex token refresh returned a malformed ${field}.`); + } + return value; +} + +function optionalTokenResponseString(data: Record, field: string): string | undefined { + const value = data[field]; + if (value === undefined || value === null) return undefined; + if (typeof value !== "string" || value.trim().length === 0) { + throw new TokenRefreshError("unknown", `Codex token refresh returned a malformed ${field}.`); + } + return value; +} + function isCredentialRecord(value: unknown): value is CodexAccountCredentialRecord { return isObject(value) && typeof value.generation === "number" @@ -229,13 +248,10 @@ export function tombstoneCodexAccount(id: string): number { }); } -const CHATGPT_TOKEN_URL = "https://auth.openai.com/oauth/token"; -const CHATGPT_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"; - export class TokenRefreshError extends Error { reason: "expired" | "revoked" | "unknown"; - constructor(reason: "expired" | "revoked" | "unknown", message: string) { - super(message); + constructor(reason: "expired" | "revoked" | "unknown", message: string, options?: ErrorOptions) { + super(message, options); this.name = "TokenRefreshError"; this.reason = reason; } @@ -331,7 +347,7 @@ function isRefreshLockStale(path: string): boolean { } } -async function withCodexRefreshFileLock(lockKey: string, signal: AbortSignal, fn: () => Promise): Promise { +export async function withCodexRefreshFileLock(lockKey: string, signal: AbortSignal, fn: () => Promise): Promise { hardenConfigDir(); const dir = getConfigDir(); if (!existsSync(dir)) mkdirSync(dir, { recursive: true, mode: 0o700 }); @@ -372,20 +388,68 @@ async function withCodexRefreshFileLock(lockKey: string, signal: AbortSignal, } } -function findFreshCredentialForGrant( - refreshGrantFingerprint: string, - excludeId: string, -): CodexAccountCredentials | null { - const now = Date.now(); +export function findFreshCredentialForGrant(query: { + refreshGrantFingerprint: string; + excludeId?: string; + now?: number; +}): CodexAccountCredentials | null { + const now = query.now ?? Date.now(); const records = loadCodexAccountRecordStore(); for (const [candidateId, candidate] of Object.entries(records)) { - if (candidateId === excludeId || candidate.deletedAt != null || !candidate.credential) continue; - if (recordGrantFingerprint(candidate) !== refreshGrantFingerprint) continue; + if (candidateId === query.excludeId || candidate.deletedAt != null || !candidate.credential) continue; + if (recordGrantFingerprint(candidate) !== query.refreshGrantFingerprint) continue; if (candidate.credential.expiresAt > now + REFRESH_SKEW_MS) return candidate.credential; } return null; } +export function findUniqueFreshCredentialForChatgptAccount(query: { + chatgptAccountId: string; + excludeId?: string; + now?: number; +}): CodexAccountCredentials | null { + const chatgptAccountId = query.chatgptAccountId.trim(); + if (!chatgptAccountId) return null; + const now = query.now ?? Date.now(); + const records = loadCodexAccountRecordStore(); + let match: CodexAccountCredentials | null = null; + for (const [candidateId, candidate] of Object.entries(records)) { + if (candidateId === query.excludeId || candidate.deletedAt != null || !candidate.credential) continue; + if (candidate.credential.chatgptAccountId !== chatgptAccountId) continue; + if (candidate.credential.expiresAt <= now + REFRESH_SKEW_MS) continue; + if (match) return null; + match = candidate.credential; + } + return match; +} + +export function publishFreshCredentialForGrant(query: { + refreshGrantFingerprint: string; + credential: CodexAccountCredentials; + excludeId?: string; +}): void { + withCredentialMutationLockSync(() => { + const store = loadCodexAccountRecordStore(); + let changed = false; + for (const [candidateId, candidate] of Object.entries(store)) { + if (candidateId === query.excludeId || candidate.deletedAt != null || !candidate.credential) continue; + if (recordGrantFingerprint(candidate) !== query.refreshGrantFingerprint) continue; + store[candidateId] = { + credential: { + ...query.credential, + chatgptAccountId: candidate.credential.chatgptAccountId || query.credential.chatgptAccountId, + }, + generation: candidate.generation + 1, + refreshGrantFingerprint: refreshGrantFingerprintForToken(query.credential.refreshToken), + replacedAt: candidate.replacedAt, + ...preservedValidationMetadata(candidate), + }; + changed = true; + } + if (changed) persist(store); + }); +} + async function notePlanFromRefreshedAccessToken( id: string, accessToken: string, @@ -471,7 +535,7 @@ export async function getValidCodexToken(id: string): Promise credential: lockedCred, }; } - const sameGrantFreshCredential = findFreshCredentialForGrant(refreshGrantFingerprint, id); + const sameGrantFreshCredential = findFreshCredentialForGrant({ refreshGrantFingerprint, excludeId: id }); if (sameGrantFreshCredential) { if (!saveCodexAccountCredentialIfGeneration(id, startGeneration, sameGrantFreshCredential)) { throw new CodexCredentialGenerationConflictError(); @@ -505,7 +569,9 @@ export async function getValidCodexToken(id: string): Promise : "unknown" as const; throw new TokenRefreshError(reason, `Codex token refresh failed (${reason}); reauthenticate the account.`); } - const data = (await res.json()) as { access_token: string; refresh_token?: string; expires_in: number }; + const data = (await res.json()) as Record; + const accessToken = requiredTokenResponseString(data, "access_token"); + const refreshToken = optionalTokenResponseString(data, "refresh_token"); // Guard against a missing/non-finite/negative expires_in (malformed upstream // response): a NaN expiry would never compare as expired, and a negative // duration would stamp an already-past expiry — both block refresh semantics. @@ -519,8 +585,8 @@ export async function getValidCodexToken(id: string): Promise const safeExpiresAt = Number.isFinite(expiresAt) ? expiresAt : Date.now() + 3600 * 1000; const updated: CodexAccountCredentials = { - accessToken: data.access_token, - refreshToken: data.refresh_token ?? lockedCred.refreshToken, + accessToken, + refreshToken: refreshToken ?? lockedCred.refreshToken, expiresAt: safeExpiresAt, chatgptAccountId: lockedCred.chatgptAccountId, }; diff --git a/src/codex/account-usability.ts b/src/codex/account-usability.ts index d508e19f4d..045d2b934c 100644 --- a/src/codex/account-usability.ts +++ b/src/codex/account-usability.ts @@ -1,6 +1,6 @@ import { getCodexAccountCredential } from "./account-store"; import { isAccountNeedsReauth } from "./account-runtime-state"; -import { MAIN_CODEX_ACCOUNT_ID, isMainAccountTokenLive } from "./main-account"; +import { MAIN_CODEX_ACCOUNT_ID, isMainAccountCredentialUsable } from "./main-account"; import { hasLegacyMainCodexPoolAccount, isSelectableCodexPoolAccount } from "./account-id"; import type { OcxConfig } from "../types"; import { isNativeMainTrafficBlocked } from "./native-profile-startup"; @@ -9,7 +9,7 @@ export interface CodexAccountUsabilityOptions { /** Route using cached runtime state only; the caller must reject selected main before auth. */ nativeMainSelectionOnly?: boolean; /** Test seam for proving whether routing attempted a physical native-token read. */ - isMainAccountTokenLive?: typeof isMainAccountTokenLive; + isMainAccountTokenLive?: typeof isMainAccountCredentialUsable; /** Confirmed account ids for an account-gated model; omitted for ordinary native models. */ modelEligibleAccountIds?: ReadonlySet; } @@ -32,8 +32,8 @@ export function isCodexAccountUsable( // before reservation or token materialization. Treat cached main as a routing // candidate without touching the credential file so affinity is not rebound. if (options.nativeMainSelectionOnly) return true; - // Main account: credential is the read-only ~/.codex/auth.json token (Option A). - return (options.isMainAccountTokenLive ?? isMainAccountTokenLive)(); + // Main account: a refresh grant is enough to route; materialization refreshes before I/O. + return (options.isMainAccountTokenLive ?? isMainAccountCredentialUsable)(); } const exists = (config.codexAccounts ?? []) .some(account => isSelectableCodexPoolAccount(account) && account.id === accountId); diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 7dd58c6b91..718aac67af 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -4,6 +4,7 @@ import { CodexCredentialRefreshLockTimeoutError, CodexCredentialRefreshBusyError, CodexCredentialRefreshStaleError, + TokenRefreshError, getValidCodexToken, isCodexAccountGenerationLive, } from "./account-store"; @@ -12,7 +13,14 @@ import { isCodexAccountPaused } from "./account-pause"; import { ConfigMutationLockError } from "../config"; import { isCodexAccountUsable } from "./account-usability"; import { reconcileMainCodexAccountRuntimeState } from "./account-lifecycle"; -import { MAIN_CODEX_ACCOUNT_ID, getMainAccountToken, isMainAccountTokenLive } from "./main-account"; +import { + MAIN_CODEX_ACCOUNT_ID, + MainAuthJsonChangedDuringRefreshError, + getMainAccountToken, + getValidMainAccountToken, + isMainAccountTokenLive, + type NativeMainRefreshDependencies, +} from "./main-account"; import { isNativeMainTrafficBlocked, nativeMainStartupGateSnapshot } from "./native-profile-startup"; import type { NativeMainStartupBlockReason } from "./native-profile-startup"; import { @@ -29,6 +37,7 @@ import { entitledCodexAccountIdsForModel, isDirectCallerEntitledToCodexModel, resolveCodexModelEntitlements, + type CodexModelEntitlementResolveOptions, type CodexModelEntitlementSnapshot, } from "./model-entitlements"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "./catalog/native-models"; @@ -318,9 +327,16 @@ export function shouldMarkAccountNeedsReauthForCodexAuthFailure(cause: unknown): && !(cause instanceof CodexCredentialRefreshLockTimeoutError) && !(cause instanceof CodexCredentialRefreshBusyError) && !(cause instanceof CodexCredentialRefreshStaleError) + && !(cause instanceof MainAuthJsonChangedDuringRefreshError) + && !(cause instanceof TokenRefreshError && cause.reason === "unknown") && !(cause instanceof ConfigMutationLockError); } +type ResolveCodexModelEntitlementsForAuthContext = ( + config: Pick, + options?: Pick, +) => Promise; + export interface ResolveCodexAuthContextOptions { excludeAccountId?: string; /** Resolve exactly this account without consulting or mutating Pool selection. */ @@ -332,17 +348,27 @@ export interface ResolveCodexAuthContextOptions { /** Test-only native credential read seams. */ isMainAccountTokenLive?: () => boolean; getMainAccountToken?: typeof getMainAccountToken; + getValidMainAccountToken?: typeof getValidMainAccountToken; + nativeMainRefreshDependencies?: NativeMainRefreshDependencies; + signal?: AbortSignal; primeCodexPoolQuotas?: (config: OcxConfig, reason: string) => Promise; /** Test seam for account-gated native model discovery. */ - resolveCodexModelEntitlements?: ( - config: Pick, - ) => Promise; + resolveCodexModelEntitlements?: ResolveCodexModelEntitlementsForAuthContext; /** Direct requests admitted with a proxy bearer substitute the stored native-main credential. */ substituteMainCredentialForDirect?: boolean; /** Test seam for a Direct request's own forwarded ChatGPT credential. */ isDirectCallerEntitledToCodexModel?: (headers: Headers, modelId: string) => Promise; } +function entitlementOptionsForAuthContext( + options: Pick, +): Pick { + return { + ...(options.nativeMainRefreshDependencies ? { nativeMainRefreshDependencies: options.nativeMainRefreshDependencies } : {}), + ...(options.signal ? { signal: options.signal } : {}), + }; +} + export interface CodexAccountSelectionAdmission { readonly mainProfileDraining: boolean; claimMainProfile(): boolean; @@ -367,7 +393,10 @@ export async function resolveCodexAuthContext( if (options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)) { const entitled = options.substituteMainCredentialForDirect ? entitledCodexAccountIdsForModel( - await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config), + await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)( + config, + entitlementOptionsForAuthContext(options), + ), options.modelId, )?.has(MAIN_CODEX_ACCOUNT_ID) === true : await (options.isDirectCallerEntitledToCodexModel ?? isDirectCallerEntitledToCodexModel)( @@ -382,7 +411,10 @@ export async function resolveCodexAuthContext( } const affinityKey = fixedAccountId === undefined ? codexPoolAffinityKey(headers) : undefined; const entitlementSnapshot = options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId) - ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config) + ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)( + config, + entitlementOptionsForAuthContext(options), + ) : undefined; const modelEligibleAccountIds = entitlementSnapshot ? entitledCodexAccountIdsForModel(entitlementSnapshot, options.modelId) @@ -486,7 +518,12 @@ export async function resolveCodexAuthContext( // best-effort prime so the NEXT routing decision has real scores. This never // blocks the current request, and the helper's single-flight guard collapses // repeated triggers into one pass. - if (fixedAccountId === undefined && !nativeMainReadsForbidden && !getAccountQuota(accountId)) { + if ( + accountId !== MAIN_CODEX_ACCOUNT_ID + && fixedAccountId === undefined + && !nativeMainReadsForbidden + && !getAccountQuota(accountId) + ) { if (options.primeCodexPoolQuotas) { void options.primeCodexPoolQuotas(config, "pre-route").catch(() => {}); } else { @@ -520,8 +557,21 @@ export async function resolveCodexAuthContext( } if (accountId === MAIN_CODEX_ACCOUNT_ID) { - // Main account in rotation: inject the read-only auth.json token and fail closed if it vanished. - const token = (options.getMainAccountToken ?? getMainAccountToken)(); + // Main account in rotation: refresh auth.json before upstream I/O and fail closed if it vanished. + let token: { accessToken: string; chatgptAccountId: string } | null; + try { + token = await (options.getValidMainAccountToken ?? getValidMainAccountToken)({ + signal: options.signal, + dependencies: options.nativeMainRefreshDependencies, + }); + } catch (cause) { + if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); + else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); + if (shouldMarkAccountNeedsReauthForCodexAuthFailure(cause)) { + markAccountNeedsReauth(accountId, writerGeneration); + } + throw new CodexAuthContextError(accountId, cause); + } if (!token) { // Nothing will reach upstream, so give the probe back instead of burning it. if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); @@ -530,6 +580,15 @@ export async function resolveCodexAuthContext( fixedAccountId !== undefined ? "Selected Codex account is unavailable" : undefined, ); } + if (fixedAccountId === undefined && !nativeMainReadsForbidden && !getAccountQuota(accountId)) { + if (options.primeCodexPoolQuotas) { + void options.primeCodexPoolQuotas(config, "pre-route").catch(() => {}); + } else { + import("./auth-api") + .then(({ primeCodexPoolQuotas }) => primeCodexPoolQuotas(config, "pre-route")) + .catch(() => {}); + } + } return { kind: "main-pool", accountId, @@ -644,6 +703,36 @@ export function materializeCodexUpstreamAuth( return selected; } +export async function materializeCodexUpstreamAuthAsync( + headers: Headers, + ctx: CodexAuthContext, + options: { + substituteMainCredential?: boolean; + signal?: AbortSignal; + getValidMainAccountToken?: typeof getValidMainAccountToken; + nativeMainRefreshDependencies?: NativeMainRefreshDependencies; + } = {}, +): Promise { + if (ctx.kind !== "main" || options.substituteMainCredential !== true) { + return materializeCodexUpstreamAuth(headers, ctx, options); + } + const selected = new Headers(); + for (const name of FORWARD_HEADERS) { + const value = headers.get(name); + if (value) selected.set(name, value); + } + const stored = await (options.getValidMainAccountToken ?? getValidMainAccountToken)({ + signal: options.signal, + dependencies: options.nativeMainRefreshDependencies, + }); + if (!stored?.accessToken) { + throw new CodexMainSubstitutionUnavailableError(); + } + selected.set("authorization", `Bearer ${stored.accessToken}`); + if (stored.chatgptAccountId) selected.set("chatgpt-account-id", stored.chatgptAccountId); + return selected; +} + /** @deprecated Prefer materializeCodexUpstreamAuth; kept for call sites without admission context. */ export function headersForCodexAuthContext(headers: Headers, ctx: CodexAuthContext): Headers { return materializeCodexUpstreamAuth(headers, ctx); diff --git a/src/codex/main-account.ts b/src/codex/main-account.ts index 30296b586a..321ddd43b9 100644 --- a/src/codex/main-account.ts +++ b/src/codex/main-account.ts @@ -1,7 +1,30 @@ +import { createHash } from "node:crypto"; +import { existsSync, linkSync, mkdirSync, readFileSync, renameSync, statSync, unlinkSync } from "node:fs"; +import { join } from "node:path"; import { readCodexTokens } from "./auth-collision"; -import { decodeJwtPayload } from "../oauth/chatgpt"; +import { clearAccountNeedsReauth, markAccountNeedsReauth } from "./account-runtime-state"; +import { + CODEX_REFRESH_SKEW_MS, + TokenRefreshError, + findFreshCredentialForGrant, + findUniqueFreshCredentialForChatgptAccount, + publishFreshCredentialForGrant, + refreshGrantFingerprintForToken, + withCodexRefreshFileLock, +} from "./account-store"; +import { + ChatGPTTokenRefreshError, + decodeJwtPayload, + extractAccountId, + refreshChatGPTToken, +} from "../oauth/chatgpt"; +import type { OAuthCredentials } from "../oauth/types"; import { extractChatgptPlanType } from "./plan"; import { MAIN_CODEX_ACCOUNT_ID } from "./account-id"; +import { atomicWriteFile } from "../config"; +import { assertNotRealCodexHomeUnderTest } from "../lib/test-home-guard"; +import { resolveCodexHomeDir } from "./home"; +import type { CodexAccountCredentials } from "../types"; export { MAIN_CODEX_ACCOUNT_ID } from "./account-id"; @@ -12,6 +35,433 @@ export { MAIN_CODEX_ACCOUNT_ID } from "./account-id"; */ let mainAccountPlan: string | null = null; let jwtPlanAttempted = false; +let beforeMainAuthJsonPublishForTests: (() => void) | null = null; +let beforeMainAuthJsonRenameForTests: (() => void) | null = null; +let beforeMainAuthJsonReplaceForTests: (() => void) | null = null; +let mainAuthJsonBackupSequence = 0; +const nativeMainRefreshFlights = new Map>(); + +interface AuthJsonSnapshot { + dev: number; + ino: number; + mtimeMs: number; + size: number; + hash: string; +} + +interface MainAuthJsonCredential { + path: string; + root: Record; + tokens: Record; + snapshot: AuthJsonSnapshot; + accessToken?: string; + refreshToken?: string; + idToken?: string; + chatgptAccountId: string; +} + +type MainAuthJsonReadResult = + | { status: "ok"; auth: MainAuthJsonCredential } + | { status: "missing" | "invalid" | "unreadable" }; + +type MainAuthPersistGuardResult = + | { status: "current"; auth: MainAuthJsonCredential } + | { status: "adopted"; token: { accessToken: string; chatgptAccountId: string } }; + +export interface NativeMainRefreshDependencies { + refreshToken?: (refreshToken: string, options: { signal?: AbortSignal }) => Promise; +} + +export class MainAuthJsonChangedDuringRefreshError extends Error { + readonly code = "CODEX_MAIN_AUTH_CHANGED"; + readonly retryable = true; + + constructor() { + super("Codex main auth.json changed during refresh"); + this.name = "MainAuthJsonChangedDuringRefreshError"; + } +} + +function mainAuthJsonPath(): string { + return join(resolveCodexHomeDir(), "auth.json"); +} + +function isObject(value: unknown): value is Record { + return !!value && typeof value === "object" && !Array.isArray(value); +} + +function nonEmptyString(value: unknown): string | undefined { + return typeof value === "string" && value.trim().length > 0 ? value : undefined; +} + +function snapshotFor(path: string, raw: string): AuthJsonSnapshot { + const stat = statSync(path); + return { + dev: stat.dev, + ino: stat.ino, + mtimeMs: stat.mtimeMs, + size: stat.size, + hash: createHash("sha256").update(raw).digest("hex"), + }; +} + +function sameSnapshot(a: AuthJsonSnapshot, b: AuthJsonSnapshot): boolean { + return a.dev === b.dev + && a.ino === b.ino + && a.mtimeMs === b.mtimeMs + && a.size === b.size + && a.hash === b.hash; +} + +function pathErrorCode(error: unknown): string | undefined { + return (error as NodeJS.ErrnoException | undefined)?.code; +} + +function isExistingPathError(error: unknown): boolean { + return pathErrorCode(error) === "EEXIST"; +} + +function isMissingPathError(error: unknown): boolean { + return pathErrorCode(error) === "ENOENT"; +} + +function assertMainAuthJsonSnapshotUnchanged(current: MainAuthJsonCredential): void { + const latest = readMainAuthJsonCredential(); + if (latest.status !== "ok" || !sameSnapshot(latest.auth.snapshot, current.snapshot)) { + throw new MainAuthJsonChangedDuringRefreshError(); + } +} + +function readAuthJsonRaw(path: string): { raw: string; snapshot: AuthJsonSnapshot } | MainAuthJsonReadResult { + try { + const firstStat = statSync(path); + const firstRaw = readFileSync(path, "utf-8"); + const secondStat = statSync(path); + if ( + firstStat.dev === secondStat.dev + && firstStat.ino === secondStat.ino + && firstStat.mtimeMs === secondStat.mtimeMs + && firstStat.size === secondStat.size + ) { + return { raw: firstRaw, snapshot: snapshotFor(path, firstRaw) }; + } + const raw = readFileSync(path, "utf-8"); + return { raw, snapshot: snapshotFor(path, raw) }; + } catch (error) { + const code = typeof error === "object" && error !== null && "code" in error + ? (error as { code?: unknown }).code + : undefined; + return { status: code === "ENOENT" ? "missing" : "unreadable" }; + } +} + +function authJsonSnapshotMatches(path: string, snapshot: AuthJsonSnapshot): boolean { + const current = readAuthJsonRaw(path); + return !("status" in current) && sameSnapshot(current.snapshot, snapshot); +} + +function readMainAuthJsonCredential(): MainAuthJsonReadResult { + const path = mainAuthJsonPath(); + const raw = readAuthJsonRaw(path); + if ("status" in raw) return raw; + try { + const parsed = JSON.parse(raw.raw) as unknown; + if (!isObject(parsed) || !isObject(parsed.tokens)) return { status: "invalid" }; + const accessToken = nonEmptyString(parsed.tokens.access_token); + const refreshToken = nonEmptyString(parsed.tokens.refresh_token); + if (!accessToken && !refreshToken) return { status: "invalid" }; + const idToken = nonEmptyString(parsed.tokens.id_token); + const accountId = nonEmptyString(parsed.tokens.account_id); + return { + status: "ok", + auth: { + path, + root: parsed, + tokens: parsed.tokens, + snapshot: raw.snapshot, + ...(accessToken ? { accessToken } : {}), + ...(refreshToken ? { refreshToken } : {}), + ...(idToken ? { idToken } : {}), + chatgptAccountId: extractAccountId(idToken, accessToken) ?? accountId ?? "", + }, + }; + } catch { + return { status: "invalid" }; + } +} + +export function mainAccessTokenFresh( + accessToken: string | undefined, + now = Date.now(), + skewMs = 0, +): boolean { + if (!accessToken) return false; + const payload = decodeJwtPayload(accessToken); + const exp = typeof payload?.exp === "number" ? payload.exp * 1000 : undefined; + return exp === undefined || exp > now + skewMs; +} + +function tokenResultFromAuth( + auth: MainAuthJsonCredential, + options: { rejectedAccessToken?: string; now?: number; skewMs?: number } = {}, +): { accessToken: string; chatgptAccountId: string } | null { + if (!auth.accessToken) return null; + if (auth.accessToken === options.rejectedAccessToken) return null; + if (!mainAccessTokenFresh(auth.accessToken, options.now ?? Date.now(), options.skewMs ?? 0)) return null; + return { accessToken: auth.accessToken, chatgptAccountId: auth.chatgptAccountId }; +} + +function tokenResultFromCredential( + credential: CodexAccountCredentials, +): { accessToken: string; chatgptAccountId: string } { + return { accessToken: credential.accessToken, chatgptAccountId: credential.chatgptAccountId }; +} + +function combineNativeRefreshSignal(signal?: AbortSignal): AbortSignal { + const timeout = AbortSignal.timeout(30_000); + return signal ? AbortSignal.any([signal, timeout]) : timeout; +} + +async function joinNativeMainRefreshFlight( + flight: Promise<{ accessToken: string; chatgptAccountId: string } | null>, + signal: AbortSignal, +): Promise<{ accessToken: string; chatgptAccountId: string } | null> { + if (signal.aborted) throw signal.reason; + let abortListener: (() => void) | undefined; + const abort = new Promise((_resolve, reject) => { + abortListener = () => reject(signal.reason); + signal.addEventListener("abort", abortListener, { once: true }); + }); + try { + return await Promise.race([flight, abort]); + } finally { + if (abortListener) signal.removeEventListener("abort", abortListener); + } +} + +function tokenRefreshReason(error: unknown): "expired" | "revoked" | "unknown" { + if (error instanceof TokenRefreshError) return error.reason; + if (error instanceof ChatGPTTokenRefreshError) { + const oauthError = error.oauthError?.toLowerCase(); + const description = error.oauthErrorDescription?.toLowerCase() ?? ""; + if (description.includes("expired")) return "expired"; + if (oauthError === "invalid_grant" || description.includes("revoked") || description.includes("invalidated")) { + return "revoked"; + } + } + return "unknown"; +} + +function tokenRefreshError(error: unknown): TokenRefreshError { + if (error instanceof TokenRefreshError) return error; + const reason = tokenRefreshReason(error); + return new TokenRefreshError( + reason, + `Codex main token refresh failed (${reason}); ${reason === "unknown" ? "retry the request" : "reauthenticate the main account"}.`, + { cause: error }, + ); +} + +function refreshedCredentialFromOAuth( + locked: MainAuthJsonCredential, + refreshed: OAuthCredentials, +): CodexAccountCredentials { + return { + accessToken: refreshed.access, + refreshToken: refreshed.refresh || locked.refreshToken!, + expiresAt: refreshed.expires, + chatgptAccountId: refreshed.accountId ?? extractAccountId(undefined, refreshed.access) ?? locked.chatgptAccountId, + }; +} + +function mainAuthJsonCredentialMatches( + auth: MainAuthJsonCredential, + credential: CodexAccountCredentials, +): boolean { + return auth.accessToken === credential.accessToken + && auth.refreshToken === credential.refreshToken + && nonEmptyString(auth.tokens.account_id) === credential.chatgptAccountId; +} + +function assertMainAuthJsonCredentialPersisted(credential: CodexAccountCredentials): void { + const current = readMainAuthJsonCredential(); + if (current.status !== "ok" || !mainAuthJsonCredentialMatches(current.auth, credential)) { + throw new MainAuthJsonChangedDuringRefreshError(); + } +} + +function restoreAuthJsonBackupWithoutReplacing(backupPath: string, targetPath: string): void { + try { + linkSync(backupPath, targetPath); + } catch (error) { + if (!isExistingPathError(error)) throw error; + } + try { + unlinkSync(backupPath); + } catch (error) { + if (!isMissingPathError(error)) throw error; + } +} + +function removePublishedAuthJsonTempBestEffort(path: string): void { + try { + unlinkSync(path); + } catch { + /* Preserve the published auth.json; throwing here would make atomic cleanup scrub a hard-linked target. */ + } +} + +function replaceMainAuthJsonWithoutClobbering( + current: MainAuthJsonCredential, + tempPath: string, + targetPath: string, +): void { + const backupPath = `${targetPath}.ocx-main-auth.${process.pid}.${++mainAuthJsonBackupSequence}.bak`; + try { + renameSync(targetPath, backupPath); + } catch (error) { + if (isMissingPathError(error)) throw new MainAuthJsonChangedDuringRefreshError(); + throw error; + } + + if (!authJsonSnapshotMatches(backupPath, current.snapshot)) { + restoreAuthJsonBackupWithoutReplacing(backupPath, targetPath); + throw new MainAuthJsonChangedDuringRefreshError(); + } + + const replacement = readAuthJsonRaw(tempPath); + if ("status" in replacement) throw new Error("Codex main auth.json replacement temp was unreadable"); + + try { + linkSync(tempPath, targetPath); + } catch (error) { + restoreAuthJsonBackupWithoutReplacing(backupPath, targetPath); + if (isExistingPathError(error)) throw new MainAuthJsonChangedDuringRefreshError(); + throw error; + } + removePublishedAuthJsonTempBestEffort(tempPath); + removePublishedAuthJsonTempBestEffort(backupPath); + if (!authJsonSnapshotMatches(targetPath, replacement.snapshot)) { + throw new MainAuthJsonChangedDuringRefreshError(); + } +} + +function persistMainAuthJson( + current: MainAuthJsonCredential, + credential: CodexAccountCredentials, +): void { + const home = resolveCodexHomeDir(); + assertNotRealCodexHomeUnderTest(home); + if (!existsSync(home)) mkdirSync(home, { recursive: true, mode: 0o700 }); + const tokens = { + ...current.tokens, + access_token: credential.accessToken, + refresh_token: credential.refreshToken, + account_id: credential.chatgptAccountId, + }; + atomicWriteFile( + current.path, + JSON.stringify({ ...current.root, tokens }, null, 2) + "\n", + undefined, + { + replace: (_tempPath, _targetPath) => { + const hook = beforeMainAuthJsonPublishForTests; + beforeMainAuthJsonPublishForTests = null; + hook?.(); + assertMainAuthJsonSnapshotUnchanged(current); + const renameHook = beforeMainAuthJsonRenameForTests; + beforeMainAuthJsonRenameForTests = null; + renameHook?.(); + assertMainAuthJsonSnapshotUnchanged(current); + const replaceHook = beforeMainAuthJsonReplaceForTests; + beforeMainAuthJsonReplaceForTests = null; + replaceHook?.(); + replaceMainAuthJsonWithoutClobbering(current, _tempPath, _targetPath); + }, + }, + ); +} + +function guardMainAuthJsonBeforePersist( + locked: MainAuthJsonCredential, + options: { rejectedAccessToken?: string }, +): MainAuthPersistGuardResult { + const current = readMainAuthJsonCredential(); + if (current.status === "ok" && sameSnapshot(current.auth.snapshot, locked.snapshot)) { + return { status: "current", auth: current.auth }; + } + if (current.status === "ok") { + const adopted = tokenResultFromAuth(current.auth, { + rejectedAccessToken: options.rejectedAccessToken, + skewMs: options.rejectedAccessToken ? 0 : CODEX_REFRESH_SKEW_MS, + }); + if (adopted) return { status: "adopted", token: adopted }; + } + throw new MainAuthJsonChangedDuringRefreshError(); +} + +function freshStoredCredentialForMain( + locked: MainAuthJsonCredential, + refreshGrantFingerprint: string, + rejectedAccessToken?: string, +): CodexAccountCredentials | null { + const sameGrantCredential = findFreshCredentialForGrant({ + refreshGrantFingerprint, + excludeId: MAIN_CODEX_ACCOUNT_ID, + }); + if (sameGrantCredential && sameGrantCredential.accessToken !== rejectedAccessToken) return sameGrantCredential; + const sameAccountCredential = findUniqueFreshCredentialForChatgptAccount({ + chatgptAccountId: locked.chatgptAccountId, + excludeId: MAIN_CODEX_ACCOUNT_ID, + }); + if (sameAccountCredential && sameAccountCredential.accessToken !== rejectedAccessToken) return sameAccountCredential; + return null; +} + +function persistMainStoredCredential( + locked: MainAuthJsonCredential, + credential: CodexAccountCredentials, + rejectedAccessToken?: string, +): { accessToken: string; chatgptAccountId: string } { + const guard = guardMainAuthJsonBeforePersist(locked, { rejectedAccessToken }); + if (guard.status === "adopted") return guard.token; + try { + persistMainAuthJson(guard.auth, credential); + assertMainAuthJsonCredentialPersisted(credential); + } catch (error) { + return adoptFreshMainAuthJsonAfterPersistRace(error, rejectedAccessToken); + } + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); + return tokenResultFromCredential(credential); +} + +function adoptFreshMainAuthJsonAfterPersistRace( + error: unknown, + rejectedAccessToken?: string, +): { accessToken: string; chatgptAccountId: string } { + if (!(error instanceof MainAuthJsonChangedDuringRefreshError)) throw error; + const current = readMainAuthJsonCredential(); + if (current.status === "ok") { + const adopted = tokenResultFromAuth(current.auth, { + rejectedAccessToken, + skewMs: rejectedAccessToken ? 0 : CODEX_REFRESH_SKEW_MS, + }); + if (adopted) return adopted; + } + throw error; +} + +export function setMainAuthJsonPublishHookForTests(hook: (() => void) | null): void { + beforeMainAuthJsonPublishForTests = hook; +} + +export function setMainAuthJsonRenameHookForTests(hook: (() => void) | null): void { + beforeMainAuthJsonRenameForTests = hook; +} + +export function setMainAuthJsonReplaceHookForTests(hook: (() => void) | null): void { + beforeMainAuthJsonReplaceForTests = hook; +} export function setMainAccountPlan(plan: string | null): void { mainAccountPlan = plan; @@ -32,9 +482,9 @@ export function getMainAccountPlan(): string | undefined { /** Read-only main account token from ~/.codex/auth.json, or null when not logged in. */ export function getMainAccountToken(): { accessToken: string; chatgptAccountId: string } | null { - const tokens = readCodexTokens(); - if (!tokens?.access_token) return null; - return { accessToken: tokens.access_token, chatgptAccountId: tokens.account_id }; + const result = readMainAuthJsonCredential(); + if (result.status !== "ok" || !result.auth.accessToken) return null; + return { accessToken: result.auth.accessToken, chatgptAccountId: result.auth.chatgptAccountId }; } /** @@ -43,11 +493,17 @@ export function getMainAccountToken(): { accessToken: string; chatgptAccountId: * actually-invalid token then surfaces via the upstream 401 → cooldown path. */ export function isMainAccountTokenLive(now = Date.now()): boolean { - const tokens = readCodexTokens(); - if (!tokens?.access_token) return false; - const payload = decodeJwtPayload(tokens.access_token); - const exp = typeof payload?.exp === "number" ? payload.exp * 1000 : undefined; - return exp === undefined || exp > now; + const result = readMainAuthJsonCredential(); + if (result.status !== "ok") return false; + return mainAccessTokenFresh(result.auth.accessToken, now); +} + +/** A main account is routeable when it has a live access token or a refresh grant. */ +export function isMainAccountCredentialUsable(now = Date.now()): boolean { + const result = readMainAuthJsonCredential(); + if (result.status !== "ok") return false; + if (result.auth.refreshToken) return true; + return mainAccessTokenFresh(result.auth.accessToken, now); } /** @@ -60,9 +516,92 @@ export function isMainAccountTokenLive(now = Date.now()): boolean { * 401 "transient" and needsReauth could never flip. */ export function isMainAccountTokenVerifiablyLive(now = Date.now()): boolean { - const tokens = readCodexTokens(); - if (!tokens?.access_token) return false; - const payload = decodeJwtPayload(tokens.access_token); + const result = readMainAuthJsonCredential(); + if (result.status !== "ok" || !result.auth.accessToken) return false; + const payload = decodeJwtPayload(result.auth.accessToken); const exp = typeof payload?.exp === "number" ? payload.exp * 1000 : undefined; return exp !== undefined && exp > now; } + +export async function forceRefreshMainAccountToken( + rejectedAccessToken?: string, + options: { signal?: AbortSignal; dependencies?: NativeMainRefreshDependencies } = {}, +): Promise<{ accessToken: string; chatgptAccountId: string } | null> { + const initial = readMainAuthJsonCredential(); + if (initial.status !== "ok") return null; + const initialFresh = tokenResultFromAuth(initial.auth, { + rejectedAccessToken, + skewMs: rejectedAccessToken ? 0 : CODEX_REFRESH_SKEW_MS, + }); + if (initialFresh && !initial.auth.refreshToken) return initialFresh; + if (!initial.auth.refreshToken) return null; + + const lockKey = refreshGrantFingerprintForToken(initial.auth.refreshToken); + const signal = combineNativeRefreshSignal(options.signal); + try { + const existingFlight = nativeMainRefreshFlights.get(lockKey); + if (existingFlight) { + const joined = await joinNativeMainRefreshFlight(existingFlight, signal); + if (joined?.accessToken === rejectedAccessToken) { + return forceRefreshMainAccountToken(rejectedAccessToken, options); + } + return joined; + } + const refreshFlight = withCodexRefreshFileLock(lockKey, signal, async () => { + const locked = readMainAuthJsonCredential(); + if (locked.status !== "ok") throw new MainAuthJsonChangedDuringRefreshError(); + if (!locked.auth.refreshToken) throw new MainAuthJsonChangedDuringRefreshError(); + if (refreshGrantFingerprintForToken(locked.auth.refreshToken) !== lockKey) { + throw new MainAuthJsonChangedDuringRefreshError(); + } + const lockedFresh = tokenResultFromAuth(locked.auth, { + rejectedAccessToken, + skewMs: rejectedAccessToken ? 0 : CODEX_REFRESH_SKEW_MS, + }); + if (lockedFresh) return lockedFresh; + + const storedCredential = freshStoredCredentialForMain(locked.auth, lockKey, rejectedAccessToken); + if (storedCredential) return persistMainStoredCredential(locked.auth, storedCredential, rejectedAccessToken); + + const refresh = options.dependencies?.refreshToken ?? refreshChatGPTToken; + const refreshed = await refresh(locked.auth.refreshToken, { signal }); + const guard = guardMainAuthJsonBeforePersist(locked.auth, { rejectedAccessToken }); + if (guard.status === "adopted") return guard.token; + const credential = refreshedCredentialFromOAuth(locked.auth, refreshed); + try { + persistMainAuthJson(guard.auth, credential); + assertMainAuthJsonCredentialPersisted(credential); + } catch (error) { + return adoptFreshMainAuthJsonAfterPersistRace(error, rejectedAccessToken); + } + publishFreshCredentialForGrant({ + refreshGrantFingerprint: lockKey, + credential, + excludeId: MAIN_CODEX_ACCOUNT_ID, + }); + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); + return tokenResultFromCredential(credential); + }).finally(() => { + if (nativeMainRefreshFlights.get(lockKey) === refreshFlight) nativeMainRefreshFlights.delete(lockKey); + }); + nativeMainRefreshFlights.set(lockKey, refreshFlight); + return await refreshFlight; + } catch (error) { + if (error instanceof MainAuthJsonChangedDuringRefreshError) throw error; + const refreshError = tokenRefreshError(error); + if (refreshError.reason !== "unknown") markAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); + throw refreshError; + } +} + +export async function getValidMainAccountToken( + options: { signal?: AbortSignal; dependencies?: NativeMainRefreshDependencies } = {}, +): Promise<{ accessToken: string; chatgptAccountId: string } | null> { + const result = readMainAuthJsonCredential(); + if (result.status !== "ok") return null; + const fresh = tokenResultFromAuth(result.auth, { skewMs: CODEX_REFRESH_SKEW_MS }); + if (fresh) return fresh; + const liveWithoutRefresh = tokenResultFromAuth(result.auth); + if (liveWithoutRefresh && !result.auth.refreshToken) return liveWithoutRefresh; + return forceRefreshMainAccountToken(result.auth.accessToken, options); +} diff --git a/src/codex/model-entitlements.ts b/src/codex/model-entitlements.ts index d06100138b..9bb3639ea6 100644 --- a/src/codex/model-entitlements.ts +++ b/src/codex/model-entitlements.ts @@ -3,7 +3,12 @@ import { readBoundedResponseBody } from "../lib/bounded-body"; import type { OcxConfig } from "../types"; import { isSelectableCodexPoolAccount } from "./account-id"; import { getValidCodexToken, readCodexAccountRecord } from "./account-store"; -import { getMainAccountToken, MAIN_CODEX_ACCOUNT_ID } from "./main-account"; +import { + getMainAccountToken, + getValidMainAccountToken, + MAIN_CODEX_ACCOUNT_ID, + type NativeMainRefreshDependencies, +} from "./main-account"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "./catalog/native-models"; const CODEX_MODELS_URL = "https://chatgpt.com/backend-api/codex/models?client_version=0.0.0"; @@ -37,7 +42,9 @@ export interface CodexModelEntitlementSnapshot { export interface CodexModelEntitlementResolveOptions { readonly fetcher?: typeof fetch; + readonly nativeMainRefreshDependencies?: NativeMainRefreshDependencies; readonly now?: number; + readonly signal?: AbortSignal; /** Test-only credential seam; production callers enumerate local main + Pool credentials. */ readonly credentials?: readonly CodexModelEntitlementCredentialSnapshot[]; } @@ -88,9 +95,15 @@ function currentCredentialIdentity(accountId: string): string | undefined { return `pool:${record.generation}:${record.credential.chatgptAccountId}`; } -async function accountCredentialSnapshot(accountId: string): Promise { +async function accountCredentialSnapshot( + accountId: string, + options: Pick, +): Promise { if (accountId === MAIN_CODEX_ACCOUNT_ID) { - const token = getMainAccountToken(); + const token = await getValidMainAccountToken({ + dependencies: options.nativeMainRefreshDependencies, + signal: options.signal, + }); return token ? { accountId, @@ -254,7 +267,7 @@ export async function resolveCodexModelEntitlements( const fetcher = options.fetcher ?? fetch; const credentials = options.credentials ? [...options.credentials] - : (await Promise.all(candidateAccountIds(config).map(accountCredentialSnapshot))) + : (await Promise.all(candidateAccountIds(config).map(accountId => accountCredentialSnapshot(accountId, options)))) .filter((value): value is CodexModelEntitlementCredentialSnapshot => value !== null); const results = await Promise.all(credentials.map(async credential => ({ credential, diff --git a/src/config/atomic-write.ts b/src/config/atomic-write.ts index b7e204580b..96afe0117e 100644 --- a/src/config/atomic-write.ts +++ b/src/config/atomic-write.ts @@ -42,6 +42,11 @@ export interface AtomicWriteIO { unlink: (path: string) => void; } +export interface AtomicWriteHooks { + beforeRename?: (tempPath: string, targetPath: string) => void; + replace?: (tempPath: string, targetPath: string, rename: () => void) => void; +} + export class AtomicWriteResidualTempError extends Error { constructor(readonly tempPath: string, readonly hardened = true, options?: ErrorOptions) { super(`Atomic config write left a ${hardened ? "hardened " : ""}zero-byte temporary file`, options); @@ -101,7 +106,7 @@ export function atomicWriteFile(path: string, content: string, io: AtomicWriteIO rename: renameAtomicFile, truncate: target => truncateSync(target, 0), unlink: unlinkSync, -}): void { +}, hooks: AtomicWriteHooks = {}): void { recordOwnedConfigPath(getConfigDir(), path); const target = resolveWriteTarget(path); assertResolvedTargetAllowed(path, target); @@ -111,7 +116,12 @@ export function atomicWriteFile(path: string, content: string, io: AtomicWriteIO io.write(tmp, content); io.harden(tmp); hardened = true; - io.rename(tmp, target); + if (hooks.replace) { + hooks.replace(tmp, target, () => io.rename(tmp, target)); + } else { + hooks.beforeRename?.(tmp, target); + io.rename(tmp, target); + } forgetEphemeralSecretPath(tmp); } catch (cause) { let scrubbed = false; diff --git a/src/lib/test-home-guard.ts b/src/lib/test-home-guard.ts index 9b8c52d0dd..493e7f8297 100644 --- a/src/lib/test-home-guard.ts +++ b/src/lib/test-home-guard.ts @@ -58,15 +58,20 @@ function canonicalize(path: string): string { * `homedir()` later would return the sandbox and leave the real home unprotected — the * guard would be perfectly inverted while its tests still looked green. */ -const PROTECTED_HOME = canonicalize( - join(process.env[REAL_HOME_ENV]?.trim() || homedir(), ".opencodex"), -); +const REAL_HOME = process.env[REAL_HOME_ENV]?.trim() || homedir(); +const PROTECTED_HOME = canonicalize(join(REAL_HOME, ".opencodex")); +const PROTECTED_CODEX_HOME = canonicalize(join(REAL_HOME, ".codex")); /** The production home this process protects. Exported for the guard's own tests. */ export function protectedHomeForTests(): string { return PROTECTED_HOME; } +/** The production Codex home this process protects when tests write native credentials. */ +export function protectedCodexHomeForTests(): string { + return PROTECTED_CODEX_HOME; +} + export function isTestHomeGuardArmed(): boolean { return process.env[GUARD_ENV] === "1"; } @@ -88,3 +93,13 @@ export function assertNotRealHomeUnderTest(dir: string): void { + "instead of calling the global writer (see devlog 260730_codex_rs_upstream_v2_live_handoff/070).", ); } + +/** Throw when an armed test process is about to write the real native Codex home. */ +export function assertNotRealCodexHomeUnderTest(dir: string): void { + if (!isTestHomeGuardArmed()) return; + if (canonicalize(dir) !== PROTECTED_CODEX_HOME) return; + throw new Error( + `refusing to write the real Codex home (${PROTECTED_CODEX_HOME}) from a test process. ` + + "Point CODEX_HOME at a temp directory for this test before writing native auth.json.", + ); +} diff --git a/src/oauth/chatgpt.ts b/src/oauth/chatgpt.ts index f4ecc7f8a9..97902dc232 100644 --- a/src/oauth/chatgpt.ts +++ b/src/oauth/chatgpt.ts @@ -2,9 +2,9 @@ import { OAuthCallbackFlow } from "./callback-server"; import type { OAuthController, OAuthCredentials } from "./types"; import { generatePKCE } from "./pkce"; -const CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"; +export const CHATGPT_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"; const AUTH_URL = "https://auth.openai.com/oauth/authorize"; -const TOKEN_URL = "https://auth.openai.com/oauth/token"; +export const CHATGPT_TOKEN_URL = "https://auth.openai.com/oauth/token"; const SCOPE = "openid profile email offline_access api.connectors.read api.connectors.invoke"; const CALLBACK_PORT = 1455; const CALLBACK_PATH = "/auth/callback"; @@ -48,7 +48,8 @@ export function extractEmail(idToken?: string, accessToken?: string): string | u function credsFromToken(data: Record): OAuthCredentials { const idToken = typeof data.id_token === "string" ? data.id_token : undefined; - const accessToken = data.access_token as string; + const accessToken = requiredTokenResponseString(data, "access_token"); + const refreshToken = optionalTokenResponseString(data, "refresh_token"); // ?? only guards null/undefined; NaN or a string expires_in would otherwise // produce a NaN expiry that never compares as expired, and a negative duration // would stamp an already-past expiry — both block refresh semantics. @@ -62,13 +63,45 @@ function credsFromToken(data: Record): OAuthCredentials { const expires = Number.isFinite(computedExpires) ? computedExpires : Date.now() + 3600 * 1000; return { access: accessToken, - refresh: (data.refresh_token as string) ?? "", + refresh: refreshToken ?? "", expires, accountId: extractAccountId(idToken, accessToken), email: extractEmail(idToken, accessToken), }; } +function requiredTokenResponseString( + data: Record, + field: "access_token", +): string { + const value = tokenResponseString(data, field); + if (value === undefined) { + throw new Error(`ChatGPT token response missing ${field}`); + } + return value; +} + +function optionalTokenResponseString( + data: Record, + field: "refresh_token", +): string | undefined { + return tokenResponseString(data, field); +} + +function tokenResponseString( + data: Record, + field: "access_token" | "refresh_token", +): string | undefined { + const value = data[field]; + if (value === undefined || value === null) { + return undefined; + } + if (typeof value !== "string" || value.trim().length === 0) { + throw new Error(`ChatGPT token response contains invalid ${field}`); + } + return value; +} + export class ChatGPTOAuthFlow extends OAuthCallbackFlow { #verifier = ""; forceLogin = false; @@ -88,7 +121,7 @@ export class ChatGPTOAuthFlow extends OAuthCallbackFlow { this.#verifier = pkce.verifier; const params = new URLSearchParams({ response_type: "code", - client_id: CLIENT_ID, + client_id: CHATGPT_CLIENT_ID, redirect_uri: redirectUri, scope: SCOPE, code_challenge: pkce.challenge, @@ -107,12 +140,12 @@ export class ChatGPTOAuthFlow extends OAuthCallbackFlow { async exchangeToken(code: string, _state: string, redirectUri: string): Promise { if (!this.#verifier) throw new Error("ChatGPT PKCE verifier not initialized"); - const resp = await fetch(TOKEN_URL, { + const resp = await fetch(CHATGPT_TOKEN_URL, { method: "POST", headers: { "Content-Type": "application/x-www-form-urlencoded" }, body: new URLSearchParams({ grant_type: "authorization_code", - client_id: CLIENT_ID, + client_id: CHATGPT_CLIENT_ID, code, redirect_uri: redirectUri, code_verifier: this.#verifier, @@ -135,6 +168,40 @@ function safeErrorDescription(resp: Response): Promise { }); } +interface OAuthErrorDetails { + description: string; + error?: string; + errorDescription?: string; +} + +async function oauthErrorDetails(resp: Response): Promise { + const text = await resp.text().catch(() => ""); + try { + const parsed = JSON.parse(text) as { error?: unknown; error_description?: unknown }; + const error = typeof parsed.error === "string" ? parsed.error : undefined; + const errorDescription = typeof parsed.error_description === "string" ? parsed.error_description : undefined; + return { + description: [error, errorDescription].filter(Boolean).join(": ") || `HTTP ${resp.status}`, + ...(error ? { error } : {}), + ...(errorDescription ? { errorDescription } : {}), + }; + } catch { + return { description: `HTTP ${resp.status}` }; + } +} + +export class ChatGPTTokenRefreshError extends Error { + constructor( + readonly status: number, + readonly oauthError: string | undefined, + readonly oauthErrorDescription: string | undefined, + message: string, + ) { + super(message); + this.name = "ChatGPTTokenRefreshError"; + } +} + export async function loginChatGPT(ctrl: OAuthController, opts?: { forceLogin?: boolean }): Promise { const flow = new ChatGPTOAuthFlow(ctrl); if (opts?.forceLogin) flow.forceLogin = true; @@ -143,19 +210,28 @@ export async function loginChatGPT(ctrl: OAuthController, opts?: { forceLogin?: // Note: uses form-urlencoded per OAuth 2.0 spec (RFC 6749 §6). // Codex-rs uses JSON for refresh — intentional divergence; both accepted by auth.openai.com. -export async function refreshChatGPTToken(refreshToken: string): Promise { - const resp = await fetch(TOKEN_URL, { +export async function refreshChatGPTToken( + refreshToken: string, + options: { signal?: AbortSignal } = {}, +): Promise { + const resp = await fetch(CHATGPT_TOKEN_URL, { method: "POST", headers: { "Content-Type": "application/x-www-form-urlencoded" }, body: new URLSearchParams({ grant_type: "refresh_token", - client_id: CLIENT_ID, + client_id: CHATGPT_CLIENT_ID, refresh_token: refreshToken, }).toString(), + signal: options.signal, }); if (!resp.ok) { - const errDesc = await safeErrorDescription(resp); - throw new Error(`ChatGPT refresh failed: ${resp.status} ${errDesc}`); + const details = await oauthErrorDetails(resp); + throw new ChatGPTTokenRefreshError( + resp.status, + details.error, + details.errorDescription, + `ChatGPT refresh failed: ${resp.status} ${details.description}`, + ); } return credsFromToken((await resp.json()) as Record); } diff --git a/src/routing/analytics.ts b/src/routing/analytics.ts index 9e7b5a0696..422770e323 100644 --- a/src/routing/analytics.ts +++ b/src/routing/analytics.ts @@ -118,6 +118,7 @@ const COOLDOWN_RECOVERY_KINDS = new Set([ "rate-limit-429", "key-429", "oauth-401", + "codex-main-401", "anthropic-oauth-429", ]); diff --git a/src/server/responses/codex-auth-error.ts b/src/server/responses/codex-auth-error.ts index cb7ce8609c..58c8c125b6 100644 --- a/src/server/responses/codex-auth-error.ts +++ b/src/server/responses/codex-auth-error.ts @@ -1,4 +1,10 @@ import { formatErrorResponse } from "../../bridge"; +import { + CodexCredentialRefreshBusyError, + CodexCredentialRefreshLockTimeoutError, + CodexCredentialRefreshStaleError, + TokenRefreshError, +} from "../../codex/account-store"; import { CodexAccountCooldownError, codexMainProfileDrainingResponse, @@ -10,12 +16,44 @@ import { CodexPoolAuthenticationError, CodexThreadAffinityExpiredError, } from "../../codex/auth-context"; +import { MAIN_CODEX_ACCOUNT_ID, MainAuthJsonChangedDuringRefreshError } from "../../codex/main-account"; export interface CodexAuthContextErrorResponseOptions { accountSelector?: string; now: number; } +export function nativeMainRefreshFailureResponse(error: unknown): Response { + if (error instanceof TokenRefreshError && error.reason !== "unknown") { + return formatErrorResponse( + 401, + "authentication_error", + "Codex main account needs reauthentication", + ); + } + if ( + error instanceof TokenRefreshError + || error instanceof MainAuthJsonChangedDuringRefreshError + || error instanceof CodexCredentialRefreshBusyError + || error instanceof CodexCredentialRefreshLockTimeoutError + || error instanceof CodexCredentialRefreshStaleError + ) { + const response = formatErrorResponse( + 503, + "server_busy", + "Codex main credential refresh did not complete; retry this request", + ); + const headers = new Headers(response.headers); + headers.set("Retry-After", "1"); + return new Response(response.body, { status: response.status, headers }); + } + return formatErrorResponse( + 401, + "authentication_error", + "No usable Codex main credential to serve this request", + ); +} + /** Shared HTTP contract for Codex auth-context failures on Responses surfaces. */ export function mapCodexAuthContextErrorToResponse( error: unknown, @@ -35,6 +73,9 @@ export function mapCodexAuthContextErrorToResponse( ); } if (error instanceof CodexAuthContextError) { + if (error.accountId === MAIN_CODEX_ACCOUNT_ID) { + return nativeMainRefreshFailureResponse(error.cause); + } return formatErrorResponse( 401, "authentication_error", @@ -51,5 +92,8 @@ export function mapCodexAuthContextErrorToResponse( "No usable Codex main credential to serve this request", ); } + if (error instanceof TokenRefreshError || error instanceof MainAuthJsonChangedDuringRefreshError) { + return nativeMainRefreshFailureResponse(error); + } return undefined; } diff --git a/src/server/responses/compact.ts b/src/server/responses/compact.ts index 4416173312..0bf5add819 100644 --- a/src/server/responses/compact.ts +++ b/src/server/responses/compact.ts @@ -44,14 +44,16 @@ import { applyCodexAuthContextToProvider, CodexMainProfileDrainingError, headersForCodexAuthContext, - materializeCodexUpstreamAuth, + materializeCodexUpstreamAuthAsync, isCodexAuthContextUsable, resolveCodexAuthContext, codexProbeLeaseId, codexProbeQuotaScope, releaseCodexAuthContextProbeLease, + stripCodexRuntimeProviderFields, type CodexAuthContext, } from "../../codex/auth-context"; +import { forceRefreshMainAccountToken, type NativeMainRefreshDependencies } from "../../codex/main-account"; import { formatCodexProviderForLog, recordCodexUpstreamOutcome, @@ -126,10 +128,14 @@ import { usesCodexForwardPoolAuth, } from "./core"; import { fetchWithHeaderTimeout, providerFetch, safeHostLabel, safeOriginLabel } from "./fetch-helpers"; -import { mapCodexAuthContextErrorToResponse } from "./codex-auth-error"; +import { mapCodexAuthContextErrorToResponse, nativeMainRefreshFailureResponse } from "./codex-auth-error"; export const COMPACT_RESPONSE_MAX_BYTES = 32 * 1024 * 1024; +export interface HandleResponsesCompactOptions { + nativeMainRefreshDependencies?: NativeMainRefreshDependencies; +} + export function compactResponseTooLargeError(): Response { return new Response(JSON.stringify({ error: { @@ -158,14 +164,17 @@ async function resolveAlternateCompactContext(args: { selectedModelId: string | undefined; excludeAccountId: string | null; turnAdmissionLease?: AdmissionLease; + nativeMainRefreshDependencies?: NativeMainRefreshDependencies; }): Promise<{ authCtx: CodexAuthContext; provider: OcxProviderConfig; headers: Headers } | null> { - const { req, config, route, selectedModelId, excludeAccountId, turnAdmissionLease } = args; + const { req, config, route, selectedModelId, excludeAccountId, turnAdmissionLease, nativeMainRefreshDependencies } = args; if (!route.codexAccountMode || !excludeAccountId) return null; try { const authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode, { ...(selectedModelId ? { modelId: selectedModelId } : {}), excludeAccountId, beginCodexAccountSelection: codexAccountSelectionForTurn(turnAdmissionLease), + signal: req.signal, + nativeMainRefreshDependencies, }); if (!authCtx.accountId || authCtx.accountId === excludeAccountId) return null; const provider = applyCodexAuthContextToProvider(route.provider, authCtx, route.codexAccountMode); @@ -195,6 +204,63 @@ async function resolveAlternateCompactContext(args: { } } +async function refreshNativeMainCompactContext(args: { + req: Request; + authCtx: CodexAuthContext; + provider: OcxProviderConfig; + codexAccountMode?: CodexAccountMode; + substituteMainCredential: boolean; + nativeMainRefreshDependencies?: NativeMainRefreshDependencies; +}): Promise< + | { ok: true; authCtx: CodexAuthContext; provider: OcxProviderConfig; headers: Headers } + | { ok: false; response: Response } +> { + const { req, authCtx, provider, codexAccountMode, substituteMainCredential, nativeMainRefreshDependencies } = args; + if (authCtx.kind !== "main-pool") { + return { ok: false, response: formatErrorResponse(401, "authentication_error", "No native main credential to refresh") }; + } + try { + const refreshed = await forceRefreshMainAccountToken(authCtx.accessToken, { + signal: req.signal, + dependencies: nativeMainRefreshDependencies, + }); + if (!refreshed) { + return { + ok: false, + response: formatErrorResponse(401, "authentication_error", "Codex main account needs reauthentication"), + }; + } + const refreshedAuthCtx: CodexAuthContext = { + ...authCtx, + accessToken: refreshed.accessToken, + chatgptAccountId: refreshed.chatgptAccountId, + }; + const refreshedProvider = applyCodexAuthContextToProvider( + stripCodexRuntimeProviderFields(provider), + refreshedAuthCtx, + codexAccountMode, + ); + const headers = new Headers({ "content-type": "application/json" }); + const selected = await materializeCodexUpstreamAuthAsync(req.headers, refreshedAuthCtx, { + substituteMainCredential, + signal: req.signal, + nativeMainRefreshDependencies, + }); + for (const name of FORWARD_HEADERS) { + const value = selected.get(name); + if (value) headers.set(name, value); + } + const override = (refreshedProvider as { _codexAccountOverride?: { accessToken: string; chatgptAccountId: string } })._codexAccountOverride; + if (override) { + headers.set("authorization", `Bearer ${override.accessToken}`); + headers.set("chatgpt-account-id", override.chatgptAccountId); + } + return { ok: true, authCtx: refreshedAuthCtx, provider: refreshedProvider, headers }; + } catch (error) { + return { ok: false, response: nativeMainRefreshFailureResponse(error) }; + } +} + /** * Headers a client needs to back off correctly after a pool rejection. The buffered * response is rebuilt from scratch, so anything not listed here is dropped — which is @@ -267,6 +333,7 @@ export async function handleResponsesCompact( logCtx: RequestLogContext, turnAdmissionLease?: AdmissionLease, admission?: DataPlaneAdmission, + options: HandleResponsesCompactOptions = {}, ): Promise { let body: unknown; try { @@ -365,7 +432,7 @@ export async function handleResponsesCompact( // headers would run compaction on the wrong account (or 401) whenever a pool account is // active for this thread while normal turns succeed. let compactProvider = route.provider; - const headers = new Headers({ "content-type": "application/json" }); + let headers = new Headers({ "content-type": "application/json" }); try { if (route.codexAccountMode) { authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode, { @@ -373,9 +440,15 @@ export async function handleResponsesCompact( modelId: selectedModelId, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(turnAdmissionLease), + signal: req.signal, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, }); logCtx.accountLogLabel = codexAuthContextLogLabel(authCtx, config); - const selected = materializeCodexUpstreamAuth(req.headers, authCtx, { substituteMainCredential }); + const selected = await materializeCodexUpstreamAuthAsync(req.headers, authCtx, { + substituteMainCredential, + signal: req.signal, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, + }); compactProvider = applyCodexAuthContextToProvider(route.provider, authCtx, route.codexAccountMode); for (const name of FORWARD_HEADERS) { const value = selected.get(name); @@ -514,6 +587,7 @@ export async function handleResponsesCompact( // actually happens, so every recorder call names the context that produced it. let outcomeCtx = authCtx; let upstream: Response; + let codexMain401ReplayAttempted = false; try { // Same connect timeout + keep-alive reset + transient-5xx recovery as /v1/responses — // compact hits the same ChatGPT host and must soft-avoid / clear affinity (#186). @@ -543,6 +617,59 @@ export async function handleResponsesCompact( return formatErrorResponse(502, "upstream_error", "Failed to connect to compact upstream"); } + if ( + upstream.status === 401 + && authCtx.kind === "main-pool" + && usesCodexForwardPoolAuth(authCtx, compactProvider) + && !codexMain401ReplayAttempted + && !req.signal.aborted + ) { + codexMain401ReplayAttempted = true; + await upstream.body?.cancel().catch(() => undefined); + const replay = await refreshNativeMainCompactContext({ + req, + authCtx, + provider: compactProvider, + codexAccountMode: route.codexAccountMode, + substituteMainCredential, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, + }); + if (!replay.ok) { + recordCompactPoolOutcome(outcomeCtx, replay.response.status === 401 ? 401 : "connect_neutral"); + return replay.response; + } + authCtx = replay.authCtx; + outcomeCtx = replay.authCtx; + compactProvider = replay.provider; + headers = replay.headers; + logCtx.accountLogLabel = codexAuthContextLogLabel(replay.authCtx, config); + try { + upstream = await sendCompactAttempt(compactProvider, headers, "single"); + } catch (err) { + if (req.signal.aborted) { + recordCompactPoolOutcome(outcomeCtx, 499); + return formatErrorResponse(499, "client_cancelled", "Client cancelled compact request"); + } + const outcome = classifyTransportFailureKind(err); + if (outcome === "connect_neutral") { + if (compactHostCircuitEnabled) { + recordUpstreamHostFailure(actualCompactHostKey, { + code: transportErrorCode(err), + threshold: config.upstreamHostCircuitThreshold, + lease: compactHostAdmissionLease, + }); + } else { + recordUpstreamHostFailure(actualCompactHostKey, { code: transportErrorCode(err) }); + } + } else { + releaseUpstreamHostAdmission(compactHostAdmissionLease); + } + compactHostAdmissionLease = null; + recordCompactPoolOutcome(outcomeCtx, outcome); + return formatErrorResponse(502, "upstream_error", "Failed to connect to compact upstream"); + } + } + // Bounded same-request alternate: the regular /v1/responses path already does this // (core.ts:319-423) and recognizes exactly 429/402. Without it a pool rejection // surfaces to the client, which retries the compact task OUTSIDE the logical request @@ -569,6 +696,7 @@ export async function handleResponsesCompact( selectedModelId, excludeAccountId: authCtx.accountId, turnAdmissionLease, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, }); // Resolution can await a credential refresh, so the client may have gone away // while we were choosing B. Re-check before spending anything: recording A, @@ -676,7 +804,12 @@ export async function handleResponsesCompact( headers: internalHeaders, body: JSON.stringify(internalBody), }); - const response = await handleResponses(internalReq, config, logCtx, { abortSignal: req.signal, turnAdmissionLease, ...(admission ? { admission } : {}) }); + const response = await handleResponses(internalReq, config, logCtx, { + abortSignal: req.signal, + turnAdmissionLease, + ...(admission ? { admission } : {}), + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, + }); if (!response.ok) return response; let json: { output?: unknown[]; status?: unknown; error?: unknown }; if (response.headers.get("content-type")?.includes("text/event-stream")) { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 67005be579..c726c9dc39 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -120,15 +120,17 @@ import { CodexPoolAuthenticationError, CodexThreadAffinityExpiredError, headersForCodexAuthContext, - materializeCodexUpstreamAuth, + materializeCodexUpstreamAuthAsync, isCodexAuthContextUsable, resolveCodexAuthContext, + shouldMarkAccountNeedsReauthForCodexAuthFailure, codexProbeLeaseId, codexProbeQuotaScope, releaseCodexAuthContextProbeLease, stripCodexRuntimeProviderFields, type CodexAuthContext, } from "../../codex/auth-context"; +import { forceRefreshMainAccountToken, type NativeMainRefreshDependencies } from "../../codex/main-account"; import { entitledCodexAccountIdsForModel, invalidateCodexModelEntitlementsForAccount, @@ -270,7 +272,7 @@ import { createResponsesModelPayloadRewrite, rewriteResponsesModelJson } from ". import type { EffectiveSubagentRoster, SpawnAgentSurface } from "../../codex/catalog"; import { buildToolBridgeMaps, collabSurface, injectDeveloperMessage, multiAgentGuidanceText } from "./collaboration"; -import { mapCodexAuthContextErrorToResponse } from "./codex-auth-error"; +import { mapCodexAuthContextErrorToResponse, nativeMainRefreshFailureResponse } from "./codex-auth-error"; import { hasUnreadableEncryptedAgentTask, looksLikeBackendCiphertext, sanitizeEncryptedContentInPlace } from "./encrypted-payload"; import { fetchWithHeaderTimeout, providerFetch, safeHostLabel, safeOriginLabel } from "./fetch-helpers"; import { classifyTransportFailureKind, transportErrorCode } from "../../lib/upstream-reachability"; @@ -1235,6 +1237,8 @@ export interface HandleResponsesOptions { responsesTerminalRepairScheduler?: ResponsesTerminalRepairScheduler; /** Internal deterministic runtime-identity seam for Codex upstream WS selection tests. */ codexWsRuntimeIdentity?: BunRuntimeGateInput; + /** Test seam for native main auth.json refresh without calling the real OAuth endpoint. */ + nativeMainRefreshDependencies?: NativeMainRefreshDependencies; /** * When true, body `prompt_cache_key` is a Claude Desktop shared cache cohort * (system/tools hash), not a per-session id — do not use it for Anthropic pool affinity. @@ -1518,6 +1522,8 @@ async function resolveResponsesCodexAuth( modelId: route.modelId, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + signal: options.abortSignal, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, }); options.onCodexAuthContextResolved?.(authCtx); } else { @@ -1534,7 +1540,11 @@ async function resolveResponsesCodexAuth( return { ok: true, authCtx, - headers: materializeCodexUpstreamAuth(req.headers, authCtx, { substituteMainCredential }), + headers: await materializeCodexUpstreamAuthAsync(req.headers, authCtx, { + substituteMainCredential, + signal: options.abortSignal, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, + }), substituteMainCredential, }; } catch (err) { @@ -1542,7 +1552,11 @@ async function resolveResponsesCodexAuth( const safeAccountLabel = route.codexAccountNamespace ? `${route.providerName}-${route.codexAccountNamespace}` : formatCodexProviderForLog(route.providerName, err.accountId, config); - console.error(`[codex-auth] Pool account ${safeAccountLabel} token failed; reauthentication required`); + const cause = (err as Error & { cause?: unknown }).cause; + const disposition = shouldMarkAccountNeedsReauthForCodexAuthFailure(cause) + ? "reauthentication required" + : "retryable refresh failure"; + console.error(`[codex-auth] Pool account ${safeAccountLabel} token failed; ${disposition}`); } if (err instanceof ForwardAdmissionCredentialError) { return { ok: false, response: formatErrorResponse(401, "authentication_error", err.message) }; @@ -1556,6 +1570,52 @@ async function resolveResponsesCodexAuth( } } +async function refreshNativeMainForwardAuth(args: { + req: Request; + route: RouteResult; + authCtx: CodexAuthContext; + substituteMainCredential: boolean; + options: HandleResponsesOptions; +}): Promise< + | { ok: true; authCtx: CodexAuthContext; provider: OcxProviderConfig; headers: Headers } + | { ok: false; response: Response } +> { + const { req, route, authCtx, substituteMainCredential, options } = args; + if (authCtx.kind !== "main-pool") { + return { ok: false, response: formatErrorResponse(401, "authentication_error", "No native main credential to refresh") }; + } + try { + const refreshed = await forceRefreshMainAccountToken(authCtx.accessToken, { + signal: options.abortSignal, + dependencies: options.nativeMainRefreshDependencies, + }); + if (!refreshed) { + return { + ok: false, + response: formatErrorResponse(401, "authentication_error", "Codex main account needs reauthentication"), + }; + } + const refreshedAuthCtx: CodexAuthContext = { + ...authCtx, + accessToken: refreshed.accessToken, + chatgptAccountId: refreshed.chatgptAccountId, + }; + const provider = applyCodexAuthContextToProvider( + stripCodexRuntimeProviderFields(route.provider), + refreshedAuthCtx, + route.codexAccountMode, + ); + const headers = await materializeCodexUpstreamAuthAsync(req.headers, refreshedAuthCtx, { + substituteMainCredential, + signal: options.abortSignal, + nativeMainRefreshDependencies: options.nativeMainRefreshDependencies, + }); + return { ok: true, authCtx: refreshedAuthCtx, provider, headers }; + } catch (error) { + return { ok: false, response: nativeMainRefreshFailureResponse(error) }; + } +} + /** * Apply every route-dependent request mutation against the final selected route. * Must run only after subagent fallback has settled the model/provider. @@ -3210,6 +3270,7 @@ async function handleResponsesInner( const opaqueBlobRecoveryGuard: OpaqueBlobRecoveryGuard = { attempted: false }; let oauth401ReplayAttempted = false; + let codexMain401ReplayAttempted = false; const rateLimitPolicy = rateLimitRetryPolicyFor(route.provider); let rateLimitRetries = 0; const rebuildAndRefetch = async ( @@ -3284,6 +3345,100 @@ async function handleResponsesInner( // their OAuth contract identical: one pre-stream 401 forces a credential refresh and one // rebuilt replay. xAI's current subscription models use this branch now that their official // Grok CLI catalog declares the Responses backend. + if ( + upstreamResponse.status === 401 + && authCtx.kind === "main-pool" + && usesCodexForwardPoolAuth(authCtx, route.provider) + && !codexMain401ReplayAttempted + ) { + codexMain401ReplayAttempted = true; + try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } + const replay = await refreshNativeMainForwardAuth({ + req, + route, + authCtx, + substituteMainCredential, + options, + }); + if (!replay.ok) { + upstream.abort(); + releaseCodexAuthContextProbeLease(authCtx); + return replay.response; + } + authCtx = replay.authCtx; + route.provider = replay.provider; + selectedForwardHeaders = replay.headers; + const replayAdapter = resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, replay.provider, inboundWire), + config.cacheRetention, + ); + if (!("passthrough" in replayAdapter) || !replayAdapter.passthrough) { + upstream.abort(); + return formatErrorResponse(502, "upstream_error", "Native main refresh changed the provider wire unexpectedly"); + } + bindRouteReasoningReplayScope({ + parsed, + providerName: route.providerName, + provider: replay.provider, + adapterName: replayAdapter.name, + codexAuthContext: authCtx, + forwardHeaders: selectedForwardHeaders, + }); + logCtx.providerAdapter = replayAdapter.name; + sealRequestAttemptIdentity( + logCtx.activeAttempt, + logCtx.provider, + replayAdapter.name, + logCtx.accountLogLabel, + ); + try { + request = await replayAdapter.buildRequest(parsed, { + headers: selectedForwardHeaders, + translatorBudget, + }); + refreshRoutedNamespaceToolAliases(request); + recordAdapterReasoning(logCtx, request); + recordAdapterTier(logCtx, request); + } catch (err) { + upstream.abort(); + if (options.abortSignal?.aborted) return clientCancelledResponse(); + const msg = err instanceof Error ? err.message : String(err); + return formatErrorResponse(400, "invalid_request_error", redactSecretString(msg)); + } + passthroughEstimate = typeof request.usageLog?.inputTokens === "number" + ? request.usageLog.inputTokens + : undefined; + if (passthroughEstimate !== undefined) logCtx.usageLogInputTokens = passthroughEstimate; + refreshUndeclaredToolGuard(request); + try { + upstreamResponse = await fetchWithTransientRetry( + recovery => { + noteAttemptSend(logCtx.activeAttempt, passthroughEstimate, recovery ?? "codex-main-401"); + return fetchWithHeaderTimeout(request.url, applyUpstreamRecoveryInit({ + method: request.method, + headers: request.headers, + body: request.body, + }, recovery), upstream.signal, connectMs, parsed.stream, + providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, + modelId: route.modelId, + }), + route.provider.authMode === "forward") + .then(res => { + settleObservedHostResponse(); + return res; + }); + }, + { abortSignal: upstream.signal, label: safeHostLabel(request.url) }, + ); + } catch (err) { + return transportFailureResponse(err); + } finally { + request.releaseBodyObservation?.(); + } + continue passthroughRecovery; + } + if ( upstreamResponse.status === 401 && isOAuth401ReplayProvider @@ -4633,6 +4788,7 @@ async function handleResponsesInner( let imageRetryAttempted = false; const opaqueBlobRecoveryGuard: OpaqueBlobRecoveryGuard = { attempted: false }; let oauth401ReplayAttempted = false; + let codexMain401ReplayAttempted = false; /** * Rebuild the request from the current parsed input (and any image-tier bias) and refetch * it once, tagging the attempt with the given recovery kind. Rebuilds are deterministic @@ -4712,6 +4868,51 @@ async function handleResponsesInner( }; // Keep recovery kinds in sync with the native Responses `passthroughRecovery:` loop above. recovery: for (;;) { + if ( + upstreamResponse.status === 401 + && authCtx.kind === "main-pool" + && usesCodexForwardPoolAuth(authCtx, route.provider) + && !codexMain401ReplayAttempted + ) { + codexMain401ReplayAttempted = true; + try { void upstreamResponse.body?.cancel().catch(() => {}); } catch { /* already consumed/closed */ } + const replay = await refreshNativeMainForwardAuth({ + req, + route, + authCtx, + substituteMainCredential, + options, + }); + if (!replay.ok) { + cleanupUpstreamAbort(); + upstream.abort(); + releaseCodexAuthContextProbeLease(authCtx); + return replay.response; + } + authCtx = replay.authCtx; + route.provider = replay.provider; + selectedForwardHeaders = replay.headers; + invalidateSameTargetRequest(); + activeAdapter = resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, replay.provider, inboundWire), + config.cacheRetention, + ); + bindRouteReasoningReplayScope({ + parsed, + providerName: route.providerName, + provider: replay.provider, + adapterName: activeAdapter.name, + codexAuthContext: authCtx, + forwardHeaders: selectedForwardHeaders, + }); + logCtx.providerAdapter = activeAdapter.name; + sealRequestAttemptIdentity(logCtx.activeAttempt, logCtx.provider, activeAdapter.name, logCtx.accountLogLabel); + const result = await rebuildAndRefetch("codex-main-401"); + if ("failed" in result) return result.failed; + upstreamResponse = result; + continue recovery; + } + if ( upstreamResponse.status === 401 && isOAuth401ReplayProvider diff --git a/src/usage/log.ts b/src/usage/log.ts index 7bca650964..ddabb2caeb 100644 --- a/src/usage/log.ts +++ b/src/usage/log.ts @@ -25,6 +25,7 @@ export type AttemptRecoveryKind = | "transient-5xx" | "connection-reset" | "oauth-401" + | "codex-main-401" | "key-429" | "rate-limit-429" | "anthropic-oauth-429" @@ -215,6 +216,7 @@ const ATTEMPT_RECOVERY_KINDS = new Set([ "transient-5xx", "connection-reset", "oauth-401", + "codex-main-401", "key-429", "rate-limit-429", "anthropic-oauth-429", diff --git a/tests/chatgpt-oauth.test.ts b/tests/chatgpt-oauth.test.ts index 161baca024..ab485f5f46 100644 --- a/tests/chatgpt-oauth.test.ts +++ b/tests/chatgpt-oauth.test.ts @@ -109,13 +109,13 @@ describe("ChatGPT OAuth constants", () => { describe("codex-account-store constants sync", () => { test("uses same auth.openai.com endpoint as chatgpt.ts", async () => { const source = await Bun.file("src/codex/account-store.ts").text(); - expect(source).toContain("auth.openai.com/oauth/token"); + expect(source).toContain("CHATGPT_TOKEN_URL"); expect(source).not.toContain("auth0.openai.com"); }); test("uses same client_id as chatgpt.ts", async () => { const source = await Bun.file("src/codex/account-store.ts").text(); - expect(source).toContain("app_EMoamEEZ73f0CkXaXp7hrann"); + expect(source).toContain("CHATGPT_CLIENT_ID"); expect(source).not.toContain("DRivsnm2Mu42T3KOpqdtwB3NYviHYzwD"); }); diff --git a/tests/chatgpt-token-expiry.test.ts b/tests/chatgpt-token-expiry.test.ts index 7b0460e1f9..6948d863e8 100644 --- a/tests/chatgpt-token-expiry.test.ts +++ b/tests/chatgpt-token-expiry.test.ts @@ -63,4 +63,22 @@ describe("ChatGPT OAuth token response parsing", () => { expect(cred.expires).toBeGreaterThan(before); expect(Math.abs(cred.expires - (before + FALLBACK_MS))).toBeLessThan(TOLERANCE_MS); }); + + test("refresh rejects a 200 token response without a usable access token", async () => { + globalThis.fetch = (async () => new Response( + JSON.stringify({ refresh_token: "rt", expires_in: 3600 }), + { status: 200 }, + )) as typeof fetch; + + await expect(refreshChatGPTToken("secret")).rejects.toThrow("access_token"); + }); + + test("refresh rejects a 200 token response with a malformed refresh token", async () => { + globalThis.fetch = (async () => new Response( + JSON.stringify({ access_token: "at", refresh_token: 7, expires_in: 3600 }), + { status: 200 }, + )) as typeof fetch; + + await expect(refreshChatGPTToken("secret")).rejects.toThrow("refresh_token"); + }); }); diff --git a/tests/codex-account-store.test.ts b/tests/codex-account-store.test.ts index 63f19cf0b3..1f2cf8dcd0 100644 --- a/tests/codex-account-store.test.ts +++ b/tests/codex-account-store.test.ts @@ -261,6 +261,39 @@ describe("codex-account-store CRUD", () => { } }); + test("refresh rejects a 200 token response without a usable access token", async () => { + const { + getCodexAccountCredential, + getValidCodexToken, + readCodexAccountRecord, + saveCodexAccountCredential, + TokenRefreshError, + } = await import("../src/codex/account-store"); + saveCodexAccountCredential("refresh-malformed-access", { + accessToken: "old", + refreshToken: "old-r", + expiresAt: 0, + chatgptAccountId: "acc", + }); + const startGeneration = readCodexAccountRecord("refresh-malformed-access")!.generation; + const originalFetch = globalThis.fetch; + globalThis.fetch = (async () => new Response(JSON.stringify({ + refresh_token: "new-r", + expires_in: 3600, + }), { status: 200 })) as typeof fetch; + + try { + await expect(getValidCodexToken("refresh-malformed-access")).rejects.toBeInstanceOf(TokenRefreshError); + expect(getCodexAccountCredential("refresh-malformed-access")).toMatchObject({ + accessToken: "old", + refreshToken: "old-r", + }); + expect(readCodexAccountRecord("refresh-malformed-access")!.generation).toBe(startGeneration); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("refresh with a non-finite expires_in falls back to the 3600s default", async () => { const { getCodexAccountCredential, diff --git a/tests/codex-main-account-refresh.test.ts b/tests/codex-main-account-refresh.test.ts new file mode 100644 index 0000000000..c76aa9799f --- /dev/null +++ b/tests/codex-main-account-refresh.test.ts @@ -0,0 +1,784 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { createHash } from "node:crypto"; +import { mkdtempSync, readFileSync, rmSync, unlinkSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { isCodexAccountUsable } from "../src/codex/account-usability"; +import { clearAccountNeedsReauth, isAccountNeedsReauth } from "../src/codex/account-runtime-state"; +import { + TokenRefreshError, + getCodexAccountCredential, + readCodexAccountRecord, + refreshGrantFingerprintForToken, + saveCodexAccountCredential, +} from "../src/codex/account-store"; +import { NATIVE_DAYBREAK_BLUE_MODEL } from "../src/codex/catalog/native-models"; +import { + MAIN_CODEX_ACCOUNT_ID, + MainAuthJsonChangedDuringRefreshError, + forceRefreshMainAccountToken, + getValidMainAccountToken, + setMainAuthJsonPublishHookForTests, + setMainAuthJsonRenameHookForTests, + setMainAuthJsonReplaceHookForTests, +} from "../src/codex/main-account"; +import { + entitledCodexAccountIdsForModel, + resolveCodexModelEntitlements, +} from "../src/codex/model-entitlements"; +import { + CodexAuthContextError, + materializeCodexUpstreamAuthAsync, + resolveCodexAuthContext, + type CodexAuthContext, +} from "../src/codex/auth-context"; +import type { CodexModelEntitlementResolveOptions } from "../src/codex/model-entitlements"; +import { ChatGPTTokenRefreshError } from "../src/oauth/chatgpt"; +import type { OAuthCredentials } from "../src/oauth/types"; +import type { OcxConfig } from "../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; + +let isolatedCodexHome: IsolatedCodexHome; +let opencodexHome: string; +let previousOpencodexHome: string | undefined; + +beforeEach(() => { + isolatedCodexHome = installIsolatedCodexHome("ocx-main-refresh-codex-"); + previousOpencodexHome = process.env.OPENCODEX_HOME; + opencodexHome = mkdtempSync(join(tmpdir(), "ocx-main-refresh-store-")); + process.env.OPENCODEX_HOME = opencodexHome; + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); +}); + +afterEach(() => { + setMainAuthJsonPublishHookForTests(null); + setMainAuthJsonRenameHookForTests(null); + setMainAuthJsonReplaceHookForTests(null); + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); + isolatedCodexHome.restore(); + rmSync(opencodexHome, { recursive: true, force: true }); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; +}); + +function jwt(name: string, expiresInSeconds: number, accountId = "main-account"): string { + const payload = Buffer.from(JSON.stringify({ + exp: Math.floor(Date.now() / 1000) + expiresInSeconds, + chatgpt_account_id: accountId, + marker: name, + })).toString("base64url"); + return `header.${payload}.signature`; +} + +function authPath(): string { + return join(isolatedCodexHome.path, "auth.json"); +} + +function writeAuth(tokens: Record, extra: Record = {}): void { + writeFileSync(authPath(), JSON.stringify({ + ...extra, + tokens, + }, null, 2)); +} + +function readAuth(): { tokens: Record; untouched?: unknown } { + return JSON.parse(readFileSync(authPath(), "utf-8")) as { tokens: Record; untouched?: unknown }; +} + +function refreshed(accessToken: string, refreshToken = "fresh-refresh", accountId = "fresh-account"): OAuthCredentials { + return { + access: accessToken, + refresh: refreshToken, + expires: Date.now() + 3_600_000, + accountId, + }; +} + +function nativeRefreshLockPath(refreshToken: string): string { + const lockKey = refreshGrantFingerprintForToken(refreshToken); + const digest = createHash("sha256").update(lockKey).digest("hex").slice(0, 32); + return join(opencodexHome, `codex-refresh-${digest}.lock`); +} + +function writeNativeRefreshLock(refreshToken: string): string { + const path = nativeRefreshLockPath(refreshToken); + writeFileSync(path, JSON.stringify({ acquiredAt: Date.now(), pid: process.pid }) + "\n", { mode: 0o600 }); + return path; +} + +function mainOnlyConfig(): OcxConfig { + return { + providers: {}, + codexAccounts: [], + activeCodexAccountId: MAIN_CODEX_ACCOUNT_ID, + autoSwitchThreshold: 0, + } as OcxConfig; +} + +describe("native main auth.json refresh", () => { + test("keeps a refresh-only auth.json selectable and refreshes before auth context materialization", async () => { + const freshAccess = jwt("fresh", 3_600); + writeAuth({ refresh_token: "refresh-only", account_id: "main-account" }, { untouched: true }); + + expect(isCodexAccountUsable(mainOnlyConfig(), MAIN_CODEX_ACCOUNT_ID)).toBe(true); + const ctx = await resolveCodexAuthContext(new Headers(), mainOnlyConfig(), "pool", { + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess), + }, + }); + + expect(ctx).toMatchObject({ + kind: "main-pool", + accountId: MAIN_CODEX_ACCOUNT_ID, + accessToken: freshAccess, + chatgptAccountId: "fresh-account", + }); + const persisted = readAuth(); + expect(persisted.untouched).toBe(true); + expect(persisted.tokens.access_token).toBe(freshAccess); + expect(persisted.tokens.refresh_token).toBe("fresh-refresh"); + expect(persisted.tokens.account_id).toBe("fresh-account"); + }); + + test("adopts an external auth.json writer instead of clobbering it after the refresh round trip", async () => { + const expiredAccess = jwt("expired", -3_600); + const externalAccess = jwt("external", 3_600, "external-account"); + const fetchedAccess = jwt("fetched", 3_600, "fetched-account"); + writeAuth({ + access_token: expiredAccess, + refresh_token: "initial-refresh", + account_id: "main-account", + }); + + const token = await forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => { + writeAuth({ + access_token: externalAccess, + refresh_token: "external-refresh", + account_id: "external-account", + }); + return refreshed(fetchedAccess, "fetched-refresh", "fetched-account"); + }, + }, + }); + + expect(token).toEqual({ accessToken: externalAccess, chatgptAccountId: "external-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(externalAccess); + expect(persisted.tokens.refresh_token).toBe("external-refresh"); + expect(persisted.tokens.account_id).toBe("external-account"); + }); + + test("adopts an external auth.json writer instead of clobbering it during final publish", async () => { + const expiredAccess = jwt("expired-final-race", -3_600); + const externalAccess = jwt("external-final-race", 3_600, "external-final-account"); + const fetchedAccess = jwt("fetched-final-race", 3_600, "fetched-final-account"); + writeAuth({ + access_token: expiredAccess, + refresh_token: "initial-final-race-refresh", + account_id: "main-account", + }); + setMainAuthJsonPublishHookForTests(() => { + writeAuth({ + access_token: externalAccess, + refresh_token: "external-final-refresh", + account_id: "external-final-account", + }); + }); + + const token = await forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => refreshed(fetchedAccess, "fetched-final-refresh", "fetched-final-account"), + }, + }); + + expect(token).toEqual({ accessToken: externalAccess, chatgptAccountId: "external-final-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(externalAccess); + expect(persisted.tokens.refresh_token).toBe("external-final-refresh"); + expect(persisted.tokens.account_id).toBe("external-final-account"); + }); + + test("adopts an external auth.json writer instead of clobbering it after final validation", async () => { + const expiredAccess = jwt("expired-rename-race", -3_600); + const externalAccess = jwt("external-rename-race", 3_600, "external-rename-account"); + const fetchedAccess = jwt("fetched-rename-race", 3_600, "fetched-rename-account"); + writeAuth({ + access_token: expiredAccess, + refresh_token: "initial-rename-race-refresh", + account_id: "main-account", + }); + setMainAuthJsonRenameHookForTests(() => { + writeAuth({ + access_token: externalAccess, + refresh_token: "external-rename-refresh", + account_id: "external-rename-account", + }); + }); + + const token = await forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => refreshed(fetchedAccess, "fetched-rename-refresh", "fetched-rename-account"), + }, + }); + + expect(token).toEqual({ accessToken: externalAccess, chatgptAccountId: "external-rename-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(externalAccess); + expect(persisted.tokens.refresh_token).toBe("external-rename-refresh"); + expect(persisted.tokens.account_id).toBe("external-rename-account"); + }); + + test("adopts an external auth.json writer instead of clobbering it between final validation and replacement", async () => { + const expiredAccess = jwt("expired-replace-race", -3_600); + const externalAccess = jwt("external-replace-race", 3_600, "external-replace-account"); + const fetchedAccess = jwt("fetched-replace-race", 3_600, "fetched-replace-account"); + writeAuth({ + access_token: expiredAccess, + refresh_token: "initial-replace-race-refresh", + account_id: "main-account", + }); + setMainAuthJsonReplaceHookForTests(() => { + writeAuth({ + access_token: externalAccess, + refresh_token: "external-replace-refresh", + account_id: "external-replace-account", + }); + }); + + const token = await forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => refreshed(fetchedAccess, "fetched-replace-refresh", "fetched-replace-account"), + }, + }); + + expect(token).toEqual({ accessToken: externalAccess, chatgptAccountId: "external-replace-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(externalAccess); + expect(persisted.tokens.refresh_token).toBe("external-replace-refresh"); + expect(persisted.tokens.account_id).toBe("external-replace-account"); + }); + + test("publishes a native refresh to stored credentials that share the old refresh grant", async () => { + const expiredAccess = jwt("expired-shared", -3_600, "native-account"); + const freshAccess = jwt("fresh-shared", 3_600, "native-account"); + saveCodexAccountCredential("pool-shared", { + accessToken: jwt("pool-expired", -3_600, "pool-account"), + refreshToken: "shared-refresh", + expiresAt: Date.now() - 3_600_000, + chatgptAccountId: "pool-account", + }); + const initialRecord = readCodexAccountRecord("pool-shared")!; + writeAuth({ + access_token: expiredAccess, + refresh_token: "shared-refresh", + account_id: "native-account", + }); + + const token = await forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => refreshed(freshAccess, "shared-refresh-2", "native-account"), + }, + }); + + expect(token).toEqual({ accessToken: freshAccess, chatgptAccountId: "native-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(freshAccess); + expect(persisted.tokens.refresh_token).toBe("shared-refresh-2"); + const stored = getCodexAccountCredential("pool-shared")!; + expect(stored.accessToken).toBe(freshAccess); + expect(stored.refreshToken).toBe("shared-refresh-2"); + expect(stored.chatgptAccountId).toBe("pool-account"); + expect(readCodexAccountRecord("pool-shared")!.generation).toBe(initialRecord.generation + 1); + expect(readCodexAccountRecord("pool-shared")!.refreshGrantFingerprint) + .toBe(refreshGrantFingerprintForToken("shared-refresh-2")); + }); + + test("adopts a fresh stored credential that shares the native refresh grant without a network refresh", async () => { + const poolAccess = jwt("pool-shared-fresh", 3_600, "shared-account"); + saveCodexAccountCredential("pool-same-grant", { + accessToken: poolAccess, + refreshToken: "same-grant-refresh", + expiresAt: Date.now() + 3_600_000, + chatgptAccountId: "shared-account", + }); + writeAuth({ + access_token: jwt("native-same-grant-expired", -3_600, "shared-account"), + refresh_token: "same-grant-refresh", + account_id: "shared-account", + }); + let refreshCalls = 0; + + const token = await getValidMainAccountToken({ + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + throw new Error("network refresh should not run"); + }, + }, + }); + + expect(refreshCalls).toBe(0); + expect(token).toEqual({ accessToken: poolAccess, chatgptAccountId: "shared-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(poolAccess); + expect(persisted.tokens.refresh_token).toBe("same-grant-refresh"); + expect(persisted.tokens.account_id).toBe("shared-account"); + }); + + test("serializes concurrent native refreshes on the same grant and reuses the winner", async () => { + const expiredAccess = jwt("expired-concurrent", -3_600, "main-account"); + const freshAccess = jwt("fresh-concurrent", 3_600, "main-account"); + let refreshCalls = 0; + let releaseRefresh!: () => void; + const refreshMayFinish = new Promise(resolve => { + releaseRefresh = resolve; + }); + let firstRefreshStarted!: () => void; + const refreshStarted = new Promise(resolve => { + firstRefreshStarted = resolve; + }); + writeAuth({ + access_token: expiredAccess, + refresh_token: "concurrent-refresh", + account_id: "main-account", + }); + + const first = forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + firstRefreshStarted(); + await refreshMayFinish; + return refreshed(freshAccess, "concurrent-refresh-2", "main-account"); + }, + }, + }); + await refreshStarted; + const second = forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + return refreshed(jwt("unexpected-concurrent", 3_600), "unexpected-refresh", "main-account"); + }, + }, + }); + releaseRefresh(); + + await expect(Promise.all([first, second])).resolves.toEqual([ + { accessToken: freshAccess, chatgptAccountId: "main-account" }, + { accessToken: freshAccess, chatgptAccountId: "main-account" }, + ]); + expect(refreshCalls).toBe(1); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(freshAccess); + expect(persisted.tokens.refresh_token).toBe("concurrent-refresh-2"); + }); + + test("normalizes transient errors for callers joining an active native refresh", async () => { + const expiredAccess = jwt("expired-concurrent-transient", -3_600, "main-account"); + let releaseRefresh!: () => void; + const refreshMayFinish = new Promise(resolve => { + releaseRefresh = resolve; + }); + let firstRefreshStarted!: () => void; + const refreshStarted = new Promise(resolve => { + firstRefreshStarted = resolve; + }); + let joinedRefreshCalls = 0; + writeAuth({ + access_token: expiredAccess, + refresh_token: "concurrent-transient-refresh", + account_id: "main-account", + }); + + const owner = getValidMainAccountToken({ + dependencies: { + refreshToken: async () => { + firstRefreshStarted(); + await refreshMayFinish; + throw new Error("temporary upstream failure while refreshing"); + }, + }, + }).catch((error: unknown) => error); + await refreshStarted; + const joined = resolveCodexAuthContext(new Headers(), mainOnlyConfig(), "pool", { + nativeMainRefreshDependencies: { + refreshToken: async () => { + joinedRefreshCalls += 1; + return refreshed(jwt("unexpected-joined-transient", 3_600), "unexpected-refresh", "main-account"); + }, + }, + }).catch((error: unknown) => error); + releaseRefresh(); + + const [ownerError, joinedError] = await Promise.all([owner, joined]); + expect(ownerError).toBeInstanceOf(TokenRefreshError); + expect((ownerError as TokenRefreshError).reason).toBe("unknown"); + expect(joinedError).toBeInstanceOf(CodexAuthContextError); + expect((joinedError as Error).cause).toBeInstanceOf(TokenRefreshError); + expect(((joinedError as Error).cause as TokenRefreshError).reason).toBe("unknown"); + expect(joinedRefreshCalls).toBe(0); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); + + test("refreshes again when a joined flight returns the caller's rejected bearer", async () => { + const expiredAccess = jwt("expired-concurrent-rejected", -3_600, "main-account"); + const rejectedFreshAccess = jwt("fresh-concurrent-rejected", 3_600, "main-account"); + const replayFreshAccess = jwt("fresh-concurrent-replay", 3_600, "main-account"); + let releaseRefresh!: () => void; + const refreshMayFinish = new Promise(resolve => { + releaseRefresh = resolve; + }); + let firstRefreshStarted!: () => void; + const refreshStarted = new Promise(resolve => { + firstRefreshStarted = resolve; + }); + let firstRefreshCalls = 0; + let secondRefreshCalls = 0; + writeAuth({ + access_token: expiredAccess, + refresh_token: "concurrent-rejected-refresh", + account_id: "main-account", + }); + + const owner = forceRefreshMainAccountToken(expiredAccess, { + dependencies: { + refreshToken: async () => { + firstRefreshCalls += 1; + firstRefreshStarted(); + await refreshMayFinish; + return refreshed(rejectedFreshAccess, "concurrent-rejected-refresh-2", "main-account"); + }, + }, + }); + await refreshStarted; + const joined = forceRefreshMainAccountToken(rejectedFreshAccess, { + dependencies: { + refreshToken: async () => { + secondRefreshCalls += 1; + return refreshed(replayFreshAccess, "concurrent-rejected-refresh-3", "main-account"); + }, + }, + }); + releaseRefresh(); + + await expect(Promise.all([owner, joined])).resolves.toEqual([ + { accessToken: rejectedFreshAccess, chatgptAccountId: "main-account" }, + { accessToken: replayFreshAccess, chatgptAccountId: "main-account" }, + ]); + expect(firstRefreshCalls).toBe(1); + expect(secondRefreshCalls).toBe(1); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(replayFreshAccess); + expect(persisted.tokens.refresh_token).toBe("concurrent-rejected-refresh-3"); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); + + test("fails retryably when the native refresh grant changes while waiting for the old grant lock", async () => { + const grantAAccess = jwt("expired-grant-a", -3_600, "main-account"); + const grantBAccess = jwt("expired-grant-b", -3_600, "main-account"); + writeAuth({ + access_token: grantAAccess, + refresh_token: "grant-a-refresh", + account_id: "main-account", + }); + const lockPath = writeNativeRefreshLock("grant-a-refresh"); + let refreshCalls = 0; + + const attempt = forceRefreshMainAccountToken(grantAAccess, { + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + return refreshed(jwt("unexpected-grant-refresh", 3_600), "unexpected-refresh", "main-account"); + }, + }, + }); + writeAuth({ + access_token: grantBAccess, + refresh_token: "grant-b-refresh", + account_id: "main-account", + }); + unlinkSync(lockPath); + + await expect(attempt).rejects.toBeInstanceOf(MainAuthJsonChangedDuringRefreshError); + expect(refreshCalls).toBe(0); + expect(readAuth().tokens.refresh_token).toBe("grant-b-refresh"); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); + + test("fails retryably instead of accepting a fresh replacement grant under the old grant lock", async () => { + const grantAAccess = jwt("expired-grant-a-fresh-b", -3_600, "main-account"); + const grantBAccess = jwt("fresh-grant-b", 3_600, "main-account"); + writeAuth({ + access_token: grantAAccess, + refresh_token: "grant-a-refresh", + account_id: "main-account", + }); + const lockPath = writeNativeRefreshLock("grant-a-refresh"); + let refreshCalls = 0; + + const attempt = forceRefreshMainAccountToken(grantAAccess, { + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + return refreshed(jwt("unexpected-fresh-grant-refresh", 3_600), "unexpected-refresh", "main-account"); + }, + }, + }); + writeAuth({ + access_token: grantBAccess, + refresh_token: "grant-b-refresh", + account_id: "main-account", + }); + unlinkSync(lockPath); + + await expect(attempt).rejects.toBeInstanceOf(MainAuthJsonChangedDuringRefreshError); + expect(refreshCalls).toBe(0); + expect(readAuth().tokens.access_token).toBe(grantBAccess); + expect(readAuth().tokens.refresh_token).toBe("grant-b-refresh"); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); + + test("adopts one fresh stored credential for the same ChatGPT account when the match is unambiguous", async () => { + const poolAccess = jwt("pool-main-fresh", 3_600, "main-account"); + saveCodexAccountCredential("pool-main", { + accessToken: poolAccess, + refreshToken: "pool-main-refresh", + expiresAt: Date.now() + 3_600_000, + chatgptAccountId: "main-account", + }); + writeAuth({ + access_token: jwt("native-main-expired", -3_600, "main-account"), + refresh_token: "native-main-refresh", + account_id: "main-account", + }); + let refreshCalls = 0; + + const token = await getValidMainAccountToken({ + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + throw new Error("network refresh should not run"); + }, + }, + }); + + expect(refreshCalls).toBe(0); + expect(token).toEqual({ accessToken: poolAccess, chatgptAccountId: "main-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(poolAccess); + expect(persisted.tokens.refresh_token).toBe("pool-main-refresh"); + expect(persisted.tokens.account_id).toBe("main-account"); + }); + + test("does not adopt an ambiguous same-account stored credential", async () => { + const poolAccessOne = jwt("pool-main-one", 3_600, "main-account"); + const poolAccessTwo = jwt("pool-main-two", 3_600, "main-account"); + const networkAccess = jwt("network-main", 3_600, "main-account"); + saveCodexAccountCredential("pool-main-one", { + accessToken: poolAccessOne, + refreshToken: "pool-main-refresh-one", + expiresAt: Date.now() + 3_600_000, + chatgptAccountId: "main-account", + }); + saveCodexAccountCredential("pool-main-two", { + accessToken: poolAccessTwo, + refreshToken: "pool-main-refresh-two", + expiresAt: Date.now() + 3_600_000, + chatgptAccountId: "main-account", + }); + writeAuth({ + access_token: jwt("native-main-ambiguous-expired", -3_600, "main-account"), + refresh_token: "native-main-refresh", + account_id: "main-account", + }); + let refreshCalls = 0; + + const token = await getValidMainAccountToken({ + dependencies: { + refreshToken: async () => { + refreshCalls += 1; + return refreshed(networkAccess, "network-main-refresh", "main-account"); + }, + }, + }); + + expect(refreshCalls).toBe(1); + expect(token).toEqual({ accessToken: networkAccess, chatgptAccountId: "main-account" }); + const persisted = readAuth(); + expect(persisted.tokens.access_token).toBe(networkAccess); + expect(persisted.tokens.refresh_token).toBe("network-main-refresh"); + expect(persisted.tokens.access_token).not.toBe(poolAccessOne); + expect(persisted.tokens.access_token).not.toBe(poolAccessTwo); + }); + + test("refreshes a Direct admission substitution before returning upstream headers", async () => { + const expiredAccess = jwt("expired-direct", -3_600); + const freshAccess = jwt("fresh-direct", 3_600, "fresh-direct-account"); + writeAuth({ + access_token: expiredAccess, + refresh_token: "direct-refresh", + account_id: "main-account", + }); + + const headers = await materializeCodexUpstreamAuthAsync( + new Headers({ authorization: "Bearer admission-secret" }), + { kind: "main", accountId: null } satisfies CodexAuthContext, + { + substituteMainCredential: true, + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess, "direct-refresh-2", "fresh-direct-account"), + }, + }, + ); + + expect(headers.get("authorization")).toBe(`Bearer ${freshAccess}`); + expect(headers.get("chatgpt-account-id")).toBe("fresh-direct-account"); + }); + + test("refreshes before native main account-gated model discovery", async () => { + const expiredAccess = jwt("expired-entitlement", -3_600); + const freshAccess = jwt("fresh-entitlement", 3_600, "fresh-entitlement-account"); + let seenAuthorization = ""; + writeAuth({ + access_token: expiredAccess, + refresh_token: "entitlement-refresh", + account_id: "main-account", + }); + + const snapshot = await resolveCodexModelEntitlements(mainOnlyConfig(), { + fetcher: (async (_input, init) => { + seenAuthorization = new Headers(init?.headers).get("authorization") ?? ""; + return Response.json({ + models: [{ slug: NATIVE_DAYBREAK_BLUE_MODEL, supported_in_api: true, visibility: "list" }], + }); + }) as typeof fetch, + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess, "entitlement-refresh-2", "fresh-entitlement-account"), + }, + }); + + expect(seenAuthorization).toBe(`Bearer ${freshAccess}`); + expect([...entitledCodexAccountIdsForModel(snapshot, NATIVE_DAYBREAK_BLUE_MODEL)!]).toEqual([MAIN_CODEX_ACCOUNT_ID]); + }); + + test("marks the main account for reauthentication when the refresh grant is revoked", async () => { + writeAuth({ + access_token: jwt("expired-revoked", -3_600), + refresh_token: "revoked-refresh", + account_id: "main-account", + }); + + await expect(getValidMainAccountToken({ + dependencies: { + refreshToken: async () => { + throw new TokenRefreshError("revoked", "revoked"); + }, + }, + })).rejects.toBeInstanceOf(TokenRefreshError); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(true); + }); + + test("marks the main account for reauthentication on structured OAuth invalid_grant", async () => { + writeAuth({ + access_token: jwt("expired-structured-revoked", -3_600), + refresh_token: "structured-revoked-refresh", + account_id: "main-account", + }); + + await expect(getValidMainAccountToken({ + dependencies: { + refreshToken: async () => { + throw new ChatGPTTokenRefreshError( + 400, + "invalid_grant", + "refresh token revoked", + "ChatGPT refresh failed: 400 invalid_grant: refresh token revoked", + ); + }, + }, + })).rejects.toBeInstanceOf(TokenRefreshError); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(true); + }); + + test("does not mark the main account for reauthentication when native refresh fails transiently", async () => { + writeAuth({ + access_token: jwt("expired-transient", -3_600), + refresh_token: "transient-refresh", + account_id: "main-account", + }); + + await expect(resolveCodexAuthContext(new Headers(), mainOnlyConfig(), "pool", { + nativeMainRefreshDependencies: { + refreshToken: async () => { + throw new TokenRefreshError("unknown", "transient"); + }, + }, + })).rejects.toThrow(); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); + + test("does not mark the main account for reauthentication from generic expired transport errors", async () => { + writeAuth({ + access_token: jwt("expired-generic-error", -3_600), + refresh_token: "generic-error-refresh", + account_id: "main-account", + }); + + await expect(resolveCodexAuthContext(new Headers(), mainOnlyConfig(), "pool", { + nativeMainRefreshDependencies: { + refreshToken: async () => { + throw new Error("TLS certificate expired while refreshing token"); + }, + }, + })).rejects.toThrow(); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); + + test("passes native refresh dependencies into account-gated auth-context entitlement discovery", async () => { + const expiredAccess = jwt("expired-auth-context-entitlement", -3_600); + const freshAccess = jwt("fresh-auth-context-entitlement", 3_600, "entitled-main-account"); + const signal = new AbortController().signal; + let seenOptions: Pick< + CodexModelEntitlementResolveOptions, + "nativeMainRefreshDependencies" | "signal" + > | undefined; + writeAuth({ + access_token: expiredAccess, + refresh_token: "auth-context-entitlement-refresh", + account_id: "main-account", + }); + + const ctx = await resolveCodexAuthContext(new Headers(), mainOnlyConfig(), "pool", { + modelId: NATIVE_DAYBREAK_BLUE_MODEL, + signal, + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess, "auth-context-entitlement-refresh-2", "entitled-main-account"), + }, + resolveCodexModelEntitlements: async (_config, options) => { + seenOptions = options; + const token = await getValidMainAccountToken({ + dependencies: options?.nativeMainRefreshDependencies, + signal: options?.signal, + }); + expect(token).toEqual({ accessToken: freshAccess, chatgptAccountId: "entitled-main-account" }); + return { + modelsByAccount: new Map([[MAIN_CODEX_ACCOUNT_ID, new Set([NATIVE_DAYBREAK_BLUE_MODEL])]]), + confirmedAccountIds: new Set([MAIN_CODEX_ACCOUNT_ID]), + credentialIdentities: new Map([[MAIN_CODEX_ACCOUNT_ID, "main:entitled-main-account"]]), + }; + }, + }); + + expect(seenOptions?.signal).toBe(signal); + expect(seenOptions?.nativeMainRefreshDependencies).toBeDefined(); + expect(ctx).toMatchObject({ + kind: "main-pool", + accountId: MAIN_CODEX_ACCOUNT_ID, + accessToken: freshAccess, + chatgptAccountId: "entitled-main-account", + }); + }); +}); diff --git a/tests/responses-compact-native-main-refresh.test.ts b/tests/responses-compact-native-main-refresh.test.ts new file mode 100644 index 0000000000..958803c931 --- /dev/null +++ b/tests/responses-compact-native-main-refresh.test.ts @@ -0,0 +1,209 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { MAIN_CODEX_ACCOUNT_ID } from "../src/codex/main-account"; +import { clearAccountNeedsReauth, isAccountNeedsReauth } from "../src/codex/account-runtime-state"; +import { TokenRefreshError } from "../src/codex/account-store"; +import { clearAccountQuota } from "../src/codex/auth-api"; +import { resetMainCodexAccountIdentityTrackingForTests } from "../src/codex/account-lifecycle"; +import { clearCodexUpstreamHealth, clearThreadAccountMap } from "../src/codex/routing"; +import { handleResponsesCompact } from "../src/server/responses"; +import type { OAuthCredentials } from "../src/oauth/types"; +import type { OcxConfig } from "../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; + +const originalFetch = globalThis.fetch; + +let isolatedCodexHome: IsolatedCodexHome; +let opencodexHome: string; +let previousOpencodexHome: string | undefined; + +beforeEach(() => { + isolatedCodexHome = installIsolatedCodexHome("ocx-compact-main-refresh-codex-"); + previousOpencodexHome = process.env.OPENCODEX_HOME; + opencodexHome = mkdtempSync(join(tmpdir(), "ocx-compact-main-refresh-store-")); + process.env.OPENCODEX_HOME = opencodexHome; + clearThreadAccountMap(); + clearCodexUpstreamHealth(); + clearAccountQuota(); + resetMainCodexAccountIdentityTrackingForTests(); + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + clearThreadAccountMap(); + clearCodexUpstreamHealth(); + clearAccountQuota(); + resetMainCodexAccountIdentityTrackingForTests(); + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); + isolatedCodexHome.restore(); + rmSync(opencodexHome, { recursive: true, force: true }); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; +}); + +function jwt(name: string, expiresInSeconds: number, accountId = "main-account"): string { + const payload = Buffer.from(JSON.stringify({ + exp: Math.floor(Date.now() / 1000) + expiresInSeconds, + chatgpt_account_id: accountId, + marker: name, + })).toString("base64url"); + return `header.${payload}.signature`; +} + +function writeAuth(accessToken: string, refreshToken = "main-refresh", accountId = "main-account"): void { + writeFileSync(join(isolatedCodexHome.path, "auth.json"), JSON.stringify({ + tokens: { access_token: accessToken, refresh_token: refreshToken, account_id: accountId }, + })); +} + +function config(): OcxConfig { + return { + defaultProvider: "openai", + activeCodexAccountId: MAIN_CODEX_ACCOUNT_ID, + providers: { + openai: { + adapter: "openai-responses", + baseUrl: "https://chatgpt.com/backend-api/codex", + authMode: "forward", + codexAccountMode: "pool", + }, + }, + codexAccounts: [], + autoSwitchThreshold: 0, + } as OcxConfig; +} + +function request(): Request { + return new Request("http://localhost/v1/responses/compact", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "gpt-5.6-sol", + input: [ + { type: "message", role: "user", content: [{ type: "input_text", text: "earlier" }] }, + { type: "compaction_trigger" }, + ], + }), + }); +} + +function success(): Response { + return Response.json({ + output: [{ type: "compaction", encrypted_content: "opaque" }], + }); +} + +function refreshed(accessToken: string, accountId = "fresh-account"): OAuthCredentials { + return { + access: accessToken, + refresh: "fresh-refresh", + expires: Date.now() + 3_600_000, + accountId, + }; +} + +describe("Responses compact native main refresh", () => { + test("substitutes a refreshed native credential before compact upstream I/O", async () => { + const freshAccess = jwt("fresh-pre", 3_600, "fresh-account"); + writeAuth(jwt("expired-pre", -3_600)); + const bearers: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return success(); + }) as typeof fetch; + + const response = await handleResponsesCompact( + request(), + config(), + { model: "", provider: "" }, + undefined, + undefined, + { nativeMainRefreshDependencies: { refreshToken: async () => refreshed(freshAccess) } }, + ); + + expect(response.status).toBe(200); + expect(bearers).toEqual([`Bearer ${freshAccess}`]); + }); + + test("replays one compact native-main 401 with the refreshed bearer", async () => { + const initialAccess = jwt("initial", 3_600); + const freshAccess = jwt("fresh", 3_600, "fresh-account"); + writeAuth(initialAccess); + const bearers: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return bearers.length === 1 + ? Response.json({ error: { message: "expired" } }, { status: 401 }) + : success(); + }) as typeof fetch; + + const response = await handleResponsesCompact( + request(), + config(), + { model: "", provider: "" }, + undefined, + undefined, + { nativeMainRefreshDependencies: { refreshToken: async () => refreshed(freshAccess) } }, + ); + + expect(response.status).toBe(200); + expect(bearers).toEqual([`Bearer ${initialAccess}`, `Bearer ${freshAccess}`]); + }); + + test("does not replay compact native-main 401 more than once", async () => { + const initialAccess = jwt("initial-always", 3_600); + const freshAccess = jwt("fresh-always", 3_600, "fresh-account"); + writeAuth(initialAccess); + const bearers: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return Response.json({ error: { message: "still unauthorized" } }, { status: 401 }); + }) as typeof fetch; + + const response = await handleResponsesCompact( + request(), + config(), + { model: "", provider: "" }, + undefined, + undefined, + { nativeMainRefreshDependencies: { refreshToken: async () => refreshed(freshAccess) } }, + ); + + expect(response.status).toBe(401); + expect(bearers).toEqual([`Bearer ${initialAccess}`, `Bearer ${freshAccess}`]); + }); + + test("retryable compact 401 replay refresh failure does not mark main for reauthentication", async () => { + const initialAccess = jwt("initial-transient", 3_600); + writeAuth(initialAccess); + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + expect(new Headers(init?.headers).get("authorization")).toBe(`Bearer ${initialAccess}`); + return Response.json({ error: { message: "expired" } }, { status: 401 }); + }) as typeof fetch; + + const response = await handleResponsesCompact( + request(), + config(), + { model: "", provider: "" }, + undefined, + undefined, + { + nativeMainRefreshDependencies: { + refreshToken: async () => { + throw new TokenRefreshError("unknown", "transient"); + }, + }, + }, + ); + + expect(response.status).toBe(503); + expect(isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + }); +}); diff --git a/tests/responses-native-main-refresh.test.ts b/tests/responses-native-main-refresh.test.ts new file mode 100644 index 0000000000..0c77d03317 --- /dev/null +++ b/tests/responses-native-main-refresh.test.ts @@ -0,0 +1,287 @@ +import { afterEach, beforeEach, describe, expect, spyOn, test } from "bun:test"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { TokenRefreshError } from "../src/codex/account-store"; +import { MAIN_CODEX_ACCOUNT_ID } from "../src/codex/main-account"; +import { clearAccountNeedsReauth } from "../src/codex/account-runtime-state"; +import { clearAccountQuota } from "../src/codex/auth-api"; +import { clearCodexUpstreamHealth, clearThreadAccountMap } from "../src/codex/routing"; +import { resetMainCodexAccountIdentityTrackingForTests } from "../src/codex/account-lifecycle"; +import { handleResponses } from "../src/server/responses"; +import * as adapterResolveModule from "../src/server/adapter-resolve"; +import type { OAuthCredentials } from "../src/oauth/types"; +import type { RequestLogContext } from "../src/server/request-log"; +import type { OcxConfig } from "../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; + +const originalFetch = globalThis.fetch; + +let isolatedCodexHome: IsolatedCodexHome; +let opencodexHome: string; +let previousOpencodexHome: string | undefined; + +beforeEach(() => { + isolatedCodexHome = installIsolatedCodexHome("ocx-responses-main-refresh-codex-"); + previousOpencodexHome = process.env.OPENCODEX_HOME; + opencodexHome = mkdtempSync(join(tmpdir(), "ocx-responses-main-refresh-store-")); + process.env.OPENCODEX_HOME = opencodexHome; + clearThreadAccountMap(); + clearCodexUpstreamHealth(); + clearAccountQuota(); + resetMainCodexAccountIdentityTrackingForTests(); + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); +}); + +afterEach(() => { + globalThis.fetch = originalFetch; + clearThreadAccountMap(); + clearCodexUpstreamHealth(); + clearAccountQuota(); + resetMainCodexAccountIdentityTrackingForTests(); + clearAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID); + isolatedCodexHome.restore(); + rmSync(opencodexHome, { recursive: true, force: true }); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; +}); + +function jwt(name: string, expiresInSeconds: number, accountId = "main-account"): string { + const payload = Buffer.from(JSON.stringify({ + exp: Math.floor(Date.now() / 1000) + expiresInSeconds, + chatgpt_account_id: accountId, + marker: name, + })).toString("base64url"); + return `header.${payload}.signature`; +} + +function writeAuth(accessToken: string, refreshToken = "main-refresh", accountId = "main-account"): void { + writeFileSync(join(isolatedCodexHome.path, "auth.json"), JSON.stringify({ + tokens: { access_token: accessToken, refresh_token: refreshToken, account_id: accountId }, + })); +} + +function config(): OcxConfig { + return { + defaultProvider: "openai", + activeCodexAccountId: MAIN_CODEX_ACCOUNT_ID, + providers: { + openai: { + adapter: "openai-responses", + baseUrl: "https://chatgpt.com/backend-api/codex", + authMode: "forward", + codexAccountMode: "pool", + }, + }, + codexAccounts: [], + autoSwitchThreshold: 0, + } as OcxConfig; +} + +function request(): Request { + return new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "gpt-5.6-sol", + stream: false, + store: false, + input: [{ type: "message", role: "user", content: [{ type: "input_text", text: "hi" }] }], + }), + }); +} + +function success(): Response { + return Response.json({ + id: "resp_1", + object: "response", + status: "completed", + model: "gpt-5.6-sol", + output: [], + }); +} + +function refreshed(accessToken: string, accountId = "fresh-account"): OAuthCredentials { + return { + access: accessToken, + refresh: "fresh-refresh", + expires: Date.now() + 3_600_000, + accountId, + }; +} + +describe("Responses native main refresh", () => { + test("replays one native-main 401 with the refreshed bearer", async () => { + const initialAccess = jwt("initial", 3_600); + const freshAccess = jwt("fresh", 3_600, "fresh-account"); + writeAuth(initialAccess); + const bearers: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return bearers.length === 1 + ? Response.json({ error: { message: "expired" } }, { status: 401 }) + : success(); + }) as typeof fetch; + + const logCtx: RequestLogContext = { model: "", provider: "" }; + const response = await handleResponses(request(), config(), logCtx, { + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess), + }, + }); + + expect(response.status).toBe(200); + expect(bearers).toEqual([`Bearer ${initialAccess}`, `Bearer ${freshAccess}`]); + expect(logCtx.activeAttempt?.recoveryKinds).toEqual(["codex-main-401"]); + }); + + test("does not replay native-main 401 more than once", async () => { + const initialAccess = jwt("initial-always", 3_600); + const freshAccess = jwt("fresh-always", 3_600, "fresh-account"); + writeAuth(initialAccess); + const bearers: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return Response.json({ error: { message: "still unauthorized" } }, { status: 401 }); + }) as typeof fetch; + + const response = await handleResponses(request(), config(), { model: "", provider: "" }, { + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess), + }, + }); + + expect(response.status).toBe(401); + expect(bearers).toEqual([`Bearer ${initialAccess}`, `Bearer ${freshAccess}`]); + }); + + test("replays one generic-adapter native-main 401 with the refreshed bearer", async () => { + const initialAccess = jwt("initial-generic", 3_600); + const freshAccess = jwt("fresh-generic", 3_600, "fresh-account"); + writeAuth(initialAccess); + const bearers: string[] = []; + const adapterSpy = spyOn(adapterResolveModule, "resolveAdapter").mockReturnValue({ + name: "test-generic", + buildRequest: (_parsed, incoming) => ({ + url: "https://fixture.test/v1/responses", + method: "POST", + headers: Object.fromEntries(incoming.headers.entries()), + body: "{}", + }), + async parseResponse() { + return [{ type: "text_delta", text: "ok" }, { type: "done" }]; + }, + async *parseStream() { + yield { type: "text_delta", text: "ok" }; + yield { type: "done" }; + }, + } as ReturnType); + try { + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return bearers.length === 1 + ? Response.json({ error: { message: "expired" } }, { status: 401 }) + : success(); + }) as typeof fetch; + + const logCtx: RequestLogContext = { model: "", provider: "" }; + const response = await handleResponses(request(), config(), logCtx, { + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess), + }, + }); + + expect(response.status).toBe(200); + expect(bearers).toEqual([`Bearer ${initialAccess}`, `Bearer ${freshAccess}`]); + expect(logCtx.activeAttempt?.recoveryKinds).toEqual(["codex-main-401"]); + } finally { + adapterSpy.mockRestore(); + } + }); + + test("does not replay generic-adapter native-main 401 more than once", async () => { + const initialAccess = jwt("initial-generic-always", 3_600); + const freshAccess = jwt("fresh-generic-always", 3_600, "fresh-account"); + writeAuth(initialAccess); + const bearers: string[] = []; + const adapterSpy = spyOn(adapterResolveModule, "resolveAdapter").mockReturnValue({ + name: "test-generic", + buildRequest: (_parsed, incoming) => ({ + url: "https://fixture.test/v1/responses", + method: "POST", + headers: Object.fromEntries(incoming.headers.entries()), + body: "{}", + }), + async parseResponse() { + return [{ type: "text_delta", text: "should not parse" }, { type: "done" }]; + }, + async *parseStream() { + yield { type: "text_delta", text: "should not parse" }; + yield { type: "done" }; + }, + } as ReturnType); + try { + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + bearers.push(new Headers(init?.headers).get("authorization") ?? ""); + return Response.json({ error: { message: "still unauthorized" } }, { status: 401 }); + }) as typeof fetch; + + const response = await handleResponses(request(), config(), { model: "", provider: "" }, { + nativeMainRefreshDependencies: { + refreshToken: async () => refreshed(freshAccess), + }, + }); + + expect(response.status).toBe(401); + expect(bearers).toEqual([`Bearer ${initialAccess}`, `Bearer ${freshAccess}`]); + } finally { + adapterSpy.mockRestore(); + } + }); + + test("logs retryable native-main refresh failures without reauthentication guidance", async () => { + writeAuth(jwt("expired-transient-log", -3_600)); + const errors: string[] = []; + const errorSpy = spyOn(console, "error").mockImplementation((...args) => { + errors.push(args.map(String).join(" ")); + }); + try { + const response = await handleResponses(request(), config(), { model: "", provider: "" }, { + nativeMainRefreshDependencies: { + refreshToken: async () => { + throw new TokenRefreshError("unknown", "transient"); + }, + }, + }); + + expect(response.status).toBe(503); + expect(errors.some(line => line.includes("retryable refresh failure"))).toBe(true); + expect(errors.some(line => line.includes("reauthentication required"))).toBe(false); + } finally { + errorSpy.mockRestore(); + } + }); + + test("returns auth failure before upstream I/O when pre-request refresh is revoked", async () => { + writeAuth(jwt("expired-revoked", -3_600)); + let sends = 0; + globalThis.fetch = (async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) return Response.json({}); + sends += 1; + return success(); + }) as typeof fetch; + + const response = await handleResponses(request(), config(), { model: "", provider: "" }, { + nativeMainRefreshDependencies: { + refreshToken: async () => { + throw new TokenRefreshError("revoked", "revoked"); + }, + }, + }); + + expect(response.status).toBe(401); + expect(sends).toBe(0); + }); +}); diff --git a/tests/routing-analytics.test.ts b/tests/routing-analytics.test.ts index 171a185948..832ea60f88 100644 --- a/tests/routing-analytics.test.ts +++ b/tests/routing-analytics.test.ts @@ -6,6 +6,7 @@ import { handleManagementAPI } from "../src/server/management-api"; import { ManagementRequest } from "./helpers/management-auth"; import { appendUsageEntry, + readUsageEntries, resetUsageReadCacheForTests, usageLogPath, type PersistedUsageEntry, @@ -225,6 +226,31 @@ describe("routing analytics (RI-03)", () => { expect(result.cooldownTriggeringFailures).toBe(1); }); + test("retains codex-main-401 recovery attempts for cooldown analytics", async () => { + appendFileSync(usageLogPath(), `${JSON.stringify(entry("native-main-replay", { + timestamp: 1, + status: 503, + durationMs: 100, + attempts: [ + { + ordinal: 1, + provider: "openai", + model: "gpt-5.6-sol", + adapter: "openai-responses", + status: 401, + durationMs: 40, + sendCount: 1, + recoveryKinds: ["codex-main-401", "unknown"], + usageStatus: "unreported", + }, + ], + }))}\n`); + + expect(readUsageEntries()[0]?.attempts?.[0]?.recoveryKinds).toEqual(["codex-main-401"]); + const result = await computeRoutingAnalytics({}); + expect(result.cooldownTriggeringFailures).toBe(1); + }); + test("ignores malformed attempts while preserving explicit 429 classification", async () => { appendUsageEntry(entry("baseline", { timestamp: 1, status: 200, durationMs: 10 })); const historicalRows = [