From e8fe6d58da55ae72aa729e8f76f3baeede8ab15d Mon Sep 17 00:00:00 2001 From: iceteaSA <171169159+iceteaSA@users.noreply.github.com> Date: Wed, 22 Jul 2026 21:08:24 +0200 Subject: [PATCH] fix(routing): skip known-exhausted accounts at admission --- packages/opencode/src/core/accounts.ts | 2 + packages/opencode/src/core/quota-manager.ts | 35 +- packages/opencode/src/index.ts | 356 +++++++++-- packages/opencode/src/sidebar-state.ts | 132 +++- .../opencode/src/tests/integration.test.ts | 604 ++++++++++++++++++ .../opencode/src/tests/quota-push.test.ts | 76 ++- .../opencode/src/tests/sidebar-state.test.ts | 373 ++++++++++- 7 files changed, 1497 insertions(+), 81 deletions(-) diff --git a/packages/opencode/src/core/accounts.ts b/packages/opencode/src/core/accounts.ts index 35fa5d6..9b0f6b5 100644 --- a/packages/opencode/src/core/accounts.ts +++ b/packages/opencode/src/core/accounts.ts @@ -1514,6 +1514,8 @@ export class FallbackAccountManager { checkedAt, }, account.access, + false, + account.accountId, ) } diff --git a/packages/opencode/src/core/quota-manager.ts b/packages/opencode/src/core/quota-manager.ts index 8a8cd1a..7512655 100644 --- a/packages/opencode/src/core/quota-manager.ts +++ b/packages/opencode/src/core/quota-manager.ts @@ -156,6 +156,11 @@ export type QuotaEntry = { quota: OAuthQuotaSnapshot refreshAfter: number // Unix ms — earliest next refresh checkedAt: number // when snapshot was fetched + // ChatGPT identity of the account this snapshot was captured under. Lets a + // policy read detect a re-login (identity switch) on a stable internal id, so + // the previous identity's quota is never applied to its replacement. Absent + // when the identity was never recorded (fail-open, best-effort). + accountId?: string } export type QuotaManagerOptions = { @@ -343,11 +348,24 @@ export class QuotaManager { /** * Non-invalidating read of a cached fallback quota for POLICY decisions - * (killswitch). Keyed by the stable internal account id, so a token refresh - * for the same account does not drop it (getFallback(id, token) would). + * (killswitch, admission). Keyed by the stable internal account id, so a + * token refresh for the same account does not drop it (getFallback(id, token) + * would). + * + * Pass the live ChatGPT identity to still drop on a genuine re-login: if the + * cached entry's identity is known and differs, return null so the previous + * identity's (possibly exhausted) quota is never judged against its + * replacement. Unknown on either side fails open, matching peekMainForPolicy. */ - peekFallbackForPolicy(accountId: string): QuotaEntry | null { - return this.fallbacks.get(accountId) ?? null + peekFallbackForPolicy( + accountId: string, + identity?: string, + ): QuotaEntry | null { + const entry = this.fallbacks.get(accountId) ?? null + if (identity && entry?.accountId && entry.accountId !== identity) { + return null + } + return entry } setFallback( @@ -355,11 +373,16 @@ export class QuotaManager { entry: QuotaEntry, accessToken?: string, completeSnapshot = false, + identity?: string, ): void { // Empty snapshots replace cache only when their transport guarantees a // complete frame; otherwise a non-quota response could erase good data. if (!completeSnapshot && !hasAnyQuotaWindow(entry.quota)) return - this.fallbacks.set(accountId, entry) + // Bind the snapshot to the identity it was captured under so a re-login on + // this stable id stays detectable. An unknown incoming identity preserves + // the prior binding (mirrors setMain) rather than erasing a known one. + const boundIdentity = identity ?? this.fallbacks.get(accountId)?.accountId + this.fallbacks.set(accountId, { ...entry, accountId: boundIdentity }) if (accessToken) { this.fallbackTokenFps.set(accountId, tokenFingerprint(accessToken)) } else { @@ -623,6 +646,8 @@ export class QuotaManager { checkedAt, }, account.access, + false, + account.accountId, ) } } diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index bf5e249..acb7942 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -66,6 +66,7 @@ import { } from './core/oauth' import { codexRefreshFn, whamUsageFn } from './core/provider' import { + type QuotaEntry, QuotaManager, resolveMidStreamRateLimitResetAt, } from './core/quota-manager' @@ -98,7 +99,10 @@ import { getRpcDir } from './rpc/rpc-dir' import { type RpcServerHandle, startRpcServer } from './rpc/rpc-server' import { type AccountQuota, + exhaustedQuotaResetAt, + getSidebarState, getSidebarStateFile, + isQuotaExhausted, removeSidebarActiveRouting, type SidebarMachineState, type SidebarState, @@ -465,11 +469,18 @@ export function buildSidebarMachineState( qm: QuotaManager, store: AccountStorage, now = Date.now(), + mainAccountIdentity = store.mainAccountId, ): SidebarMachineState { - const mainQuota = qm.getMain()?.quota + const mainEntry = qm.getMain() + const mainQuota = mainEntry?.quota return { main: { - quota: (mainQuota as AccountQuota | undefined) ?? null, + quota: mainQuota + ? { ...(mainQuota as AccountQuota), checkedAt: mainEntry.checkedAt } + : null, + ...(typeof mainAccountIdentity === 'string' + ? { mainAccountId: mainAccountIdentity } + : {}), killed: false, ...(mainQuota?.resetCreditsAvailable !== undefined ? { resetCredits: mainQuota.resetCreditsAvailable } @@ -478,11 +489,22 @@ export function buildSidebarMachineState( fallbacks: store.accounts .filter((account) => account.enabled) .map((account) => { - const fallbackQuota = qm.getFallback(account.id)?.quota + const fallbackEntry = qm.getFallback(account.id) + const fallbackQuota = fallbackEntry?.quota return { id: account.id, label: (account as { label?: string }).label, - quota: (fallbackQuota as AccountQuota | undefined) ?? null, + // Identity follows the cached snapshot (the identity that quota was + // captured under), falling back to the live account — so a re-login + // never pairs the previous identity's stale quota with the new one. + accountId: + fallbackEntry?.accountId ?? (account as OAuthAccount).accountId, + quota: fallbackQuota + ? { + ...(fallbackQuota as AccountQuota), + checkedAt: fallbackEntry.checkedAt, + } + : null, killed: false, enabled: true, ...(fallbackQuota?.resetCreditsAvailable !== undefined @@ -1217,7 +1239,18 @@ export async function CodexAuthPlugin( refreshAfter: now + 5 * 60 * 1000, checkedAt: now, } + let resolvedMainIdentity: string | undefined if (accountId && accountId !== 'main') { + // Bind the pushed snapshot to the ChatGPT identity of the token + // that carried it (same derivation as main below), so a re-login on + // this stable id stays detectable and the sidebar never pairs a + // stale identity's quota with the new one. + const fallbackClaims = accessToken + ? parseJwtClaims(accessToken) + : null + const fallbackIdentity = fallbackClaims + ? extractAccountIdFromClaims(fallbackClaims) + : undefined // Quota-bearing transports report every live window rather than a // partial subset, so an absent slot means the wire dropped it. // Only the wham-only reset-credit metadata carries forward. @@ -1233,9 +1266,10 @@ export async function CodexAuthPlugin( entry, accessToken, completeSnapshot, + fallbackIdentity, ) } else { - let resolvedMainIdentity = mainAccountIdentity + resolvedMainIdentity = mainAccountIdentity if (!resolvedMainIdentity && accessToken) { const claims = parseJwtClaims(accessToken) resolvedMainIdentity = claims @@ -1274,7 +1308,13 @@ export async function CodexAuthPlugin( snapshot: entry.quota, }) const latestStorage = await loadRequestAccounts() - await writeMachineSidebarState(quotaManager, latestStorage) + await writeMachineSidebarState( + quotaManager, + latestStorage, + accountId && accountId !== 'main' + ? latestStorage?.mainAccountId + : resolvedMainIdentity, + ) } // The pure resolver prefers an admission error's explicit reset and @@ -1365,10 +1405,16 @@ export async function CodexAuthPlugin( async function writeMachineSidebarState( qm: QuotaManager, store: Awaited>, + mainAccountIdentity = store?.mainAccountId, ) { if (!store) return await setSidebarMachineState( - buildSidebarMachineState(qm, store), + buildSidebarMachineState( + qm, + store, + Date.now(), + mainAccountIdentity, + ), boundSidebarFile, ) } @@ -1661,24 +1707,101 @@ export async function CodexAuthPlugin( return quotaManager.peekMainForPolicy(mainAccountIdentity)?.quota } + // ----------------------------------------------------------------- + // Admission quota consultation: before probing an account, check the + // shared sidebar file (written by every plugin process on the machine) + // against the in-process cache, so a fresh process never spends a + // doomed request on an account already known to be at 100%. + // ----------------------------------------------------------------- + type AdmissionQuotaDecision = + | { exhausted: false } + | { + exhausted: true + source: 'memory' | 'file' + resetsAt: string + resetAtMs: number + } + + // Freshness key for a quota snapshot: the window's own checkedAt, then + // the snapshot-level checkedAt, then the cache entry's checkedAt. + function quotaCheckedAt( + quota: AccountQuota | null | undefined, + entryCheckedAt?: number, + ): number | undefined { + for (const checkedAt of [ + quota?.primary?.checkedAt, + quota?.checkedAt, + entryCheckedAt, + ]) { + if (typeof checkedAt === 'number' && Number.isFinite(checkedAt)) { + return checkedAt + } + } + return undefined + } + + // Selects the fresher quota source and judges it. The file wins only + // when strictly newer; memory wins ties; empty memory defers to a + // valid file. A file row stamped with a DIFFERENT account identity is + // treated as absent (fail-open) — a re-login must never be judged by + // the previous account's exhaustion. + function admissionQuotaDecision( + memoryEntry: QuotaEntry | null, + fileQuota: AccountQuota | null | undefined, + now: number, + fileAccountId?: string, + currentAccountId?: string, + ): AdmissionQuotaDecision { + const memoryQuota = memoryEntry?.quota as AccountQuota | undefined + const memoryCheckedAt = quotaCheckedAt( + memoryQuota, + memoryEntry?.checkedAt, + ) + const fileCheckedAt = quotaCheckedAt(fileQuota) + const useFile = + (fileAccountId === undefined || + fileAccountId === currentAccountId) && + fileQuota != null && + (memoryQuota === undefined || + (fileCheckedAt !== undefined && + (memoryCheckedAt === undefined || + fileCheckedAt > memoryCheckedAt))) + const source = useFile ? 'file' : 'memory' + const quota = useFile ? fileQuota : memoryQuota + if (!isQuotaExhausted(quota, now)) return { exhausted: false } + + const reset = exhaustedQuotaResetAt(quota, now) + if (!reset) return { exhausted: false } + return { exhausted: true, source, ...reset } + } + + function logAdmissionQuotaSkip( + accountId: string, + decision: Extract, + ) { + logQ.debug('admission skip: exhausted account', { + accountId, + source: decision.source, + resetsAt: decision.resetsAt, + }) + } + // Synthetic provider-shaped 429 with a Retry-After derived from the // earliest known quota reset across all accounts. Returned when the - // killswitch blocks the primary and no surviving account can serve — - // or when a mid-stream rate-limit mark blocks it instead (killswitch - // may be OFF in that case), hence the parametrized reason: the - // retry-after math is the same earliest-known-reset computation - // either way, only the human-readable message differs. + // killswitch blocks the primary and no surviving account can serve, or + // when admission knows the primary is temporarily unavailable for one + // of the unconditional quota reasons. // - // markResetAtMs carries the mid-stream mark's OWN stored reset (~60s, - // DEFAULT_MID_STREAM_RATE_LIMIT_RESET_MS) when the block is a mark, not - // a killswitch decision. killswitchRetryAfterSeconds falls back to a - // 300s default when no account reset is known — without this, a bare - // mid-stream block (no cached quota yet) would tell the client to wait - // 4 extra minutes past the mark's actual expiry. The sooner of the two - // always wins — this only ever tightens the wait, never lengthens it. + // markResetAtMs carries the blocking account's own reset when the reason + // has one. A known-exhausted quota uses that reset exactly; a transient + // mid-stream mark takes the sooner of its bounded reset and any cached + // quota reset so a missing quota snapshot cannot overstate the wait. function killswitchBlockedResponse( storage: AccountStorage | null, - reason: 'killswitch' | 'mid-stream-rate-limit' = 'killswitch', + reason: + | 'killswitch' + | 'mid-stream-rate-limit' + | 'quota-exhausted' = 'killswitch', markResetAtMs?: number, ): Response { const now = Date.now() @@ -1698,7 +1821,9 @@ export async function CodexAuthPlugin( now, storage, ) - if ( + if (reason === 'quota-exhausted' && markResetAtMs !== undefined) { + retryAfter = Math.max(1, Math.ceil((markResetAtMs - now) / 1000)) + } else if ( reason === 'mid-stream-rate-limit' && markResetAtMs !== undefined ) { @@ -1713,7 +1838,9 @@ export async function CodexAuthPlugin( const message = reason === 'mid-stream-rate-limit' ? `OpenAI rate limit reached — retrying on another account. Retry in ${mins}m ${secs}s.` - : `Killswitch: all OpenAI accounts are below their configured quota threshold. Retry in ${mins}m ${secs}s.` + : reason === 'quota-exhausted' + ? `OpenAI quota exhausted — retrying on another account. Retry in ${mins}m ${secs}s.` + : `Killswitch: all OpenAI accounts are below their configured quota threshold. Retry in ${mins}m ${secs}s.` return new Response( JSON.stringify({ error: { @@ -1745,9 +1872,19 @@ export async function CodexAuthPlugin( fallback: FallbackAccount } + type FallbackCandidateSelection = { + current: FallbackCandidate[] + retained: FallbackCandidate[] + skipped: Array<{ + candidate: FallbackCandidate + decision: Extract + }> + } + async function usableFallbackCandidates( fallbackStorage: Awaited>, - ): Promise { + sidebarState: SidebarState, + ): Promise { const usableFallbacks = await fallbackManager.getUsableFallbackAccounts(fallbackStorage) const candidates: FallbackCandidate[] = [] @@ -1773,15 +1910,66 @@ export async function CodexAuthPlugin( // threshold so routing never spends on a killed account. Opt-in — a // no-op when disabled. Non-invalidating peek so a token refresh does // not flip a killed account to "unknown". - if (!isKillswitchEnabled(fallbackStorage)) return notRateLimited - return notRateLimited.filter((c) => - killswitchPassesPolicy( - quotaManager.peekFallbackForPolicy(c.quotaAccountId)?.quota, - fallbackStorage, - c.quotaAccountId, - Date.now(), - ), + const current = isKillswitchEnabled(fallbackStorage) + ? notRateLimited.filter((c) => + killswitchPassesPolicy( + quotaManager.peekFallbackForPolicy( + c.quotaAccountId, + c.accountId, + )?.quota, + fallbackStorage, + c.quotaAccountId, + Date.now(), + ), + ) + : notRateLimited + // Admission: drop candidates whose freshest known quota (memory or + // the shared sidebar file) is exhausted, so a doomed probe is never + // spent on them. + const fileQuotas = new Map( + sidebarState.fallbacks.map((account) => [ + account.id, + { quota: account.quota, accountId: account.accountId }, + ]), ) + const retained: FallbackCandidate[] = [] + const skipped: FallbackCandidateSelection['skipped'] = [] + const now = Date.now() + for (const candidate of current) { + const fileEntry = fileQuotas.get(candidate.quotaAccountId) + const decision = admissionQuotaDecision( + quotaManager.peekFallbackForPolicy( + candidate.quotaAccountId, + candidate.accountId, + ), + fileEntry?.quota, + now, + fileEntry?.accountId, + candidate.accountId, + ) + if (decision.exhausted) { + skipped.push({ candidate, decision }) + } else { + retained.push(candidate) + } + } + return { current, retained, skipped } + } + + function applyAdmissionQuotaSafety( + selection: FallbackCandidateSelection, + mainRemainsCandidate: boolean, + ): FallbackCandidate[] { + // The shared file is advisory. If filtering it would remove the last + // admission path, keep probing in the original wire-authority order — + // a stale or corrupt file can never brick routing. + if (selection.retained.length === 0 && !mainRemainsCandidate) { + return selection.current + } + for (const { candidate, decision } of selection.skipped) { + logAdmissionQuotaSkip(candidate.quotaAccountId, decision) + } + return selection.retained } async function pushFailedFallbackQuota( @@ -1811,6 +1999,7 @@ export async function CodexAuthPlugin( requestInput: RequestInfo | URL, init: RequestInit | undefined, fallbackStorage: Awaited>, + candidates: FallbackCandidate[], ): Promise< | { response: Response @@ -1820,7 +2009,6 @@ export async function CodexAuthPlugin( } | undefined > { - const candidates = await usableFallbackCandidates(fallbackStorage) for (const candidate of candidates) { let response: Response try { @@ -1877,14 +2065,13 @@ export async function CodexAuthPlugin( requestInput: RequestInfo | URL, init: RequestInit | undefined, primaryResponse: Response, - _unused?: string, + fallbackStorage: Awaited>, + candidates: FallbackCandidate[], ) { if (!isReplayableRequest(requestInput, init)) { return { response: primaryResponse } } - const fallbackStorage = await loadRequestAccounts() - const candidates = await usableFallbackCandidates(fallbackStorage) if (!candidates.length) return { response: primaryResponse } // Keep the returned response body live; only cancel a response after a @@ -2060,6 +2247,40 @@ export async function CodexAuthPlugin( currentMainIdentity = mainAccountIdentity } const mode: RoutingMode = reqStorage?.routing?.mode ?? 'main-first' + // One shared sidebar read per request: admission decisions for main + // and the fallbacks must all judge the same snapshot. + let requestSidebarStatePromise: Promise | undefined + const requestSidebarState = () => { + requestSidebarStatePromise ??= getSidebarState(boundSidebarFile) + return requestSidebarStatePromise + } + const sidebarState = await requestSidebarState() + const mainQuotaDecision = admissionQuotaDecision( + quotaManager.peekMainForPolicy(mainAccountIdentity), + sidebarState.main.quota, + Date.now(), + sidebarState.main.mainAccountId, + mainAccountIdentity, + ) + const killswitchBlocksMain = + isKillswitchEnabled(reqStorage) && + !killswitchPassesPolicy( + killswitchMainQuota(mainAccountIdentity), + reqStorage, + undefined, + Date.now(), + ) + const mainRateLimited = quotaManager.isRateLimited('main') + let fallbackSelectionPromise: + | Promise + | undefined + const requestFallbackSelection = () => { + fallbackSelectionPromise ??= usableFallbackCandidates( + reqStorage, + sidebarState, + ) + return fallbackSelectionPromise + } // fallback-first (proactive): try usable fallbacks BEFORE main. If one // serves, use it and skip main entirely; otherwise fall through to the @@ -2072,12 +2293,25 @@ export async function CodexAuthPlugin( // True when the proactive gate already tried every usable fallback, // so the reactive path below must not re-try (and re-spend on) them. let fallbacksAlreadyTried = false + let fallbackCandidates: FallbackCandidate[] | undefined if ( mode === 'fallback-first' && isReplayableRequest(requestInput, init) ) { fallbacksAlreadyTried = true - const pre = await tryFallbackFirst(requestInput, init, reqStorage) + const selection = await requestFallbackSelection() + fallbackCandidates = applyAdmissionQuotaSafety( + selection, + !killswitchBlocksMain && + !mainRateLimited && + !mainQuotaDecision.exhausted, + ) + const pre = await tryFallbackFirst( + requestInput, + init, + reqStorage, + fallbackCandidates, + ) if (pre) { response = pre.response servedFallback = { @@ -2103,21 +2337,25 @@ export async function CodexAuthPlugin( // cached quota may not yet reflect it, so block main here too until // the mark's reset passes. if (!response) { - const killswitchBlocksMain = - isKillswitchEnabled(reqStorage) && - !killswitchPassesPolicy( - killswitchMainQuota(mainAccountIdentity), - reqStorage, - undefined, - Date.now(), - ) - const mainRateLimited = quotaManager.isRateLimited('main') - if (killswitchBlocksMain || mainRateLimited) { - const blockReason = - mainRateLimited && !killswitchBlocksMain + // Admission: a main the shared data already knows to be exhausted + // is blocked only when a non-exhausted fallback survives to serve + // — otherwise the probe stands (the wire is the final authority). + const quotaBlocksMain = + mainQuotaDecision.exhausted && + (await requestFallbackSelection()).retained.length > 0 + if (killswitchBlocksMain || mainRateLimited || quotaBlocksMain) { + const blockReason = killswitchBlocksMain + ? 'killswitch' + : mainRateLimited ? 'mid-stream-rate-limit' - : 'killswitch' - logA.debug('killswitch blocked primary', { + : 'quota-exhausted' + if ( + blockReason === 'quota-exhausted' && + mainQuotaDecision.exhausted + ) { + logAdmissionQuotaSkip('main', mainQuotaDecision) + } + logA.debug('admission blocked primary', { pid: process.pid, activeId: 'main', reason: blockReason, @@ -2125,9 +2363,12 @@ export async function CodexAuthPlugin( response = killswitchBlockedResponse( reqStorage, blockReason, - mainRateLimited - ? quotaManager.rateLimitedUntil('main') - : undefined, + blockReason === 'quota-exhausted' && + mainQuotaDecision.exhausted + ? mainQuotaDecision.resetAtMs + : mainRateLimited + ? quotaManager.rateLimitedUntil('main') + : undefined, ) } else { // Send through the main account. @@ -2164,11 +2405,18 @@ export async function CodexAuthPlugin( pid: process.pid, status: response.status, }) + // Main already failed, so it is no longer a surviving candidate + // for the admission safety valve. + fallbackCandidates ??= applyAdmissionQuotaSafety( + await requestFallbackSelection(), + false, + ) const fallbackResult = await tryFallbackAccounts( requestInput, init, response, - undefined, + reqStorage, + fallbackCandidates, ) const fallbackResponse = fallbackResult.response if (fallbackResponse !== response) { diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 83a25a1..b35bf92 100644 --- a/packages/opencode/src/sidebar-state.ts +++ b/packages/opencode/src/sidebar-state.ts @@ -1,11 +1,13 @@ export interface QuotaWindow { usedPercent: number remainingPercent: number + checkedAt?: number resetsAt?: string windowMinutes?: number } export interface AccountQuota { + checkedAt?: number primary?: QuotaWindow secondary?: QuotaWindow resetCreditsAvailable?: number @@ -80,6 +82,8 @@ export function getPresentQuotaWindows( export interface SidebarAccountState { id: string label: string | undefined + /** ChatGPT identity of the account this quota belongs to. */ + accountId?: string quota: AccountQuota | null killed: boolean enabled: boolean @@ -97,6 +101,8 @@ export type ActiveRoutingMap = Record export interface SidebarState { main: { quota: AccountQuota | null + /** ChatGPT identity of the main account this quota belongs to. */ + mainAccountId?: string killed: boolean quotaBackedOff?: boolean quotaBackoffUntil?: number @@ -212,6 +218,9 @@ export function normalizeSidebarState(raw: unknown): SidebarState { main = { quota: ('quota' in m ? m.quota : null) as AccountQuota | null, killed: typeof m.killed === 'boolean' ? m.killed : false, + ...(typeof m.mainAccountId === 'string' + ? { mainAccountId: m.mainAccountId } + : {}), // Preserve optional backoff fields if present ...(typeof m.quotaBackedOff === 'boolean' ? { quotaBackedOff: m.quotaBackedOff } @@ -247,6 +256,9 @@ export function normalizeSidebarState(raw: unknown): SidebarState { .map((e) => ({ id: e.id as string, label: typeof e.label === 'string' ? e.label : undefined, + ...(typeof e.accountId === 'string' + ? { accountId: e.accountId } + : {}), quota: ('quota' in e ? e.quota : null) as AccountQuota | null, killed: typeof e.killed === 'boolean' ? e.killed : false, enabled: typeof e.enabled === 'boolean' ? e.enabled : true, @@ -280,9 +292,11 @@ export function normalizeSidebarState(raw: unknown): SidebarState { } } -export async function getSidebarState(): Promise { +export async function getSidebarState( + stateFile = getSidebarStateFile(), +): Promise { try { - const raw = await readFile(getSidebarStateFile(), 'utf8') + const raw = await readFile(stateFile, 'utf8') return normalizeSidebarState(JSON.parse(raw)) } catch { return DEFAULT_SIDEBAR_STATE @@ -317,20 +331,40 @@ export function isUsableRoutingEntry( ) } +// Earliest future reset among the quota's exhausted windows, or undefined when +// no present window is exhausted. Every present window is evaluated — matching +// the admission policy, which rejects an account when ANY live window is below +// its threshold — and each check fails open: a missing/malformed usage, a +// missing/unparsable reset, or a reset already in the past never counts as +// exhausted, so a stale or corrupt snapshot can never block routing. +export function exhaustedQuotaResetAt( + quota: AccountQuota | null | undefined, + now = Date.now(), +): { resetsAt: string; resetAtMs: number } | undefined { + let earliest: { resetsAt: string; resetAtMs: number } | undefined + for (const key of QUOTA_WINDOW_KEYS) { + const window = quota?.[key] + if ( + typeof window?.resetsAt !== 'string' || + !Number.isFinite(window.usedPercent) || + window.usedPercent < 100 + ) { + continue + } + const resetAtMs = Date.parse(window.resetsAt) + if (!Number.isFinite(resetAtMs) || resetAtMs <= now) continue + if (!earliest || resetAtMs < earliest.resetAtMs) { + earliest = { resetsAt: window.resetsAt, resetAtMs } + } + } + return earliest +} + export function isQuotaExhausted( quota: AccountQuota | null | undefined, now = Date.now(), ): boolean { - const primary = quota?.primary - if ( - typeof primary?.resetsAt !== 'string' || - !Number.isFinite(primary.usedPercent) || - primary.usedPercent < 100 - ) { - return false - } - const resetsAt = Date.parse(primary.resetsAt) - return Number.isFinite(resetsAt) && resetsAt > now + return exhaustedQuotaResetAt(quota, now) !== undefined } export function resolveSessionSidebarRouting( @@ -496,6 +530,30 @@ export type SidebarMachineState = Pick< 'main' | 'fallbacks' | 'planType' | 'credits' | 'lastUpdated' > & { route: string } +function primaryQuotaCheckedAt(quota: AccountQuota | null): number | undefined { + for (const checkedAt of [quota?.primary?.checkedAt, quota?.checkedAt]) { + if (typeof checkedAt === 'number' && Number.isFinite(checkedAt)) { + return checkedAt + } + } + return undefined +} + +function freshestQuota( + incoming: AccountQuota | null, + existing: AccountQuota | null, +): AccountQuota | null { + const incomingCheckedAt = primaryQuotaCheckedAt(incoming) + const existingCheckedAt = primaryQuotaCheckedAt(existing) + if ( + existingCheckedAt !== undefined && + (incomingCheckedAt === undefined || existingCheckedAt > incomingCheckedAt) + ) { + return existing + } + return incoming +} + export function setSidebarMachineState( machineState: SidebarMachineState, file = getSidebarStateFile(), @@ -504,13 +562,49 @@ export function setSidebarMachineState( return enqueueSidebarWrite(async () => { await writeMergedSidebarState( file, - (latest) => ({ - ...latest, - ...machineState, - activeId: latest.activeId, - activeRouting: latest.activeRouting, - lastUpdated: Math.max(Date.now(), latest.lastUpdated + 1), - }), + (latest) => { + const latestFallbacks = new Map( + latest.fallbacks.map((account) => [account.id, account]), + ) + const mergedMainQuota = freshestQuota( + machineState.main.quota, + latest.main.quota, + ) + return { + ...latest, + ...machineState, + main: { + ...machineState.main, + quota: mergedMainQuota, + // Identity follows the snapshot that wins the freshness merge, so a + // reader never pairs one account's id with another account's quota + // (a re-login race would otherwise resurrect the stale-account bug). + mainAccountId: + mergedMainQuota === latest.main.quota && + mergedMainQuota !== machineState.main.quota + ? latest.main.mainAccountId + : machineState.main.mainAccountId, + }, + fallbacks: machineState.fallbacks.map((account) => { + const existing = latestFallbacks.get(account.id) + const mergedQuota = freshestQuota( + account.quota, + existing?.quota ?? null, + ) + return { + ...account, + quota: mergedQuota, + accountId: + mergedQuota === existing?.quota && mergedQuota !== account.quota + ? existing?.accountId + : account.accountId, + } + }), + activeId: latest.activeId, + activeRouting: latest.activeRouting, + lastUpdated: Math.max(Date.now(), latest.lastUpdated + 1), + } + }, hooks, ) }) diff --git a/packages/opencode/src/tests/integration.test.ts b/packages/opencode/src/tests/integration.test.ts index 92bd4d7..d25ab5e 100644 --- a/packages/opencode/src/tests/integration.test.ts +++ b/packages/opencode/src/tests/integration.test.ts @@ -2570,6 +2570,7 @@ describe('integration: active fallback routing', () => { mainExpires: number, experimentalWebSockets = false, responsesLite = false, + mainAccountId?: string, ) { const hooks = await CodexAuthPlugin(input, { experimentalWebSockets, @@ -2584,6 +2585,7 @@ describe('integration: active fallback routing', () => { access: 'main-stale-token', refresh: 'main-refresh-token', expires: mainExpires, + ...(mainAccountId ? { accountId: mainAccountId } : {}), }), { id: 'openai', @@ -2782,6 +2784,89 @@ describe('integration: active fallback routing', () => { ) } + function seedAdmissionAccounts( + ids: string[], + mode: 'main-first' | 'fallback-first' = 'fallback-first', + quotas: Record = {}, + ) { + writeFileSync( + configFile, + JSON.stringify({ + version: 1, + main: { type: 'opencode', provider: 'openai' }, + accounts: ids.map((id) => ({ + id, + type: 'oauth', + label: id, + enabled: true, + access: `${id}-token`, + refresh: `${id}-refresh`, + expires: Date.now() + 24 * 3600_000, + accountId: `chatgpt-${id}`, + ...(quotas[id] ? { quota: quotas[id] } : {}), + })), + refresh: { refreshBeforeExpiryMinutes: 5 }, + routing: { mode }, + }), + ) + } + + function admissionQuota( + usedPercent: number, + resetsAt: string, + checkedAt: number, + ): NonNullable { + return { + primary: { + usedPercent, + remainingPercent: 100 - usedPercent, + resetsAt, + checkedAt, + windowMinutes: 10_080, + }, + } + } + + function writeAdmissionSidebarState(input: { + fallbackIds: string[] + fallbackQuotas?: Record + fallbackAccountIds?: Record + mainQuota?: SidebarState['main']['quota'] + mainAccountId?: string + activeId?: string + route?: 'main-first' | 'fallback-first' + }) { + const state: SidebarState = { + main: { quota: input.mainQuota ?? null, killed: false }, + fallbacks: input.fallbackIds.map((id) => ({ + id, + label: id, + ...(input.fallbackAccountIds?.[id] !== undefined + ? { accountId: input.fallbackAccountIds[id] } + : {}), + quota: input.fallbackQuotas?.[id] ?? null, + killed: false, + enabled: true, + })), + activeId: input.activeId, + route: input.route ?? 'fallback-first', + lastUpdated: Date.now(), + } + if (input.mainAccountId !== undefined) { + state.main.mainAccountId = input.mainAccountId + } + writeFileSync(sidebarFile, JSON.stringify(state)) + } + + function mockAdmissionFetch(seenAuth: string[], status = 200) { + return (async (url: unknown, init?: unknown) => { + if (String(url).includes('responses')) { + seenAuth.push(headerValue(init, 'authorization')) + } + return new Response('{}', { status }) + }) as unknown as typeof globalThis.fetch + } + test.each([ [ { @@ -3519,6 +3604,525 @@ describe('integration: active fallback routing', () => { } }) + it('admission quota skips an exhausted first fallback from the shared sidebar state', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt', 'client-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt', 'client-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now), + 'client-alt': admissionQuota(20, reset, now), + }, + activeId: 'work-alt', + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer client-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota skips a file-exhausted fallback with an empty process quota cache', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now), + }, + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer main-stale-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota uses fresher healthy memory instead of a stale exhausted file row', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt'], 'fallback-first', { + 'work-alt': admissionQuota(20, reset, now), + }) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now - 60_000), + }, + }) + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(seenAuth).toEqual(['Bearer work-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota uses a fresher exhausted file row instead of stale healthy memory', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt', 'client-alt'], 'fallback-first', { + 'work-alt': admissionQuota(20, reset, now - 60_000), + }) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt', 'client-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now), + 'client-alt': admissionQuota(20, reset, now), + }, + }) + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(seenAuth).toEqual(['Bearer client-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota skips from fresher exhausted memory instead of a stale healthy file row', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt', 'client-alt'], 'fallback-first', { + 'work-alt': admissionQuota(100, reset, now), + }) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt', 'client-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(20, reset, now - 60_000), + 'client-alt': admissionQuota(20, reset, now), + }, + }) + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(seenAuth).toEqual(['Bearer client-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + test.each([ + ['missing quota', null], + ['missing reset', { primary: { usedPercent: 100, remainingPercent: 0 } }], + [ + 'malformed reset', + { + primary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: 'not-a-date', + }, + }, + ], + [ + 'malformed usage', + { + primary: { + usedPercent: '100', + remainingPercent: 0, + resetsAt: new Date(Date.now() + 3600_000).toISOString(), + }, + }, + ], + ])('admission quota retains a fallback with %s', async (_label, quota) => { + seedAdmissionAccounts(['work-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + Date.now() + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + 'work-alt': quota as SidebarState['main']['quota'], + }, + }) + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(seenAuth).toEqual(['Bearer work-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota preserves probe order when every account is exhausted', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt', 'client-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth, 429) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt', 'client-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now), + 'client-alt': admissionQuota(100, reset, now), + }, + mainQuota: admissionQuota(100, reset, now), + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(429) + expect(seenAuth).toEqual([ + 'Bearer work-alt-token', + 'Bearer client-alt-token', + 'Bearer main-stale-token', + ]) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota reroutes a file-exhausted main without probing it', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['client-alt'], 'main-first') + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + await new Promise((resolve) => setTimeout(resolve, 2)) + const checkedAt = Date.now() + writeAdmissionSidebarState({ + fallbackIds: ['client-alt'], + fallbackQuotas: { + 'client-alt': admissionQuota(20, reset, checkedAt), + }, + mainQuota: admissionQuota(100, reset, checkedAt), + route: 'main-first', + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer client-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota ignores an exhausted main row from a different account', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['client-alt'], 'main-first') + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + false, + false, + 'new-account', + ) + hooks = loaded.hooks + await new Promise((resolve) => setTimeout(resolve, 2)) + const checkedAt = Date.now() + writeAdmissionSidebarState({ + fallbackIds: ['client-alt'], + fallbackQuotas: { + 'client-alt': admissionQuota(20, reset, checkedAt), + }, + mainQuota: admissionQuota(100, reset, checkedAt), + mainAccountId: 'old-account', + route: 'main-first', + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer main-stale-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota ignores an exhausted fallback row stamped with a different account identity', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + // Live account identity is chatgpt-work-alt (see seedAdmissionAccounts). + seedAdmissionAccounts(['work-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now), + }, + // The file row belongs to a previous login of this stable id; the live + // account is a different ChatGPT identity, so the exhausted row must be + // treated as absent (fail-open) rather than blocking the replacement. + fallbackAccountIds: { 'work-alt': 'chatgpt-stale' }, + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer work-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota honors an exhausted fallback row matching the live account identity', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota(100, reset, now), + }, + // Identity matches the live account, so the exhausted row is honored + // and the fallback is skipped in favor of main. + fallbackAccountIds: { 'work-alt': 'chatgpt-work-alt' }, + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer main-stale-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota skips a fallback exhausted only on its secondary window', async () => { + const now = Date.now() + const reset = new Date(now + 7 * 24 * 3600_000).toISOString() + seedAdmissionAccounts(['work-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + // No primary window; the secondary window alone is exhausted. + 'work-alt': { + secondary: { + usedPercent: 100, + remainingPercent: 0, + resetsAt: reset, + checkedAt: now, + windowMinutes: 10_080, + }, + }, + }, + }) + + const response = await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(response.status).toBe(200) + expect(seenAuth).toEqual(['Bearer main-stale-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + + it('admission quota retains an exhausted-looking fallback after its reset passes', async () => { + const now = Date.now() + seedAdmissionAccounts(['work-alt']) + const seenAuth: string[] = [] + const originalFetch = globalThis.fetch + globalThis.fetch = mockAdmissionFetch(seenAuth) + + let hooks: Hooks | undefined + try { + const loaded = await loadFetchOverride( + createMockPluginInput(), + now + 3600_000, + ) + hooks = loaded.hooks + writeAdmissionSidebarState({ + fallbackIds: ['work-alt'], + fallbackQuotas: { + 'work-alt': admissionQuota( + 100, + new Date(now - 60_000).toISOString(), + now, + ), + }, + }) + + await loaded.fetchOverride( + 'https://api.openai.com/v1/responses', + requestInit(), + ) + + expect(seenAuth).toEqual(['Bearer work-alt-token']) + } finally { + globalThis.fetch = originalFetch + await hooks?.dispose?.() + } + }) + it('refreshes an expired active fallback without writing the auth slot', async () => { seedStorage({ access: 'fallback-stale-token', diff --git a/packages/opencode/src/tests/quota-push.test.ts b/packages/opencode/src/tests/quota-push.test.ts index 602a5b0..afd86ad 100644 --- a/packages/opencode/src/tests/quota-push.test.ts +++ b/packages/opencode/src/tests/quota-push.test.ts @@ -1,6 +1,10 @@ import { afterEach, beforeEach, describe, expect, it } from 'bun:test' import { randomUUID } from 'node:crypto' -import type { AccountStorage, OAuthQuotaSnapshot } from '../core/accounts.ts' +import type { + AccountStorage, + OAuthAccount, + OAuthQuotaSnapshot, +} from '../core/accounts.ts' import type { QuotaManager } from '../core/quota-manager.ts' import { buildSidebarMachineState, @@ -160,6 +164,72 @@ describe('QuotaManager push', () => { expect(qm.getFallback('fb-1', newToken)).toBeNull() }) + it('peekFallbackForPolicy drops a quota bound to a different identity (re-login on a stable id)', async () => { + const { QuotaManager } = await import('../core/quota-manager.ts') + const snapshot = goodSnapshot() + + const qm = new QuotaManager({ + storage: null, + fetchQuotaFn: () => { + throw new Error('must not be called') + }, + }) + + // Cached under the OLD ChatGPT identity on stable id fb-1. + qm.setFallback( + 'fb-1', + { + quota: snapshot, + refreshAfter: Date.now() + 60_000, + checkedAt: Date.now(), + }, + `fb-old-${randomUUID()}`, + false, + 'old', + ) + + // Same identity (or no identity supplied) still sees the cached quota. + expect(qm.peekFallbackForPolicy('fb-1', 'old')).not.toBeNull() + expect(qm.peekFallbackForPolicy('fb-1')).not.toBeNull() + // A re-login (different ChatGPT identity on the same stable id) drops the + // policy view, so the old identity's quota never blocks its replacement. + expect(qm.peekFallbackForPolicy('fb-1', 'new')).toBeNull() + }) + + it('seedFallbacksFromAccounts binds persisted quota to the account identity', async () => { + const { QuotaManager } = await import('../core/quota-manager.ts') + const now = Date.now() + + const qm = new QuotaManager({ + storage: null, + fetchQuotaFn: () => { + throw new Error('must not be called') + }, + }) + + const account: OAuthAccount = { + id: 'fb-1', + type: 'oauth', + access: `fb-seed-${randomUUID()}`, + refresh: 'fb-seed-refresh', + expires: now + 3600_000, + enabled: true, + accountId: 'old', + quota: { + primary: { + usedPercent: 100, + remainingPercent: 0, + checkedAt: now, + resetsAt: new Date(now + 60_000).toISOString(), + }, + }, + } + qm.seedFallbacksFromAccounts([account]) + + expect(qm.peekFallbackForPolicy('fb-1', 'old')).not.toBeNull() + expect(qm.peekFallbackForPolicy('fb-1', 'new')).toBeNull() + }) + it('conditional push: empty snapshot does NOT overwrite a valid cached one', async () => { const { QuotaManager } = await import('../core/quota-manager.ts') const token = `access-${randomUUID()}` @@ -406,6 +476,7 @@ describe('QuotaManager push', () => { const store: AccountStorage = { version: 1, main: { type: 'opencode', provider: 'openai' }, + mainAccountId: 'chatgpt-main', accounts: [ { id: 'fallback-1', @@ -422,6 +493,7 @@ describe('QuotaManager push', () => { const expectedMachine = { main: { quota: { + checkedAt: 1, primary: { usedPercent: 20, remainingPercent: 80, @@ -430,6 +502,7 @@ describe('QuotaManager push', () => { }, resetCreditsAvailable: 4, }, + mainAccountId: 'chatgpt-main', killed: false, resetCredits: 4, }, @@ -438,6 +511,7 @@ describe('QuotaManager push', () => { id: 'fallback-1', label: undefined, quota: { + checkedAt: 1, primary: { usedPercent: 30, remainingPercent: 70, diff --git a/packages/opencode/src/tests/sidebar-state.test.ts b/packages/opencode/src/tests/sidebar-state.test.ts index 783b985..9576ca9 100644 --- a/packages/opencode/src/tests/sidebar-state.test.ts +++ b/packages/opencode/src/tests/sidebar-state.test.ts @@ -10,11 +10,13 @@ import { computeQuotaPacing, DEFAULT_SIDEBAR_STATE, drainSidebarWrites, + exhaustedQuotaResetAt, formatWindowLabel, getCollapsedQuotaSummary, getPresentQuotaWindows, getSidebarState, getSidebarStateFile, + isQuotaExhausted, isUsableRoutingEntry, normalizeSidebarState, pruneActiveRouting, @@ -35,8 +37,12 @@ import { FLOOR_SIDEBAR_STATE_FILE } from './setup-env.ts' const FIVE_HOUR_MS = 5 * 60 * 60 * 1000 const SEVEN_DAY_MS = 7 * 24 * 60 * 60 * 1000 -const quota = (used: number): AccountQuota => ({ - primary: { usedPercent: used, remainingPercent: 100 - used }, +const quota = (used: number, checkedAt?: number): AccountQuota => ({ + primary: { + usedPercent: used, + remainingPercent: 100 - used, + ...(checkedAt === undefined ? {} : { checkedAt }), + }, secondary: { usedPercent: used, remainingPercent: 100 - used }, }) @@ -97,6 +103,30 @@ describe('normalizeSidebarState', () => { expect(result.main.killed).toBe(false) }) + test('main account identity is preserved only when it is a string', () => { + const valid = normalizeSidebarState({ + main: { quota: null, mainAccountId: 'chatgpt-main' }, + }) + const malformed = normalizeSidebarState({ + main: { quota: null, mainAccountId: 42 }, + }) + + expect(valid.main.mainAccountId).toBe('chatgpt-main') + expect(malformed.main.mainAccountId).toBeUndefined() + }) + + test('fallback account identity is preserved only when it is a string', () => { + const result = normalizeSidebarState({ + fallbacks: [ + { id: 'fb1', accountId: 'chatgpt-fb1' }, + { id: 'fb2', accountId: 42 }, + ], + }) + + expect(result.fallbacks[0]?.accountId).toBe('chatgpt-fb1') + expect(result.fallbacks[1]?.accountId).toBeUndefined() + }) + // (d) fallbacks is a non-array value test('{"fallbacks":"notarray"} — fallbacks coerced to []', () => { const result = normalizeSidebarState({ fallbacks: 'notarray' }) @@ -132,6 +162,7 @@ describe('normalizeSidebarState', () => { }, secondary: { usedPercent: 17, remainingPercent: 83 }, }, + mainAccountId: 'chatgpt-main', killed: true, quotaBackedOff: true, quotaBackoffUntil: 1234567890, @@ -143,6 +174,7 @@ describe('normalizeSidebarState', () => { { id: 'fb1', label: 'work', + accountId: 'chatgpt-fb1', quota: { primary: { usedPercent: 5, remainingPercent: 95 } }, killed: false, enabled: true, @@ -172,10 +204,12 @@ describe('normalizeSidebarState', () => { expect(result.main.refreshBackedOff).toBe(false) expect(result.main.refreshBackoffUntil).toBe(9876543210) expect(result.main.resetCredits).toBe(4) + expect(result.main.mainAccountId).toBe('chatgpt-main') expect(result.fallbacks).toHaveLength(1) const fb0 = result.fallbacks[0]! expect(fb0.id).toBe('fb1') expect(fb0.label).toBe('work') + expect(fb0.accountId).toBe('chatgpt-fb1') expect(fb0.resetCredits).toBe(2) expect(result.activeId).toBe('fb1') expect(result.route).toBe('fallback') @@ -1324,6 +1358,260 @@ test('machine writes preserve routing while retaining reset-credit fields', asyn expect(written.activeRouting?.['sess-a']?.activeId).toBe('fallback-1') }) +test('machine writes cannot clobber fresher main and fallback quota from disk', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-quota-fresh-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + const stale = now - 10 * 60_000 + await setSidebarState( + make({ + main: main(quota(10, now)), + fallbacks: [fb({ id: 'fallback-1', quota: quota(20, now) })], + activeRouting: { + session: { + activeId: 'fallback-1', + route: 'fallback-first', + updatedAt: now, + }, + }, + }), + file, + ) + + await setSidebarMachineState( + { + main: main(quota(90, stale)), + fallbacks: [fb({ id: 'fallback-1', quota: quota(80, stale) })], + route: 'main-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + + const written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.primary).toMatchObject({ + usedPercent: 10, + checkedAt: now, + }) + expect(written.fallbacks[0]?.quota?.primary).toMatchObject({ + usedPercent: 20, + checkedAt: now, + }) + expect(written.route).toBe('main-first') + expect(written.activeRouting?.session?.activeId).toBe('fallback-1') +}) + +test('machine write keeps the existing identity when the existing quota wins the merge (re-login race)', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-identity-keep-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + // On disk: the OLD account's quota with a fresh checkedAt. + await setSidebarState( + make({ + main: { ...main(quota(10, now)), mainAccountId: 'account-old' }, + fallbacks: [ + fb({ id: 'fallback-1', accountId: 'fb-old', quota: quota(20, now) }), + ], + }), + file, + ) + + // Incoming: the re-logged-in process has no quota yet (null) but a NEW + // identity. The existing quota is fresher so it wins the merge — the + // identity must follow it rather than be overwritten with the new + // account's id, or a reader would judge the new account by the old + // account's quota. + await setSidebarMachineState( + { + main: { ...main(null), mainAccountId: 'account-new' }, + fallbacks: [fb({ id: 'fallback-1', accountId: 'fb-new', quota: null })], + route: 'main-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + + const written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.primary?.usedPercent).toBe(10) + expect(written.main.mainAccountId).toBe('account-old') + expect(written.fallbacks[0]?.quota?.primary?.usedPercent).toBe(20) + expect(written.fallbacks[0]?.accountId).toBe('fb-old') +}) + +test('machine write carries the incoming identity when the incoming quota wins the merge', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-identity-take-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + const stale = now - 10 * 60_000 + // On disk: stale quota under the OLD identity. + await setSidebarState( + make({ + main: { ...main(quota(10, stale)), mainAccountId: 'account-old' }, + }), + file, + ) + + // Incoming: a fresher snapshot under the NEW identity wins the merge, and + // the identity follows it. + await setSidebarMachineState( + { + main: { ...main(quota(50, now)), mainAccountId: 'account-new' }, + fallbacks: [], + route: 'main-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + + const written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.primary?.usedPercent).toBe(50) + expect(written.main.mainAccountId).toBe('account-new') +}) + +test('concurrent machine writes under different identities never pair an identity with the wrong quota', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-identity-race-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + const stale = now - 10 * 60_000 + // Two writers race an account switch: whichever quota wins the freshness + // merge, its OWN identity must win with it — the new identity can never + // be paired with the old quota. + await setSidebarState( + make({ + main: { ...main(quota(10, stale)), mainAccountId: 'account-a' }, + }), + file, + ) + + // Writer B lands a FRESHER snapshot — B's quota and B's identity win. + await setSidebarMachineState( + { + main: { ...main(quota(60, now)), mainAccountId: 'account-b' }, + fallbacks: [], + route: 'main-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + let written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.primary?.usedPercent).toBe(60) + expect(written.main.mainAccountId).toBe('account-b') + + // A stale write from A arrives afterwards — it cannot resurrect A's quota + // over B's fresher snapshot, and cannot attach A's identity to B's quota. + await setSidebarMachineState( + { + main: { ...main(quota(10, stale)), mainAccountId: 'account-a' }, + fallbacks: [], + route: 'main-first', + lastUpdated: now + 2, + }, + file, + ) + await drainSidebarWrites() + written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.primary?.usedPercent).toBe(60) + expect(written.main.mainAccountId).toBe('account-b') +}) + +test('machine writes select quota freshness independently per account', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-quota-mixed-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + const stale = now - 10 * 60_000 + await setSidebarState( + make({ + fallbacks: [ + fb({ id: 'fallback-a', quota: quota(10, stale) }), + fb({ id: 'fallback-b', quota: quota(20, now) }), + ], + }), + file, + ) + + await setSidebarMachineState( + { + main: main(null), + fallbacks: [ + fb({ id: 'fallback-a', quota: quota(30, now) }), + fb({ id: 'fallback-b', quota: quota(40, stale) }), + ], + route: 'fallback-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + + const written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect( + written.fallbacks.find((row) => row.id === 'fallback-a')?.quota?.primary + ?.usedPercent, + ).toBe(30) + expect( + written.fallbacks.find((row) => row.id === 'fallback-b')?.quota?.primary + ?.usedPercent, + ).toBe(20) +}) + +test('valid checkedAt beats missing or invalid values while incoming wins ties without timestamps', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-quota-invalid-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + const invalidQuota = { + primary: { + usedPercent: 90, + remainingPercent: 10, + checkedAt: 'not-a-number', + }, + } as unknown as AccountQuota + await setSidebarState( + make({ + main: main(quota(10, now)), + fallbacks: [ + fb({ id: 'incoming-valid', quota: invalidQuota }), + fb({ id: 'disk-valid', quota: quota(20, now) }), + fb({ id: 'both-missing', quota: quota(30) }), + ], + }), + file, + ) + + await setSidebarMachineState( + { + main: main(null), + fallbacks: [ + fb({ id: 'incoming-valid', quota: quota(40, now) }), + fb({ id: 'disk-valid', quota: invalidQuota }), + fb({ id: 'both-missing', quota: quota(50) }), + ], + route: 'main-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + + const written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.primary?.usedPercent).toBe(10) + expect( + written.fallbacks.find((row) => row.id === 'incoming-valid')?.quota?.primary + ?.usedPercent, + ).toBe(40) + expect( + written.fallbacks.find((row) => row.id === 'disk-valid')?.quota?.primary + ?.usedPercent, + ).toBe(20) + expect( + written.fallbacks.find((row) => row.id === 'both-missing')?.quota?.primary + ?.usedPercent, + ).toBe(50) +}) + test('headerless request compatibility write does not create a session entry', async () => { const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-legacy-')) const file = join(tempDir, 'sidebar-state.json') @@ -1364,3 +1652,84 @@ test('machine refresh cannot replace request-authored legacy active id', async ( expect(written.route).toBe('main-first') expect(written.main.resetCredits).toBe(4) }) + +describe('isQuotaExhausted / exhaustedQuotaResetAt', () => { + const now = 1_750_000_000_000 + const future = new Date(now + 3600_000).toISOString() + const laterFuture = new Date(now + 7200_000).toISOString() + const past = new Date(now - 3600_000).toISOString() + + const windowAt = ( + usedPercent: number, + resetsAt?: string, + ): AccountQuota['primary'] => ({ + usedPercent, + remainingPercent: 100 - usedPercent, + ...(resetsAt === undefined ? {} : { resetsAt }), + }) + + test('null, undefined, and empty quotas are never exhausted', () => { + expect(isQuotaExhausted(null, now)).toBe(false) + expect(isQuotaExhausted(undefined, now)).toBe(false) + expect(isQuotaExhausted({}, now)).toBe(false) + expect(exhaustedQuotaResetAt(null, now)).toBeUndefined() + }) + + test('a primary window at 100% with a future reset is exhausted', () => { + const quota: AccountQuota = { primary: windowAt(100, future) } + expect(isQuotaExhausted(quota, now)).toBe(true) + expect(exhaustedQuotaResetAt(quota, now)).toEqual({ + resetsAt: future, + resetAtMs: Date.parse(future), + }) + }) + + test('an exhausted SECONDARY window alone exhausts the account', () => { + const quota: AccountQuota = { + primary: windowAt(20, future), + secondary: windowAt(100, laterFuture), + } + expect(isQuotaExhausted(quota, now)).toBe(true) + expect(exhaustedQuotaResetAt(quota, now)).toEqual({ + resetsAt: laterFuture, + resetAtMs: Date.parse(laterFuture), + }) + }) + + test('with both windows exhausted, the earliest future reset wins', () => { + const quota: AccountQuota = { + primary: windowAt(100, laterFuture), + secondary: windowAt(100, future), + } + expect(exhaustedQuotaResetAt(quota, now)).toEqual({ + resetsAt: future, + resetAtMs: Date.parse(future), + }) + }) + + test.each([ + ['usage below 100%', { primary: windowAt(99, future) }], + ['missing reset', { primary: windowAt(100) }], + ['malformed reset', { primary: windowAt(100, 'not-a-date') }], + ['reset already past', { primary: windowAt(100, past) }], + [ + 'malformed usage', + { + primary: { + usedPercent: Number.NaN, + remainingPercent: 0, + resetsAt: future, + }, + }, + ], + ])('fails open on %s', (_label, quota) => { + expect(isQuotaExhausted(quota, now)).toBe(false) + expect(exhaustedQuotaResetAt(quota, now)).toBeUndefined() + }) + + test('an absent window is not applicable, not unknown', () => { + // Only secondary present and healthy — no primary to judge. + const quota: AccountQuota = { secondary: windowAt(30, future) } + expect(isQuotaExhausted(quota, now)).toBe(false) + }) +})