Skip to content

Commit 1399ceb

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

5 files changed

Lines changed: 382 additions & 37 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: 57 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ export interface RefreshAllQuotaDeps {
5050
isOAuthAccountFn: typeof isOAuthAccount
5151
whamFn?: typeof whamUsageFn
5252
respectBackoff?: boolean
53+
skipFresherThanMs?: number
5354
}
5455

5556
export interface RefreshAllQuotaResult {
@@ -65,53 +66,65 @@ export async function refreshAllQuota(
6566
if (!whamFn) throw new Error('whamFn is required for refreshAllQuota')
6667

6768
const results: RefreshAllQuotaResult[] = []
69+
const isFresh = (checkedAt: number | undefined) =>
70+
deps.skipFresherThanMs !== undefined &&
71+
checkedAt !== undefined &&
72+
deps.now() - checkedAt < deps.skipFresherThanMs
6873

6974
// --- MAIN ---
7075
try {
7176
let auth = await deps.getAuth()
7277
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,
78+
const freshMainQuota = isFresh(
79+
deps.quotaManager.peekMainForPolicy(deps.storageMainAccountId)
80+
?.checkedAt,
81+
)
82+
if (freshMainQuota) {
83+
results.push({ account: 'main', ok: true })
84+
} else {
85+
if (!auth.access || (auth.expires ?? 0) < deps.now()) {
86+
const tokens = await deps.codexRefreshFn({
87+
refreshToken: auth.refresh ?? '',
9788
fetchImpl: deps.fetchImpl,
9889
now: deps.now,
99-
accountId: deps.storageMainAccountId,
10090
})
101-
deps.quotaManager.setMain(
102-
auth.access,
103-
{
104-
quota: snap,
105-
refreshAfter: deps.now() + 5 * 60 * 1000,
106-
checkedAt: deps.now(),
91+
await deps.client.auth.set({
92+
path: { id: 'openai' },
93+
body: {
94+
type: 'oauth',
95+
access: tokens.access,
96+
refresh: tokens.refresh,
97+
expires: tokens.expires,
10798
},
108-
undefined,
109-
true,
110-
)
111-
results.push({ account: 'main', ok: true })
99+
})
100+
auth = { ...auth, access: tokens.access, expires: tokens.expires }
101+
}
102+
103+
if (auth.access) {
104+
if (deps.respectBackoff && deps.quotaManager.isBackedOff()) {
105+
results.push({ account: 'main', ok: true })
106+
} else {
107+
const snap = await whamFn({
108+
accessToken: auth.access,
109+
fetchImpl: deps.fetchImpl,
110+
now: deps.now,
111+
accountId: deps.storageMainAccountId,
112+
})
113+
deps.quotaManager.setMain(
114+
auth.access,
115+
{
116+
quota: snap,
117+
refreshAfter: deps.now() + 5 * 60 * 1000,
118+
checkedAt: deps.now(),
119+
},
120+
undefined,
121+
true,
122+
)
123+
results.push({ account: 'main', ok: true })
124+
}
125+
} else {
126+
results.push({ account: 'main', ok: false, error: 'no access token' })
112127
}
113-
} else {
114-
results.push({ account: 'main', ok: false, error: 'no access token' })
115128
}
116129
} else {
117130
results.push({
@@ -135,6 +148,13 @@ export async function refreshAllQuota(
135148
if (acct.enabled === false || !deps.isOAuthAccountFn(acct)) continue
136149

137150
try {
151+
if (
152+
isFresh(deps.quotaManager.peekFallbackForPolicy(acct.id)?.checkedAt)
153+
) {
154+
results.push({ account: acct.id, ok: true })
155+
continue
156+
}
157+
138158
if (
139159
deps.respectBackoff &&
140160
deps.quotaManager.isFallbackBackedOff(

‎packages/opencode/src/index.ts‎

Lines changed: 42 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,
@@ -152,6 +156,7 @@ const DEFAULT_MID_STREAM_RATE_LIMIT_RESET_MS = 60_000
152156
const HANDLED_SENTINEL = '__OPENCODE_OPENAI_AUTH_COMMAND_HANDLED__'
153157

154158
let bootQuotaSeedStarted = false
159+
const backgroundQuotaRefresh = new BackgroundQuotaRefresh()
155160
const logModels = createLogger('models')
156161
let loggedCostRestoration = false
157162
let warnedCostCatalogUnavailable = false
@@ -716,6 +721,7 @@ export async function CodexAuthPlugin(
716721

717722
return {
718723
async dispose() {
724+
backgroundQuotaRefresh.stop()
719725
for (const websocketFetch of websocketFetches) websocketFetch.close()
720726
websocketFetches.length = 0
721727
if (activeRpcServer) {
@@ -1974,6 +1980,41 @@ export async function CodexAuthPlugin(
19741980
}).catch(() => {})
19751981
}
19761982

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

0 commit comments

Comments
 (0)