Skip to content

Commit e67f58c

Browse files
iceteaSAiceteaSA
committed
feat(quota): background poll keeps idle account quota fresh
1 parent 6559eb4 commit e67f58c

6 files changed

Lines changed: 536 additions & 39 deletions

File tree

‎packages/opencode/src/core/background-quota-refresh.ts‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import {
2+
type RefreshAllQuotaDeps,
3+
type RefreshAllQuotaResult,
4+
refreshAllQuota,
5+
} from './refresh-all-quota'
6+
7+
export const BACKGROUND_QUOTA_REFRESH_INTERVAL_MS = 5 * 60_000
8+
export const BACKGROUND_QUOTA_REFRESH_JITTER_MS = 30_000
9+
export const BACKGROUND_QUOTA_FRESHNESS_MS = 4 * 60_000
10+
11+
type TimerHandle = ReturnType<typeof setInterval>
12+
13+
interface BackgroundQuotaRefreshOptions {
14+
setIntervalFn?: (callback: () => void, intervalMs: number) => TimerHandle
15+
clearIntervalFn?: (timer: TimerHandle) => void
16+
random?: () => number
17+
onError?: (error: unknown) => void
18+
}
19+
20+
type RefreshAllQuotaFn = (
21+
deps: RefreshAllQuotaDeps,
22+
) => Promise<RefreshAllQuotaResult[]>
23+
24+
export function refreshQuotaInBackground(
25+
deps: RefreshAllQuotaDeps,
26+
refreshFn: RefreshAllQuotaFn = refreshAllQuota,
27+
): Promise<RefreshAllQuotaResult[]> {
28+
return refreshFn({
29+
...deps,
30+
respectBackoff: true,
31+
skipFresherThanMs: BACKGROUND_QUOTA_FRESHNESS_MS,
32+
})
33+
}
34+
35+
export class BackgroundQuotaRefresh {
36+
private readonly setIntervalFn: NonNullable<
37+
BackgroundQuotaRefreshOptions['setIntervalFn']
38+
>
39+
private readonly clearIntervalFn: NonNullable<
40+
BackgroundQuotaRefreshOptions['clearIntervalFn']
41+
>
42+
private readonly random: () => number
43+
private onError: ((error: unknown) => void) | undefined
44+
private run: (() => Promise<void>) | undefined
45+
private timer: TimerHandle | undefined
46+
47+
constructor(options: BackgroundQuotaRefreshOptions = {}) {
48+
this.setIntervalFn = options.setIntervalFn ?? setInterval
49+
this.clearIntervalFn = options.clearIntervalFn ?? clearInterval
50+
this.random = options.random ?? Math.random
51+
this.onError = options.onError
52+
}
53+
54+
start(
55+
run: () => Promise<void>,
56+
onError: ((error: unknown) => void) | undefined = this.onError,
57+
): void {
58+
this.run = run
59+
this.onError = onError
60+
if (this.timer) return
61+
62+
const jitter = Math.round(
63+
(this.random() * 2 - 1) * BACKGROUND_QUOTA_REFRESH_JITTER_MS,
64+
)
65+
this.timer = this.setIntervalFn(() => {
66+
const currentRun = this.run
67+
if (!currentRun) return
68+
void currentRun().catch((error) => {
69+
try {
70+
this.onError?.(error)
71+
} catch {}
72+
})
73+
}, BACKGROUND_QUOTA_REFRESH_INTERVAL_MS + jitter)
74+
if ('unref' in this.timer) this.timer.unref()
75+
}
76+
77+
stop(): void {
78+
this.run = undefined
79+
if (!this.timer) return
80+
this.clearIntervalFn(this.timer)
81+
this.timer = undefined
82+
}
83+
}

‎packages/opencode/src/core/refresh-all-quota.ts‎

Lines changed: 80 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { getSidebarState, type SidebarState } from '../sidebar-state'
12
import type {
23
FallbackAccountManager,
34
isOAuthAccount,
@@ -50,6 +51,8 @@ export interface RefreshAllQuotaDeps {
5051
isOAuthAccountFn: typeof isOAuthAccount
5152
whamFn?: typeof whamUsageFn
5253
respectBackoff?: boolean
54+
skipFresherThanMs?: number
55+
readSidebarState?: () => Promise<SidebarState>
5356
}
5457

5558
export interface RefreshAllQuotaResult {
@@ -65,53 +68,83 @@ export async function refreshAllQuota(
6568
if (!whamFn) throw new Error('whamFn is required for refreshAllQuota')
6669

6770
const results: RefreshAllQuotaResult[] = []
71+
const freshnessMs = deps.skipFresherThanMs
72+
let sharedSidebarState: SidebarState | undefined
73+
if (freshnessMs !== undefined) {
74+
try {
75+
sharedSidebarState = await (deps.readSidebarState ?? getSidebarState)()
76+
} catch {}
77+
}
78+
const sharedFallbackCheckedAt = new Map(
79+
sharedSidebarState?.fallbacks.map((account) => [
80+
account.id,
81+
account.quota?.primary?.checkedAt,
82+
]) ?? [],
83+
)
84+
const isFresh = (...checkedAts: unknown[]) =>
85+
freshnessMs !== undefined &&
86+
checkedAts.some(
87+
(checkedAt) =>
88+
typeof checkedAt === 'number' &&
89+
Number.isFinite(checkedAt) &&
90+
deps.now() - checkedAt < freshnessMs,
91+
)
6892

6993
// --- MAIN ---
7094
try {
7195
let auth = await deps.getAuth()
7296
if (auth.type === 'oauth') {
73-
if (!auth.access || (auth.expires ?? 0) < deps.now()) {
74-
const tokens = await deps.codexRefreshFn({
75-
refreshToken: auth.refresh ?? '',
76-
fetchImpl: deps.fetchImpl,
77-
now: deps.now,
78-
})
79-
await deps.client.auth.set({
80-
path: { id: 'openai' },
81-
body: {
82-
type: 'oauth',
83-
access: tokens.access,
84-
refresh: tokens.refresh,
85-
expires: tokens.expires,
86-
},
87-
})
88-
auth = { ...auth, access: tokens.access, expires: tokens.expires }
89-
}
90-
91-
if (auth.access) {
92-
if (deps.respectBackoff && deps.quotaManager.isBackedOff()) {
93-
results.push({ account: 'main', ok: true })
94-
} else {
95-
const snap = await whamFn({
96-
accessToken: auth.access,
97+
const freshMainQuota = isFresh(
98+
deps.quotaManager.peekMainForPolicy(deps.storageMainAccountId)
99+
?.checkedAt,
100+
sharedSidebarState?.main.quota?.primary?.checkedAt,
101+
)
102+
if (freshMainQuota) {
103+
results.push({ account: 'main', ok: true })
104+
} else {
105+
if (!auth.access || (auth.expires ?? 0) < deps.now()) {
106+
const tokens = await deps.codexRefreshFn({
107+
refreshToken: auth.refresh ?? '',
97108
fetchImpl: deps.fetchImpl,
98109
now: deps.now,
99-
accountId: deps.storageMainAccountId,
100110
})
101-
deps.quotaManager.setMain(
102-
auth.access,
103-
{
104-
quota: snap,
105-
refreshAfter: deps.now() + 5 * 60 * 1000,
106-
checkedAt: deps.now(),
111+
await deps.client.auth.set({
112+
path: { id: 'openai' },
113+
body: {
114+
type: 'oauth',
115+
access: tokens.access,
116+
refresh: tokens.refresh,
117+
expires: tokens.expires,
107118
},
108-
undefined,
109-
true,
110-
)
111-
results.push({ account: 'main', ok: true })
119+
})
120+
auth = { ...auth, access: tokens.access, expires: tokens.expires }
121+
}
122+
123+
if (auth.access) {
124+
if (deps.respectBackoff && deps.quotaManager.isBackedOff()) {
125+
results.push({ account: 'main', ok: true })
126+
} else {
127+
const snap = await whamFn({
128+
accessToken: auth.access,
129+
fetchImpl: deps.fetchImpl,
130+
now: deps.now,
131+
accountId: deps.storageMainAccountId,
132+
})
133+
deps.quotaManager.setMain(
134+
auth.access,
135+
{
136+
quota: snap,
137+
refreshAfter: deps.now() + 5 * 60 * 1000,
138+
checkedAt: deps.now(),
139+
},
140+
undefined,
141+
true,
142+
)
143+
results.push({ account: 'main', ok: true })
144+
}
145+
} else {
146+
results.push({ account: 'main', ok: false, error: 'no access token' })
112147
}
113-
} else {
114-
results.push({ account: 'main', ok: false, error: 'no access token' })
115148
}
116149
} else {
117150
results.push({
@@ -135,6 +168,16 @@ export async function refreshAllQuota(
135168
if (acct.enabled === false || !deps.isOAuthAccountFn(acct)) continue
136169

137170
try {
171+
if (
172+
isFresh(
173+
deps.quotaManager.peekFallbackForPolicy(acct.id)?.checkedAt,
174+
sharedFallbackCheckedAt.get(acct.id),
175+
)
176+
) {
177+
results.push({ account: acct.id, ok: true })
178+
continue
179+
}
180+
138181
if (
139182
deps.respectBackoff &&
140183
deps.quotaManager.isFallbackBackedOff(

‎packages/opencode/src/index.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ import {
3939
type RoutingMode,
4040
shouldFallbackStatus,
4141
} from './core/accounts'
42+
import {
43+
BackgroundQuotaRefresh,
44+
refreshQuotaInBackground,
45+
} from './core/background-quota-refresh'
4246
import {
4347
buildRefreshOperationError,
4448
formatRefreshBackoffMessage,
@@ -98,6 +102,7 @@ import { getRpcDir } from './rpc/rpc-dir'
98102
import { type RpcServerHandle, startRpcServer } from './rpc/rpc-server'
99103
import {
100104
type AccountQuota,
105+
getSidebarState,
101106
getSidebarStateFile,
102107
removeSidebarActiveRouting,
103108
type SidebarMachineState,
@@ -152,6 +157,7 @@ const DEFAULT_MID_STREAM_RATE_LIMIT_RESET_MS = 60_000
152157
const HANDLED_SENTINEL = '__OPENCODE_OPENAI_AUTH_COMMAND_HANDLED__'
153158

154159
let bootQuotaSeedStarted = false
160+
const backgroundQuotaRefresh = new BackgroundQuotaRefresh()
155161
const logModels = createLogger('models')
156162
let loggedCostRestoration = false
157163
let warnedCostCatalogUnavailable = false
@@ -716,6 +722,7 @@ export async function CodexAuthPlugin(
716722

717723
return {
718724
async dispose() {
725+
backgroundQuotaRefresh.stop()
719726
for (const websocketFetch of websocketFetches) websocketFetch.close()
720727
websocketFetches.length = 0
721728
if (activeRpcServer) {
@@ -1974,6 +1981,42 @@ export async function CodexAuthPlugin(
19741981
}).catch(() => {})
19751982
}
19761983

1984+
backgroundQuotaRefresh.start(
1985+
async () => {
1986+
const results = await refreshQuotaInBackground({
1987+
getAuth,
1988+
codexRefreshFn,
1989+
fallbackManager,
1990+
quotaManager,
1991+
loadAccounts,
1992+
writeSidebarState: writeMachineSidebarState,
1993+
client: input.client as Parameters<
1994+
typeof refreshAllQuota
1995+
>[0]['client'],
1996+
fetchImpl: fetch,
1997+
now: Date.now,
1998+
configPath: getConfigPath(),
1999+
storageMainAccountId: storage?.mainAccountId,
2000+
isOAuthAccountFn: isOAuthAccount,
2001+
whamFn: whamUsageFn,
2002+
readSidebarState: () => getSidebarState(boundSidebarFile),
2003+
})
2004+
const failures = results.filter((result) => !result.ok)
2005+
if (failures.length > 0) {
2006+
logQ.warn('background quota refresh completed with failures', {
2007+
pid: process.pid,
2008+
failures,
2009+
})
2010+
}
2011+
},
2012+
(error) => {
2013+
logQ.warn('background quota refresh failed', {
2014+
pid: process.pid,
2015+
error: error instanceof Error ? error.message : String(error),
2016+
})
2017+
},
2018+
)
2019+
19772020
// -------------------------------------------------------------------
19782021
// Fetch override that selects the active account, refreshes if
19792022
// needed, sends the transformed Codex request, and records quota.
@@ -2198,6 +2241,7 @@ export async function CodexAuthPlugin(
21982241
return finalResponse
21992242
},
22002243
async dispose() {
2244+
backgroundQuotaRefresh.stop()
22012245
cacheKeepManager.stop()
22022246
if (
22032247
cacheKeepGlobal.__openaiAuthCacheKeepManager === cacheKeepManager

‎packages/opencode/src/sidebar-state.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
export interface QuotaWindow {
22
usedPercent: number
33
remainingPercent: number
4+
checkedAt?: number
45
resetsAt?: string
56
windowMinutes?: number
67
}
@@ -280,9 +281,11 @@ export function normalizeSidebarState(raw: unknown): SidebarState {
280281
}
281282
}
282283

283-
export async function getSidebarState(): Promise<SidebarState> {
284+
export async function getSidebarState(
285+
stateFile = getSidebarStateFile(),
286+
): Promise<SidebarState> {
284287
try {
285-
const raw = await readFile(getSidebarStateFile(), 'utf8')
288+
const raw = await readFile(stateFile, 'utf8')
286289
return normalizeSidebarState(JSON.parse(raw))
287290
} catch {
288291
return DEFAULT_SIDEBAR_STATE

0 commit comments

Comments
 (0)