diff --git a/packages/opencode/src/core/background-quota-refresh.ts b/packages/opencode/src/core/background-quota-refresh.ts new file mode 100644 index 0000000..f5be7ce --- /dev/null +++ b/packages/opencode/src/core/background-quota-refresh.ts @@ -0,0 +1,140 @@ +import { + type RefreshAllQuotaDeps, + type RefreshAllQuotaResult, + refreshAllQuota, +} from './refresh-all-quota' +import { acquireRefreshFileLock } from './refresh-file-lock' + +export const BACKGROUND_QUOTA_REFRESH_INTERVAL_MS = 5 * 60_000 +export const BACKGROUND_QUOTA_REFRESH_JITTER_MS = 30_000 +export const BACKGROUND_QUOTA_FRESHNESS_MS = 4 * 60_000 +export const BACKGROUND_QUOTA_REFRESH_LOCK_NAME = 'bg-quota-refresh' +// The lock must outlive the worst-case serial refresh: each fallback account +// costs a token refresh plus a wham fetch (up to ~15s each), so a fleet of +// accounts can run well past a minute. A TTL that expires mid-run lets a second +// process start a concurrent refresh — 120s covers a multi-account pass with +// margin while still releasing a crashed owner's lock within one poll interval. +export const BACKGROUND_QUOTA_REFRESH_LOCK_TTL_MS = 120_000 + +type TimerHandle = ReturnType + +interface BackgroundQuotaRefreshOptions { + setIntervalFn?: (callback: () => void, intervalMs: number) => TimerHandle + clearIntervalFn?: (timer: TimerHandle) => void + random?: () => number + onError?: (error: unknown) => void +} + +type RefreshAllQuotaFn = ( + deps: RefreshAllQuotaDeps, +) => Promise + +type BackgroundLockHandle = { release: () => Promise } +export type BackgroundLockAcquirer = () => Promise + +export function refreshQuotaInBackground( + deps: RefreshAllQuotaDeps, + refreshFn: RefreshAllQuotaFn = refreshAllQuota, + acquireLock: BackgroundLockAcquirer = () => + acquireRefreshFileLock({ + name: BACKGROUND_QUOTA_REFRESH_LOCK_NAME, + ttlMs: BACKGROUND_QUOTA_REFRESH_LOCK_TTL_MS, + path: deps.configPath, + }), +): Promise { + // The lock is advisory and only serializes the fetch across processes. A + // clean null means another live owner is already refreshing, so this tick + // skips — that process will update the shared sidebar file. A lock-mechanism + // failure fails open and refreshes anyway, matching the pre-lock behavior + // rather than stranding quota freshness on a broken lock. + return claimBackgroundRefresh(deps, refreshFn, acquireLock) +} + +async function claimBackgroundRefresh( + deps: RefreshAllQuotaDeps, + refreshFn: RefreshAllQuotaFn, + acquireLock: BackgroundLockAcquirer, +): Promise { + let lock: BackgroundLockHandle | null | undefined + try { + lock = await acquireLock() + } catch { + lock = undefined + } + if (lock === null) return [] + try { + return await refreshFn({ + ...deps, + respectBackoff: true, + skipFresherThanMs: BACKGROUND_QUOTA_FRESHNESS_MS, + }) + } finally { + await lock?.release().catch(() => {}) + } +} + +export class BackgroundQuotaRefresh { + private readonly setIntervalFn: NonNullable< + BackgroundQuotaRefreshOptions['setIntervalFn'] + > + private readonly clearIntervalFn: NonNullable< + BackgroundQuotaRefreshOptions['clearIntervalFn'] + > + private readonly random: () => number + private onError: ((error: unknown) => void) | undefined + private run: (() => Promise) | undefined + private timer: TimerHandle | undefined + private tickPromise: Promise | undefined + private stopped = false + + constructor(options: BackgroundQuotaRefreshOptions = {}) { + this.setIntervalFn = options.setIntervalFn ?? setInterval + this.clearIntervalFn = options.clearIntervalFn ?? clearInterval + this.random = options.random ?? Math.random + this.onError = options.onError + } + + start( + run: () => Promise, + onError: ((error: unknown) => void) | undefined = this.onError, + ): void { + this.run = run + this.onError = onError + this.stopped = false + if (this.timer) return + + const jitter = Math.round( + (this.random() * 2 - 1) * BACKGROUND_QUOTA_REFRESH_JITTER_MS, + ) + this.timer = this.setIntervalFn(() => { + const currentRun = this.run + if (this.stopped || !currentRun || this.tickPromise) return + const tickPromise = currentRun() + .catch((error) => { + try { + this.onError?.(error) + } catch {} + }) + .finally(() => { + if (this.tickPromise === tickPromise) this.tickPromise = undefined + }) + this.tickPromise = tickPromise + }, BACKGROUND_QUOTA_REFRESH_INTERVAL_MS + jitter) + if ('unref' in this.timer) this.timer.unref() + } + + // Non-blocking stop: clears the timer and bars further ticks, but lets an + // in-flight run finish naturally (it checks isStopped() before committing + // sidebar state, so a stopping poller never writes a stale snapshot). + stop(): void { + this.stopped = true + this.run = undefined + if (!this.timer) return + this.clearIntervalFn(this.timer) + this.timer = undefined + } + + isStopped(): boolean { + return this.stopped + } +} diff --git a/packages/opencode/src/core/refresh-all-quota.ts b/packages/opencode/src/core/refresh-all-quota.ts index c483dd7..6b32aea 100644 --- a/packages/opencode/src/core/refresh-all-quota.ts +++ b/packages/opencode/src/core/refresh-all-quota.ts @@ -1,4 +1,5 @@ import { createLogger } from '../logger' +import { getSidebarState, type SidebarState } from '../sidebar-state' import { errorMessage } from '../util/error' import type { FallbackAccountManager, @@ -28,6 +29,11 @@ export interface RefreshAllQuotaDeps { refresh: string expires: number }> + refreshMainWithLease: () => Promise<{ + access: string + refresh: string + expires: number + }> fallbackManager: FallbackAccountManager quotaManager: QuotaManager loadAccounts: typeof loadAccounts @@ -56,6 +62,12 @@ export interface RefreshAllQuotaDeps { whamFn?: typeof whamUsageFn respectBackoff?: boolean logger?: QuotaLogger + skipFresherThanMs?: number + readSidebarState?: () => Promise +} + +export interface RefreshAllQuotaOptions { + accountKey?: string } export interface RefreshAllQuotaResult { @@ -66,6 +78,7 @@ export interface RefreshAllQuotaResult { export async function refreshAllQuota( deps: RefreshAllQuotaDeps, + options: RefreshAllQuotaOptions = {}, ): Promise { const whamFn = deps.whamFn if (!whamFn) throw new Error('whamFn is required for refreshAllQuota') @@ -84,80 +97,152 @@ export async function refreshAllQuota( else logger.warn('quota refresh failed', payload) } - // --- MAIN --- - try { - let auth = await deps.getAuth() - if (auth.type === 'oauth') { - if (!auth.access || (auth.expires ?? 0) < deps.now()) { - const tokens = await deps.codexRefreshFn({ - refreshToken: auth.refresh ?? '', - fetchImpl: deps.fetchImpl, - now: deps.now, - }) - await deps.client.auth.set({ - path: { id: 'openai' }, - body: { - type: 'oauth', - access: tokens.access, - refresh: tokens.refresh, - expires: tokens.expires, - }, - }) - auth = { ...auth, access: tokens.access, expires: tokens.expires } - } + const freshnessMs = deps.skipFresherThanMs + let quotaUpdated = false + let sharedSidebarState: SidebarState | undefined + if (freshnessMs !== undefined) { + try { + sharedSidebarState = await (deps.readSidebarState ?? getSidebarState)() + } catch {} + } + const sharedFallbacks = new Map( + sharedSidebarState?.fallbacks.map((account) => [account.id, account]) ?? [], + ) + const isFresh = (...checkedAts: unknown[]) => + freshnessMs !== undefined && + checkedAts.some( + (checkedAt) => + typeof checkedAt === 'number' && + Number.isFinite(checkedAt) && + checkedAt <= deps.now() && + deps.now() - checkedAt < freshnessMs, + ) - if (auth.access) { - if (deps.respectBackoff && deps.quotaManager.isBackedOff()) { + // Load the live storage once, up front, so the freshness gate judges identity + // by the account logged in NOW rather than the id captured when the loader + // initialized. A re-login within the same process changes mainAccountId on + // disk; comparing against the stale captured id would let the previous + // account's fresh quota suppress polling for the new one. A load failure fails + // open to the captured id so the main refresh still runs. + const storage = await deps + .loadAccounts(deps.configPath) + .catch(() => undefined) + const liveMainAccountId = storage?.mainAccountId ?? deps.storageMainAccountId + + if (!options.accountKey || options.accountKey === 'main') { + // --- MAIN --- + try { + let auth = await deps.getAuth() + if (auth.type === 'oauth') { + const sharedMainQuota = + sharedSidebarState && + liveMainAccountId !== undefined && + sharedSidebarState.main.mainAccountId === liveMainAccountId + ? sharedSidebarState.main.quota + : undefined + const freshMainQuota = isFresh( + deps.quotaManager.peekMainForPolicy(liveMainAccountId)?.checkedAt, + sharedMainQuota?.primary?.checkedAt, + sharedMainQuota?.secondary?.checkedAt, + sharedMainQuota?.checkedAt, + ) + if (freshMainQuota) { recordOutcome({ account: 'main', ok: true }) } else { - const snap = await whamFn({ - accessToken: auth.access, - fetchImpl: deps.fetchImpl, - now: deps.now, - accountId: deps.storageMainAccountId, - accountKey: 'main', - }) - deps.quotaManager.setMain( - auth.access, - { - quota: snap, - refreshAfter: deps.now() + 5 * 60 * 1000, - checkedAt: deps.now(), - }, - undefined, - true, - ) - recordOutcome({ account: 'main', ok: true }) + if (!auth.access || (auth.expires ?? 0) < deps.now()) { + const tokens = await deps.refreshMainWithLease() + auth = { ...auth, access: tokens.access, expires: tokens.expires } + } + + if (auth.access) { + if (deps.respectBackoff && deps.quotaManager.isBackedOff()) { + recordOutcome({ account: 'main', ok: true }) + } else { + const snap = await whamFn({ + accessToken: auth.access, + fetchImpl: deps.fetchImpl, + now: deps.now, + accountId: liveMainAccountId, + accountKey: 'main', + }) + deps.quotaManager.setMain( + auth.access, + { + quota: snap, + refreshAfter: deps.now() + 5 * 60 * 1000, + checkedAt: deps.now(), + }, + undefined, + true, + ) + quotaUpdated = true + recordOutcome({ account: 'main', ok: true }) + } + } else { + recordOutcome({ + account: 'main', + ok: false, + error: 'no access token', + }) + } } } else { recordOutcome({ account: 'main', ok: false, - error: 'no access token', + error: 'auth type is not oauth', }) } - } else { + } catch (e) { recordOutcome({ account: 'main', ok: false, - error: 'auth type is not oauth', + error: errorMessage(e), }) } - } catch (e) { - recordOutcome({ - account: 'main', - ok: false, - error: errorMessage(e), - }) } // --- FALLBACKS --- - const storage = await deps.loadAccounts(deps.configPath) if (storage) { for (const acct of storage.accounts) { - if (acct.enabled === false || !deps.isOAuthAccountFn(acct)) continue + if ( + options.accountKey && + (options.accountKey === 'main' || acct.id !== options.accountKey) + ) { + continue + } + if (acct.enabled === false || !deps.isOAuthAccountFn(acct)) { + if (options.accountKey) { + recordOutcome({ + account: acct.id, + ok: false, + error: 'account is not an enabled OAuth fallback', + }) + } + continue + } try { + const sharedFb = sharedFallbacks.get(acct.id) + const currentAccountId = (acct as OAuthAccount).accountId + const sharedFbQuota = + sharedFb && + currentAccountId !== undefined && + sharedFb.accountId === currentAccountId + ? sharedFb.quota + : undefined + if ( + isFresh( + deps.quotaManager.peekFallbackForPolicy(acct.id)?.checkedAt, + sharedFbQuota?.primary?.checkedAt, + sharedFbQuota?.secondary?.checkedAt, + sharedFbQuota?.checkedAt, + ) + ) { + recordOutcome({ account: acct.id, ok: true }) + continue + } + if ( deps.respectBackoff && deps.quotaManager.isFallbackBackedOff( @@ -202,6 +287,7 @@ export async function refreshAllQuota( refreshed.access, true, ) + quotaUpdated = true recordOutcome({ account: acct.id, ok: true }) } catch (e) { recordOutcome({ @@ -213,9 +299,22 @@ export async function refreshAllQuota( } } - // Refresh sidebar after all fetches - const freshStorage = await deps.loadAccounts(deps.configPath) - await deps.writeSidebarState(deps.quotaManager, freshStorage) + if ( + options.accountKey && + options.accountKey !== 'main' && + !results.some((result) => result.account === options.accountKey) + ) { + results.push({ + account: options.accountKey, + ok: false, + error: 'account not found', + }) + } + + if (freshnessMs === undefined || quotaUpdated) { + const freshStorage = await deps.loadAccounts(deps.configPath) + await deps.writeSidebarState(deps.quotaManager, freshStorage) + } return results } diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index bf5e249..b651268 100644 --- a/packages/opencode/src/index.ts +++ b/packages/opencode/src/index.ts @@ -39,6 +39,10 @@ import { type RoutingMode, shouldFallbackStatus, } from './core/accounts' +import { + BackgroundQuotaRefresh, + refreshQuotaInBackground, +} from './core/background-quota-refresh' import { buildRefreshOperationError, formatRefreshBackoffMessage, @@ -98,6 +102,7 @@ import { getRpcDir } from './rpc/rpc-dir' import { type RpcServerHandle, startRpcServer } from './rpc/rpc-server' import { type AccountQuota, + getSidebarState, getSidebarStateFile, removeSidebarActiveRouting, type SidebarMachineState, @@ -466,10 +471,14 @@ export function buildSidebarMachineState( store: AccountStorage, now = Date.now(), ): 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, + mainAccountId: store.mainAccountId, killed: false, ...(mainQuota?.resetCreditsAvailable !== undefined ? { resetCredits: mainQuota.resetCreditsAvailable } @@ -478,11 +487,18 @@ 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, + accountId: (account as OAuthAccount).accountId, + quota: fallbackQuota + ? { + ...(fallbackQuota as AccountQuota), + checkedAt: fallbackEntry.checkedAt, + } + : null, killed: false, enabled: true, ...(fallbackQuota?.resetCreditsAvailable !== undefined @@ -692,6 +708,12 @@ export async function CodexAuthPlugin( let activeRpcServer: RpcServerHandle | null = null let sidebarStateFileForEvents: string | undefined + // Per-loader poller: each plugin invocation owns its timer and callback, so + // one loader disposing or re-starting never stops or overwrites another's + // background refresh (a module-level singleton let the last loader win and + // let any disposal kill the shared poller). + const backgroundQuotaRefresh = new BackgroundQuotaRefresh() + async function sendIgnoredMessage(sessionId: string, text: string) { const session = input.client.session as | { promptAsync?: (req: unknown) => Promise } @@ -717,6 +739,7 @@ export async function CodexAuthPlugin( return { async dispose() { + backgroundQuotaRefresh.stop() for (const websocketFetch of websocketFetches) websocketFetch.close() websocketFetches.length = 0 if (activeRpcServer) { @@ -1426,6 +1449,7 @@ export async function CodexAuthPlugin( refreshAllQuota({ getAuth, codexRefreshFn, + refreshMainWithLease, fallbackManager, quotaManager, loadAccounts, @@ -1967,6 +1991,7 @@ export async function CodexAuthPlugin( void refreshAllQuota({ getAuth, codexRefreshFn, + refreshMainWithLease, fallbackManager, quotaManager, loadAccounts, @@ -1989,6 +2014,46 @@ export async function CodexAuthPlugin( ) } + backgroundQuotaRefresh.start( + async () => { + const results = await refreshQuotaInBackground({ + getAuth, + codexRefreshFn, + refreshMainWithLease, + fallbackManager, + quotaManager, + loadAccounts, + writeSidebarState: (qm, store) => + backgroundQuotaRefresh.isStopped() + ? Promise.resolve() + : writeMachineSidebarState(qm, store), + client: input.client as Parameters< + typeof refreshAllQuota + >[0]['client'], + fetchImpl: fetch, + now: Date.now, + configPath: getConfigPath(), + storageMainAccountId: storage?.mainAccountId, + isOAuthAccountFn: isOAuthAccount, + whamFn: whamUsageFn, + readSidebarState: () => getSidebarState(boundSidebarFile), + }) + const failures = results.filter((result) => !result.ok) + if (failures.length > 0) { + logQ.warn('background quota refresh completed with failures', { + pid: process.pid, + failures, + }) + } + }, + (error) => { + logQ.warn('background quota refresh failed', { + pid: process.pid, + error: error instanceof Error ? error.message : String(error), + }) + }, + ) + // ------------------------------------------------------------------- // Fetch override that selects the active account, refreshes if // needed, sends the transformed Codex request, and records quota. @@ -2215,6 +2280,7 @@ export async function CodexAuthPlugin( return finalResponse }, async dispose() { + backgroundQuotaRefresh.stop() cacheKeepManager.stop() if ( cacheKeepGlobal.__openaiAuthCacheKeepManager === cacheKeepManager diff --git a/packages/opencode/src/sidebar-state.ts b/packages/opencode/src/sidebar-state.ts index 83a25a1..fc4c538 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 @@ -496,6 +510,39 @@ export type SidebarMachineState = Pick< 'main' | 'fallbacks' | 'planType' | 'credits' | 'lastUpdated' > & { route: string } +// The freshest signal across every timestamp a snapshot carries: either window +// (primary/secondary) or the legacy top-level stamp. A retired primary window +// (null) with a fresh secondary must still outrank an older incoming primary, so +// the comparison takes the max rather than the first present value. +function latestQuotaCheckedAt(quota: AccountQuota | null): number | undefined { + let latest: number | undefined + for (const checkedAt of [ + quota?.primary?.checkedAt, + quota?.secondary?.checkedAt, + quota?.checkedAt, + ]) { + if (typeof checkedAt === 'number' && Number.isFinite(checkedAt)) { + latest = latest === undefined ? checkedAt : Math.max(latest, checkedAt) + } + } + return latest +} + +function freshestQuota( + incoming: AccountQuota | null, + existing: AccountQuota | null, +): AccountQuota | null { + const incomingCheckedAt = latestQuotaCheckedAt(incoming) + const existingCheckedAt = latestQuotaCheckedAt(existing) + if ( + existingCheckedAt !== undefined && + (incomingCheckedAt === undefined || existingCheckedAt > incomingCheckedAt) + ) { + return existing + } + return incoming +} + export function setSidebarMachineState( machineState: SidebarMachineState, file = getSidebarStateFile(), @@ -504,13 +551,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/background-quota-refresh.test.ts b/packages/opencode/src/tests/background-quota-refresh.test.ts new file mode 100644 index 0000000..b7fa19e --- /dev/null +++ b/packages/opencode/src/tests/background-quota-refresh.test.ts @@ -0,0 +1,389 @@ +import { describe, expect, mock, test } from 'bun:test' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + BACKGROUND_QUOTA_FRESHNESS_MS, + BACKGROUND_QUOTA_REFRESH_INTERVAL_MS, + BACKGROUND_QUOTA_REFRESH_JITTER_MS, + BACKGROUND_QUOTA_REFRESH_LOCK_NAME, + BACKGROUND_QUOTA_REFRESH_LOCK_TTL_MS, + BackgroundQuotaRefresh, + refreshQuotaInBackground, +} from '../core/background-quota-refresh' +import type { + RefreshAllQuotaDeps, + RefreshAllQuotaResult, +} from '../core/refresh-all-quota' +import { acquireRefreshFileLock } from '../core/refresh-file-lock' + +function timerHarness() { + let callback: (() => void) | undefined + let active = false + const handle = { + unref: mock(() => {}), + } as unknown as ReturnType + const setIntervalFn = mock((next: () => void, _intervalMs: number) => { + callback = next + active = true + return handle + }) + const clearIntervalFn = mock((timer: ReturnType) => { + expect(timer).toBe(handle) + active = false + }) + + return { + setIntervalFn, + clearIntervalFn, + handle, + tick: () => { + if (active) callback?.() + }, + } +} + +async function drainMicrotasks() { + await Promise.resolve() + await Promise.resolve() +} + +describe('BackgroundQuotaRefresh', () => { + test('background ticks force backoff respect and the freshness gate', async () => { + const deps = {} as RefreshAllQuotaDeps + const refreshFn = mock(async () => []) + + await refreshQuotaInBackground(deps, refreshFn, async () => ({ + release: async () => {}, + })) + + expect(refreshFn).toHaveBeenCalledWith({ + respectBackoff: true, + skipFresherThanMs: BACKGROUND_QUOTA_FRESHNESS_MS, + }) + }) + + test('concurrent background ticks claim a lock so only one refreshes', async () => { + const deps = { configPath: '/tmp/test-config.json' } as RefreshAllQuotaDeps + let held = false + const release = mock(async () => { + held = false + }) + const acquireLock = mock(async () => { + if (held) return null + held = true + return { release } + }) + let resolveRefresh!: () => void + const refreshFn = mock( + () => + new Promise((resolve) => { + resolveRefresh = () => resolve([]) + }), + ) + + const first = refreshQuotaInBackground(deps, refreshFn, acquireLock) + const second = refreshQuotaInBackground(deps, refreshFn, acquireLock) + await drainMicrotasks() + + // The first tick holds the lock and refreshes; the second skips this tick. + expect(refreshFn).toHaveBeenCalledTimes(1) + + resolveRefresh() + const [firstResults, secondResults] = await Promise.all([first, second]) + + expect(firstResults).toEqual([]) + expect(secondResults).toEqual([]) + expect(release).toHaveBeenCalledTimes(1) + }) + + test('a lock-mechanism failure fails open and still refreshes', async () => { + const deps = { configPath: '/tmp/test-config.json' } as RefreshAllQuotaDeps + const acquireLock = mock(async () => { + throw new Error('lock filesystem unavailable') + }) + const refreshFn = mock(async () => []) + + await refreshQuotaInBackground(deps, refreshFn, acquireLock) + + expect(refreshFn).toHaveBeenCalledTimes(1) + }) + + test('two concurrent ticks against a real lock file refresh only once', async () => { + const dir = mkdtempSync(join(tmpdir(), 'oai-bg-lock-')) + const configPath = join(dir, 'config.json') + try { + const deps = { configPath } as RefreshAllQuotaDeps + let resolveRefresh!: () => void + const refreshFn = mock( + () => + new Promise((resolve) => { + resolveRefresh = () => resolve([]) + }), + ) + + const first = refreshQuotaInBackground(deps, refreshFn) + const second = refreshQuotaInBackground(deps, refreshFn) + await new Promise((resolve) => setTimeout(resolve, 100)) + + expect(refreshFn).toHaveBeenCalledTimes(1) + + resolveRefresh() + await Promise.all([first, second]) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + test('timer runs the registered quota refresh with the fixed cadence', async () => { + const timer = timerHarness() + const run = mock(async () => {}) + const poller = new BackgroundQuotaRefresh({ + ...timer, + random: () => 0.5, + }) + + poller.start(run) + timer.tick() + await drainMicrotasks() + + expect(timer.setIntervalFn).toHaveBeenCalledWith( + expect.any(Function), + BACKGROUND_QUOTA_REFRESH_INTERVAL_MS, + ) + expect(run).toHaveBeenCalledTimes(1) + expect( + (timer.handle as unknown as { unref: ReturnType }).unref, + ).toHaveBeenCalledTimes(1) + }) + + test('jitter shifts the cadence within the +/- bound', () => { + const low = timerHarness() + new BackgroundQuotaRefresh({ ...low, random: () => 0 }).start( + mock(async () => {}), + ) + expect(low.setIntervalFn).toHaveBeenCalledWith( + expect.any(Function), + BACKGROUND_QUOTA_REFRESH_INTERVAL_MS - BACKGROUND_QUOTA_REFRESH_JITTER_MS, + ) + + const high = timerHarness() + new BackgroundQuotaRefresh({ ...high, random: () => 1 }).start( + mock(async () => {}), + ) + expect(high.setIntervalFn).toHaveBeenCalledWith( + expect.any(Function), + BACKGROUND_QUOTA_REFRESH_INTERVAL_MS + BACKGROUND_QUOTA_REFRESH_JITTER_MS, + ) + }) + + test('double start keeps one timer and uses the latest loader callback', async () => { + const timer = timerHarness() + const firstRun = mock(async () => {}) + const secondRun = mock(async () => {}) + const poller = new BackgroundQuotaRefresh({ + ...timer, + random: () => 0.5, + }) + + poller.start(firstRun) + poller.start(secondRun) + timer.tick() + await drainMicrotasks() + + expect(timer.setIntervalFn).toHaveBeenCalledTimes(1) + expect(firstRun).not.toHaveBeenCalled() + expect(secondRun).toHaveBeenCalledTimes(1) + }) + + test('stop clears the timer and prevents later ticks', async () => { + const timer = timerHarness() + const run = mock(async () => {}) + const poller = new BackgroundQuotaRefresh({ + ...timer, + random: () => 0.5, + }) + + poller.start(run) + poller.stop() + timer.tick() + await drainMicrotasks() + + expect(timer.clearIntervalFn).toHaveBeenCalledTimes(1) + expect(run).not.toHaveBeenCalled() + }) + + test('a failed poll is contained and the next tick still runs', async () => { + const timer = timerHarness() + const onError = mock(() => {}) + let attempts = 0 + const run = mock(async () => { + attempts += 1 + if (attempts === 1) throw new Error('wham unavailable') + }) + const poller = new BackgroundQuotaRefresh({ + ...timer, + random: () => 0.5, + onError, + }) + + poller.start(run) + timer.tick() + await drainMicrotasks() + timer.tick() + await drainMicrotasks() + + expect(run).toHaveBeenCalledTimes(2) + expect(onError).toHaveBeenCalledTimes(1) + expect(onError).toHaveBeenCalledWith(expect.any(Error)) + }) + + test('a tick is skipped while the previous refresh is still in flight', async () => { + const timer = timerHarness() + const resolvers: Array<() => void> = [] + const run = mock( + () => + new Promise((resolve) => { + resolvers.push(resolve) + }), + ) + const poller = new BackgroundQuotaRefresh({ + ...timer, + random: () => 0.5, + }) + + poller.start(run) + timer.tick() + timer.tick() + expect(run).toHaveBeenCalledTimes(1) + + resolvers[0]?.() + await drainMicrotasks() + timer.tick() + expect(run).toHaveBeenCalledTimes(2) + resolvers[1]?.() + await drainMicrotasks() + }) + + test('stop during an in-flight run lets it finish but fires no further ticks', async () => { + const timer = timerHarness() + let resolveRun!: () => void + let completed = false + const run = mock( + () => + new Promise((resolve) => { + resolveRun = () => { + completed = true + resolve() + } + }), + ) + const poller = new BackgroundQuotaRefresh({ ...timer, random: () => 0.5 }) + + poller.start(run) + expect(poller.isStopped()).toBe(false) + timer.tick() + expect(run).toHaveBeenCalledTimes(1) + + poller.stop() + expect(poller.isStopped()).toBe(true) + // The in-flight run is not cancelled... + expect(completed).toBe(false) + // ...and no further tick starts a new run. + timer.tick() + expect(run).toHaveBeenCalledTimes(1) + + resolveRun() + await drainMicrotasks() + expect(completed).toBe(true) + expect(run).toHaveBeenCalledTimes(1) + }) + + test('the stopped flag blocks a stray tick even if the timer fires once more', async () => { + let callback: (() => void) | undefined + const handle = { + unref: () => {}, + } as unknown as ReturnType + const setIntervalFn = mock((next: () => void) => { + callback = next + return handle + }) + const clearIntervalFn = mock(() => {}) + const run = mock(async () => {}) + const poller = new BackgroundQuotaRefresh({ + setIntervalFn, + clearIntervalFn, + random: () => 0.5, + }) + + poller.start(run) + poller.stop() + callback?.() + await drainMicrotasks() + + expect(run).not.toHaveBeenCalled() + expect(clearIntervalFn).toHaveBeenCalledTimes(1) + }) + + test('the lock TTL outlasts a slow multi-account refresh', async () => { + const dir = mkdtempSync(join(tmpdir(), 'oai-bg-ttl-')) + const configPath = join(dir, 'config.json') + try { + let nowMs = 1_000_000 + const opts = { + name: BACKGROUND_QUOTA_REFRESH_LOCK_NAME, + ttlMs: BACKGROUND_QUOTA_REFRESH_LOCK_TTL_MS, + path: configPath, + now: () => nowMs, + } + + // The first acquirer holds the lock for the whole (slow) refresh. + const holder = await acquireRefreshFileLock(opts) + expect(holder).not.toBeNull() + + // A serial pass over many accounts can run 35s; a contender arriving then + // must still be locked out rather than start a concurrent refresh. + nowMs += 35_000 + const contender = await acquireRefreshFileLock(opts) + expect(contender).toBeNull() + + // Discriminator: the null above means "held", not a broken acquirer — once + // the holder releases, acquisition succeeds again. + await holder?.release() + const afterRelease = await acquireRefreshFileLock(opts) + expect(afterRelease).not.toBeNull() + await afterRelease?.release() + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + test('separate poller instances run and dispose independently', async () => { + const timer1 = timerHarness() + const timer2 = timerHarness() + const run1 = mock(async () => {}) + const run2 = mock(async () => {}) + const loader1 = new BackgroundQuotaRefresh({ ...timer1, random: () => 0.5 }) + const loader2 = new BackgroundQuotaRefresh({ ...timer2, random: () => 0.5 }) + + // Two loaders start → two independent timers. + loader1.start(run1) + loader2.start(run2) + expect(timer1.setIntervalFn).toHaveBeenCalledTimes(1) + expect(timer2.setIntervalFn).toHaveBeenCalledTimes(1) + + // Dispose loader 1 → loader 2's timer still fires. + loader1.stop() + timer1.tick() + timer2.tick() + await drainMicrotasks() + expect(run1).not.toHaveBeenCalled() + expect(run2).toHaveBeenCalledTimes(1) + + // Dispose loader 2 → no further tick fires. + loader2.stop() + timer2.tick() + await drainMicrotasks() + expect(run2).toHaveBeenCalledTimes(1) + }) +}) diff --git a/packages/opencode/src/tests/quota-push.test.ts b/packages/opencode/src/tests/quota-push.test.ts index 602a5b0..b81a001 100644 --- a/packages/opencode/src/tests/quota-push.test.ts +++ b/packages/opencode/src/tests/quota-push.test.ts @@ -422,6 +422,7 @@ describe('QuotaManager push', () => { const expectedMachine = { main: { quota: { + checkedAt: 1, primary: { usedPercent: 20, remainingPercent: 80, @@ -438,6 +439,7 @@ describe('QuotaManager push', () => { id: 'fallback-1', label: undefined, quota: { + checkedAt: 1, primary: { usedPercent: 30, remainingPercent: 70, diff --git a/packages/opencode/src/tests/refresh-all-quota.test.ts b/packages/opencode/src/tests/refresh-all-quota.test.ts index ca550b3..bfc4c32 100644 --- a/packages/opencode/src/tests/refresh-all-quota.test.ts +++ b/packages/opencode/src/tests/refresh-all-quota.test.ts @@ -1,4 +1,7 @@ import { describe, expect, mock, test } from 'bun:test' +import { mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' import type { AccountQuotaWindow, FallbackAccount, @@ -9,6 +12,11 @@ import { type RefreshAllQuotaDeps, refreshAllQuota, } from '../core/refresh-all-quota' +import { + DEFAULT_SIDEBAR_STATE, + getSidebarState, + normalizeSidebarState, +} from '../sidebar-state' function makeQuotaSnapshot(usedPercent: number): OAuthQuotaSnapshot { const window: AccountQuotaWindow = { @@ -19,6 +27,41 @@ function makeQuotaSnapshot(usedPercent: number): OAuthQuotaSnapshot { return { primary: window } } +function sidebarSnapshot(checkedAt: unknown) { + return { + main: { + mainAccountId: 'chatgpt-main', + quota: { primary: { checkedAt } }, + }, + fallbacks: [ + { + id: 'fb-1', + accountId: 'chatgpt-fb1', + quota: { primary: { checkedAt } }, + }, + { + id: 'fb-2', + accountId: 'chatgpt-fb2', + quota: { primary: { checkedAt } }, + }, + ], + } +} + +async function withSidebarFile( + contents: string | undefined, + run: (path: string) => Promise, +) { + const dir = mkdtempSync(join(tmpdir(), 'oai-quota-refresh-')) + const path = join(dir, 'sidebar.json') + try { + if (contents !== undefined) writeFileSync(path, contents) + await run(path) + } finally { + rmSync(dir, { recursive: true, force: true }) + } +} + interface MakeDepsOptions extends Partial { accounts?: FallbackAccount[] } @@ -69,6 +112,11 @@ function makeDeps(opts: MakeDepsOptions = {}): RefreshAllQuotaDeps { refresh: 'refresh-new', expires: Date.now() + 7200_000, })), + refreshMainWithLease: mock(async () => ({ + access: 'access-refreshed', + refresh: 'refresh-new', + expires: Date.now() + 7200_000, + })), fallbackManager: { refreshAccount: mock(async (acct) => acct), } as unknown as RefreshAllQuotaDeps['fallbackManager'], @@ -88,6 +136,7 @@ function makeDeps(opts: MakeDepsOptions = {}): RefreshAllQuotaDeps { (a as { type?: string })?.type === 'oauth') as RefreshAllQuotaDeps['isOAuthAccountFn'], whamFn: mock(async () => makeQuotaSnapshot(30)), + readSidebarState: mock(async () => DEFAULT_SIDEBAR_STATE), } const { accounts: _a, ...rest } = opts @@ -184,7 +233,7 @@ describe('refreshAllQuota', () => { expect(serialized).not.toContain('access-fb2') }) - test('expired main token → codexRefreshFn called before wham', async () => { + test('expired main token → refreshMainWithLease called, not the direct refresh', async () => { const deps = makeDeps({ getAuth: mock(async () => ({ type: 'oauth' as const, @@ -196,8 +245,12 @@ describe('refreshAllQuota', () => { const results = await refreshAllQuota(deps) - expect(deps.codexRefreshFn).toHaveBeenCalled() - expect(deps.client.auth.set).toHaveBeenCalled() + // Main refresh routes through the cross-process lease; the direct + // codexRefreshFn path (and its client.auth.set mirror, which the lease + // performs internally) must not run here. + expect(deps.refreshMainWithLease).toHaveBeenCalled() + expect(deps.codexRefreshFn).not.toHaveBeenCalled() + expect(deps.client.auth.set).not.toHaveBeenCalled() expect(results[0]).toEqual({ account: 'main', ok: true }) expect(deps.quotaManager.getMain()?.quota?.primary?.usedPercent).toBe(30) }) @@ -412,4 +465,366 @@ describe('refreshAllQuota', () => { expect(results[1]).toEqual({ account: 'fb-1', ok: true }) expect(deps.whamFn).toHaveBeenCalledTimes(1) // only main, not fb-1 }) + + test('skipFresherThanMs skips quota fetched within the freshness window', async () => { + const now = Date.now() + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + readSidebarState: mock(async () => + normalizeSidebarState(sidebarSnapshot(now - 10 * 60_000)), + ), + }) + const freshEntry = { + quota: makeQuotaSnapshot(20), + refreshAfter: now + 4 * 60_000, + checkedAt: now - 60_000, + } + deps.quotaManager.setMain('access-main', freshEntry, 'chatgpt-main') + deps.quotaManager.setFallback('fb-1', freshEntry, 'access-fb1') + + await refreshAllQuota(deps) + + expect(deps.whamFn).toHaveBeenCalledTimes(1) + expect(deps.whamFn).toHaveBeenCalledWith( + expect.objectContaining({ accessToken: 'access-fb2' }), + ) + }) + + test('skipFresherThanMs fetches stale and missing quota snapshots', async () => { + const now = Date.now() + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + }) + deps.quotaManager.setMain( + 'access-main', + { + quota: makeQuotaSnapshot(20), + refreshAfter: now - 5 * 60_000, + checkedAt: now - 10 * 60_000, + }, + 'chatgpt-main', + ) + + await refreshAllQuota(deps) + + expect(deps.whamFn).toHaveBeenCalledTimes(3) + }) + + test('fresh machine-global sidebar quota skips polling with an empty in-memory cache', async () => { + const now = Date.now() + await withSidebarFile( + JSON.stringify(sidebarSnapshot(now - 60_000)), + async (path) => { + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + readSidebarState: () => getSidebarState(path), + }) + + await refreshAllQuota(deps) + + expect(deps.whamFn).not.toHaveBeenCalled() + expect(deps.writeSidebarState).not.toHaveBeenCalled() + }, + ) + }) + + test('stale machine-global sidebar quota is polled with an empty in-memory cache', async () => { + const now = Date.now() + await withSidebarFile( + JSON.stringify(sidebarSnapshot(now - 10 * 60_000)), + async (path) => { + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + readSidebarState: () => getSidebarState(path), + }) + + await refreshAllQuota(deps) + + expect(deps.whamFn).toHaveBeenCalledTimes(3) + }, + ) + }) + + test('missing, corrupt, and malformed shared quota fail open to polling', async () => { + const now = Date.now() + const cases = [ + undefined, + '{not-json', + JSON.stringify(sidebarSnapshot('recent-but-not-numeric')), + ] + + for (const contents of cases) { + await withSidebarFile(contents, async (path) => { + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + readSidebarState: () => getSidebarState(path), + }) + + await refreshAllQuota(deps) + + expect(deps.whamFn).toHaveBeenCalledTimes(3) + }) + } + + const nonFiniteDeps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + readSidebarState: mock(async () => + normalizeSidebarState(sidebarSnapshot(Number.NaN)), + ), + }) + await refreshAllQuota(nonFiniteDeps) + expect(nonFiniteDeps.whamFn).toHaveBeenCalledTimes(3) + }) + + test('manual refresh ignores fresh machine-global sidebar quota', async () => { + const now = Date.now() + await withSidebarFile( + JSON.stringify(sidebarSnapshot(now - 60_000)), + async (path) => { + const readSidebarState = mock(() => getSidebarState(path)) + const deps = makeDeps({ + now: () => now, + readSidebarState, + }) + + await refreshAllQuota(deps) + + expect(readSidebarState).not.toHaveBeenCalled() + expect(deps.whamFn).toHaveBeenCalledTimes(3) + }, + ) + }) + + // -- identity-bound freshness (account switch / re-login) -- + + test('main re-login (mainAccountId change) fetches despite a fresh sidebar checkedAt', async () => { + const now = Date.now() + // The loader captured the OLD account id and the sidebar still carries that + // account's fresh quota under the same id — so the captured id alone would + // match the file and suppress the poll. Only a fresh storage load reveals the + // re-login (mainAccountId now differs); the gate must use that live id. + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + accounts: [], + storageMainAccountId: 'chatgpt-main-old', + loadAccounts: mock(async () => ({ + version: 1 as const, + accounts: [], + mainAccountId: 'chatgpt-main-new', + })), + readSidebarState: mock(async () => + normalizeSidebarState({ + main: { + mainAccountId: 'chatgpt-main-old', + quota: { + primary: { + usedPercent: 5, + remainingPercent: 95, + checkedAt: now - 60_000, + }, + }, + }, + fallbacks: [], + }), + ), + }) + + await refreshAllQuota(deps) + + expect(deps.whamFn).toHaveBeenCalledTimes(1) + expect(deps.whamFn).toHaveBeenCalledWith( + expect.objectContaining({ accountId: 'chatgpt-main-new' }), + ) + }) + + test('re-login polls even when the in-memory cache is fresh under the stale identity', async () => { + const now = Date.now() + // Both the in-memory cache and the sidebar hold the OLD account's fresh + // quota, and the loader-captured id still matches them. Without the live + // storage load, both freshness signals agree and the poll is wrongly + // suppressed for the account that just logged in. + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + accounts: [], + storageMainAccountId: 'chatgpt-main-old', + loadAccounts: mock(async () => ({ + version: 1 as const, + accounts: [], + mainAccountId: 'chatgpt-main-new', + })), + readSidebarState: mock(async () => + normalizeSidebarState({ + main: { + mainAccountId: 'chatgpt-main-old', + quota: { + primary: { + usedPercent: 5, + remainingPercent: 95, + checkedAt: now - 60_000, + }, + }, + }, + fallbacks: [], + }), + ), + }) + deps.quotaManager.setMain( + 'access-main', + { + quota: makeQuotaSnapshot(5), + refreshAfter: now + 4 * 60_000, + checkedAt: now - 60_000, + }, + 'chatgpt-main-old', + ) + + await refreshAllQuota(deps) + + expect(deps.whamFn).toHaveBeenCalledTimes(1) + expect(deps.whamFn).toHaveBeenCalledWith( + expect.objectContaining({ accountId: 'chatgpt-main-new' }), + ) + }) + + test('fallback re-login (accountId change) fetches despite a fresh sidebar checkedAt', async () => { + const now = Date.now() + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + accounts: [ + { + id: 'fb-1', + type: 'oauth' as const, + access: 'access-fb1', + refresh: 'refresh-fb1', + expires: now + 3600_000, + enabled: true, + accountId: 'chatgpt-fb1-new', + }, + ], + readSidebarState: mock(async () => + normalizeSidebarState({ + main: { + mainAccountId: 'chatgpt-main', + quota: { + primary: { + usedPercent: 5, + remainingPercent: 95, + checkedAt: now - 60_000, + }, + }, + }, + fallbacks: [ + { + id: 'fb-1', + accountId: 'chatgpt-fb1-old', + quota: { + primary: { + usedPercent: 5, + remainingPercent: 95, + checkedAt: now - 60_000, + }, + }, + }, + ], + }), + ), + }) + + await refreshAllQuota(deps) + + // main matches identity + fresh → skipped; fb-1 identity changed → fetched. + expect(deps.whamFn).toHaveBeenCalledTimes(1) + expect(deps.whamFn).toHaveBeenCalledWith( + expect.objectContaining({ accessToken: 'access-fb1' }), + ) + }) + + // -- secondary-window freshness -- + + test('fresh secondary window alone skips polling when the primary window is absent (main)', async () => { + const now = Date.now() + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + accounts: [], + readSidebarState: mock(async () => + normalizeSidebarState({ + main: { + mainAccountId: 'chatgpt-main', + quota: { + secondary: { + usedPercent: 40, + remainingPercent: 60, + checkedAt: now - 60_000, + }, + }, + }, + fallbacks: [], + }), + ), + }) + + await refreshAllQuota(deps) + + expect(deps.whamFn).not.toHaveBeenCalled() + }) + + test('fresh secondary window alone skips polling when the primary window is absent (fallback)', async () => { + const now = Date.now() + const deps = makeDeps({ + now: () => now, + skipFresherThanMs: 4 * 60_000, + accounts: [ + { + id: 'fb-1', + type: 'oauth' as const, + access: 'access-fb1', + refresh: 'refresh-fb1', + expires: now + 3600_000, + enabled: true, + accountId: 'chatgpt-fb1', + }, + ], + readSidebarState: mock(async () => + normalizeSidebarState({ + main: { + mainAccountId: 'chatgpt-main', + quota: { + primary: { + usedPercent: 5, + remainingPercent: 95, + checkedAt: now - 60_000, + }, + }, + }, + fallbacks: [ + { + id: 'fb-1', + accountId: 'chatgpt-fb1', + quota: { + secondary: { + usedPercent: 40, + remainingPercent: 60, + checkedAt: now - 60_000, + }, + }, + }, + ], + }), + ), + }) + + await refreshAllQuota(deps) + + expect(deps.whamFn).not.toHaveBeenCalled() + }) }) diff --git a/packages/opencode/src/tests/sidebar-state.test.ts b/packages/opencode/src/tests/sidebar-state.test.ts index 783b985..981928b 100644 --- a/packages/opencode/src/tests/sidebar-state.test.ts +++ b/packages/opencode/src/tests/sidebar-state.test.ts @@ -35,8 +35,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 }, }) @@ -132,6 +136,7 @@ describe('normalizeSidebarState', () => { }, secondary: { usedPercent: 17, remainingPercent: 83 }, }, + mainAccountId: 'chatgpt-main', killed: true, quotaBackedOff: true, quotaBackoffUntil: 1234567890, @@ -143,6 +148,7 @@ describe('normalizeSidebarState', () => { { id: 'fb1', label: 'work', + accountId: 'chatgpt-fb1', quota: { primary: { usedPercent: 5, remainingPercent: 95 } }, killed: false, enabled: true, @@ -172,10 +178,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 +1332,254 @@ 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('machine write ranks a fresh secondary window above an older incoming primary', async () => { + const tempDir = mkdtempSync(join(tmpdir(), 'oai-sb-secondary-fresh-')) + const file = join(tempDir, 'sidebar-state.json') + const now = Date.now() + const stale = now - 10 * 60_000 + // On disk the primary window is retired (absent) but the secondary window is + // fresh; the incoming primary is older than that secondary. The merge must rank + // the disk's fresh secondary above the incoming primary and keep the disk quota, + // not discard it just because its primary slot is null. + await setSidebarState( + make({ + main: { + ...main(null), + quota: { + secondary: { usedPercent: 25, remainingPercent: 75, checkedAt: now }, + }, + }, + }), + file, + ) + + await setSidebarMachineState( + { + main: { + ...main(null), + quota: { + primary: { usedPercent: 80, remainingPercent: 20, checkedAt: stale }, + }, + }, + fallbacks: [], + route: 'main-first', + lastUpdated: now + 1, + }, + file, + ) + await drainSidebarWrites() + + const written = normalizeSidebarState(JSON.parse(readFileSync(file, 'utf8'))) + expect(written.main.quota?.secondary?.usedPercent).toBe(25) + expect(written.main.quota?.primary).toBeUndefined() +}) + +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')