From 0deb16fbb12ab8746e03c05b058612ac472ac21c Mon Sep 17 00:00:00 2001 From: ndycode Date: Thu, 5 Mar 2026 22:41:24 +0800 Subject: [PATCH 1/8] fix: enforce automatic quota-based account rotation Rotate immediately on 429 responses and bootstrap runtime selection from persisted quota cache (accountId-first, normalized email fallback) with a 30m lazy TTL and env escape hatch. Co-authored-by: Codex --- index.ts | 316 ++++++++++++++++++++++++++++++--------- test/index-retry.test.ts | 10 ++ test/index.test.ts | 269 +++++++++++++++++++++++++++++++++ 3 files changed, 521 insertions(+), 74 deletions(-) diff --git a/index.ts b/index.ts index 7db88088a..0f2860dfa 100644 --- a/index.ts +++ b/index.ts @@ -114,6 +114,7 @@ import { shouldUpdateAccountIdFromToken, resolveRequestAccountId, parseRateLimitReason, + getQuotaKey, lookupCodexCliTokensByEmail, isCodexCliSyncEnabled, } from "./lib/accounts.js"; @@ -151,7 +152,6 @@ import { import { applyFastSessionDefaults } from "./lib/request/request-transformer.js"; import { getRateLimitBackoff, - RATE_LIMIT_SHORT_RETRY_THRESHOLD_MS, resetRateLimitBackoff, } from "./lib/request/rate-limit-backoff.js"; import { isEmptyResponse } from "./lib/request/response-handler.js"; @@ -170,8 +170,14 @@ import { import { PreemptiveQuotaScheduler, readQuotaSchedulerSnapshot, + type QuotaSchedulerSnapshot, } from "./lib/preemptive-quota-scheduler.js"; import { CapabilityPolicyStore } from "./lib/capability-policy.js"; +import { + loadQuotaCache, + type QuotaCacheData, + type QuotaCacheEntry, +} from "./lib/quota-cache.js"; import { withStreamingFailover } from "./lib/request/stream-failover.js"; import { buildTableHeader, buildTableRow, type TableOptions } from "./lib/table-formatter.js"; import { setUiRuntimeOptions, type UiRuntimeOptions } from "./lib/ui/runtime.js"; @@ -300,6 +306,179 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { return null; }; + const QUOTA_CACHE_BOOTSTRAP_TTL_MS = 30 * 60_000; + const QUOTA_CACHE_BOOTSTRAP_MIN_WAIT_MS = 1_000; + let quotaCacheBootstrapData: QuotaCacheData | null = null; + let quotaCacheBootstrapLoadedAt = 0; + let quotaCacheBootstrapLoadPromise: Promise | null = null; + + type QuotaBootstrapAccountCandidate = { + index: number; + accountId?: string; + email?: string; + }; + + const normalizeQuotaCacheEmail = (value: string | undefined): string | null => { + const normalized = sanitizeEmail(value); + if (typeof normalized === "string" && normalized.trim().length > 0) { + return normalized.trim().toLowerCase(); + } + if (typeof value !== "string") return null; + const fallback = value.trim().toLowerCase(); + return fallback.length > 0 ? fallback : null; + }; + + const getQuotaCacheEntryForCandidate = ( + cache: QuotaCacheData, + candidate: QuotaBootstrapAccountCandidate, + ): QuotaCacheEntry | null => { + if (candidate.accountId) { + const byAccountId = cache.byAccountId[candidate.accountId]; + if (byAccountId) return byAccountId; + } + const email = normalizeQuotaCacheEmail(candidate.email); + if (!email) return null; + return cache.byEmail[email] ?? null; + }; + + const getQuotaCacheEntryFutureResetAtMs = ( + entry: QuotaCacheEntry, + now: number, + ): number | null => { + const candidates = [entry.primary.resetAtMs, entry.secondary.resetAtMs] + .filter((value): value is number => + typeof value === "number" && Number.isFinite(value) && value > now, + ); + if (candidates.length === 0) return null; + return Math.min(...candidates); + }; + + const getQuotaCacheBootstrapWaitMs = ( + entry: QuotaCacheEntry, + now: number, + ): number => { + const futureResetAtMs = getQuotaCacheEntryFutureResetAtMs(entry, now); + if (typeof futureResetAtMs === "number") { + return Math.max(0, Math.floor(futureResetAtMs - now)); + } + const maxAgeUntil = entry.updatedAt + QUOTA_CACHE_BOOTSTRAP_TTL_MS; + if (!Number.isFinite(maxAgeUntil) || maxAgeUntil <= now) { + return 0; + } + return Math.max(0, Math.floor(maxAgeUntil - now)); + }; + + const isQuotaCacheEntryExhausted = (entry: QuotaCacheEntry): boolean => { + const primaryUsedPercent = entry.primary.usedPercent; + if (typeof primaryUsedPercent === "number" && Number.isFinite(primaryUsedPercent)) { + if (primaryUsedPercent >= 100) return true; + } + const secondaryUsedPercent = entry.secondary.usedPercent; + if (typeof secondaryUsedPercent === "number" && Number.isFinite(secondaryUsedPercent)) { + if (secondaryUsedPercent >= 100) return true; + } + return false; + }; + + const shouldApplyQuotaCacheEntry = ( + entry: QuotaCacheEntry, + now: number, + ): boolean => { + const isRateLimited = entry.status === 429; + if (!isRateLimited && !isQuotaCacheEntryExhausted(entry)) { + return false; + } + return getQuotaCacheBootstrapWaitMs(entry, now) > 0; + }; + + const toQuotaSchedulerSnapshot = (entry: QuotaCacheEntry): QuotaSchedulerSnapshot => ({ + status: entry.status, + primary: { + usedPercent: entry.primary.usedPercent, + resetAtMs: entry.primary.resetAtMs, + }, + secondary: { + usedPercent: entry.secondary.usedPercent, + resetAtMs: entry.secondary.resetAtMs, + }, + updatedAt: entry.updatedAt, + }); + + const loadQuotaCacheForBootstrap = async (): Promise => { + const now = Date.now(); + if (quotaCacheBootstrapData && now - quotaCacheBootstrapLoadedAt < QUOTA_CACHE_BOOTSTRAP_TTL_MS) { + return quotaCacheBootstrapData; + } + if (quotaCacheBootstrapLoadPromise) { + return quotaCacheBootstrapLoadPromise; + } + quotaCacheBootstrapLoadPromise = (async () => { + const loaded = await loadQuotaCache(); + quotaCacheBootstrapData = loaded; + quotaCacheBootstrapLoadedAt = Date.now(); + return loaded; + })(); + try { + return await quotaCacheBootstrapLoadPromise; + } finally { + quotaCacheBootstrapLoadPromise = null; + } + }; + + const applyQuotaCacheBootstrapForModel = async ( + accountManager: AccountManager, + accountSnapshots: QuotaBootstrapAccountCandidate[], + modelFamily: ModelFamily, + model: string | undefined, + enabled: boolean, + ): Promise => { + if (!enabled || accountSnapshots.length === 0) return; + const cache = await loadQuotaCacheForBootstrap(); + const now = Date.now(); + const modelKey = getQuotaKey(modelFamily, model); + const baseKey = getQuotaKey(modelFamily); + + for (const snapshotCandidate of accountSnapshots) { + const entry = getQuotaCacheEntryForCandidate(cache, snapshotCandidate); + if (!entry) continue; + if (getModelFamily(entry.model) !== modelFamily) { + continue; + } + if (!shouldApplyQuotaCacheEntry(entry, now)) { + continue; + } + + const waitMs = Math.max( + QUOTA_CACHE_BOOTSTRAP_MIN_WAIT_MS, + getQuotaCacheBootstrapWaitMs(entry, now), + ); + const account = accountManager.getAccountByIndex(snapshotCandidate.index); + if (!account) continue; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + + const existingBaseResetAt = account.rateLimitResetTimes[baseKey] ?? 0; + const existingModelResetAt = account.rateLimitResetTimes[modelKey] ?? 0; + const existingResetAt = Math.max(existingBaseResetAt, existingModelResetAt); + const nextResetAt = now + waitMs; + if (existingResetAt >= nextResetAt) { + continue; + } + + accountManager.markRateLimitedWithReason( + account, + waitMs, + modelFamily, + "quota", + model, + ); + const quotaScheduleKey = `${resolveEntitlementAccountKey(snapshotCandidate)}:${model ?? modelFamily}`; + preemptiveQuotaScheduler.update( + quotaScheduleKey, + toQuotaSchedulerSnapshot(entry), + ); + } + }; + const sanitizeResponseHeadersForLog = (headers: Headers): Record => { const allowed = new Set([ "content-type", @@ -1135,6 +1314,8 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { const streamStallTimeoutMs = getStreamStallTimeoutMs(pluginConfig); const networkErrorCooldownMs = getNetworkErrorCooldownMs(pluginConfig); const serverErrorCooldownMs = getServerErrorCooldownMs(pluginConfig); + const quotaCacheBootstrapEnabled = + (process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP ?? "1").trim() !== "0"; const failoverMode = parseFailoverMode(process.env.CODEX_AUTH_FAILOVER_MODE); const streamFailoverMax = Math.max( 0, @@ -1439,6 +1620,13 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { } } } + await applyQuotaCacheBootstrapForModel( + accountManager, + accountSnapshotList, + modelFamily, + model ?? undefined, + quotaCacheBootstrapEnabled, + ); for (const candidate of accountSnapshotList) { const accountKey = resolveEntitlementAccountKey(candidate); capabilityBoostByAccount[candidate.index] = capabilityPolicyStore.getBoost( @@ -2005,83 +2193,63 @@ while (attempted.size < Math.max(1, accountCount)) { } if (rateLimit) { - runtimeMetrics.rateLimitedResponses++; - const { attempt, delayMs } = getRateLimitBackoff( - account.index, - quotaKey, - rateLimit.retryAfterMs, - ); - preemptiveQuotaScheduler.markRateLimited( - quotaScheduleKey, - delayMs, - ); - const waitLabel = formatWaitTime(delayMs); - - if (delayMs <= RATE_LIMIT_SHORT_RETRY_THRESHOLD_MS) { - if ( - accountManager.shouldShowAccountToast( - account.index, - rateLimitToastDebounceMs, - ) - ) { - await showToast( - `Rate limited. Retrying in ${waitLabel} (attempt ${attempt})...`, - "warning", - { duration: toastDurationMs }, - ); - accountManager.markToastShown(account.index); - } - - await sleep(addJitter(Math.max(MIN_BACKOFF_MS, delayMs), 0.2)); - continue; - } - - accountManager.markRateLimitedWithReason( - account, - delayMs, - modelFamily, - parseRateLimitReason(rateLimit.code), - model, - ); - accountManager.recordRateLimit(account, modelFamily, model); - account.lastSwitchReason = "rate-limit"; - sessionAffinityStore?.forgetSession(sessionAffinityKey); - runtimeMetrics.accountRotations++; - accountManager.saveToDiskDebounced(); + runtimeMetrics.rateLimitedResponses++; + const { attempt, delayMs } = getRateLimitBackoff( + account.index, + quotaKey, + rateLimit.retryAfterMs, + ); + preemptiveQuotaScheduler.markRateLimited( + quotaScheduleKey, + delayMs, + ); + const waitLabel = formatWaitTime(delayMs); + accountManager.markRateLimitedWithReason( + account, + delayMs, + modelFamily, + parseRateLimitReason(rateLimit.code), + model, + ); + accountManager.recordRateLimit(account, modelFamily, model); + account.lastSwitchReason = "rate-limit"; + sessionAffinityStore?.forgetSession(sessionAffinityKey); + runtimeMetrics.accountRotations++; + accountManager.saveToDiskDebounced(); logWarn( `Rate limited. Rotating account ${account.index + 1} (${account.email ?? "unknown"}).`, ); - if ( - accountManager.getAccountCount() > 1 && - accountManager.shouldShowAccountToast( - account.index, - rateLimitToastDebounceMs, - ) - ) { - await showToast( - `Rate limited. Switching accounts (retry in ${waitLabel}).`, - "warning", - { duration: toastDurationMs }, - ); - accountManager.markToastShown(account.index); - } - break; - } - if ( - !rateLimit && - !unsupportedModelInfo.isUnsupported && - errorResponse.status !== 403 - ) { - capabilityPolicyStore.recordFailure( - entitlementAccountKey, - capabilityModelKey, - ); - } - runtimeMetrics.failedRequests++; - runtimeMetrics.lastError = `HTTP ${response.status}`; - return errorResponse; - } + if ( + accountManager.getAccountCount() > 1 && + accountManager.shouldShowAccountToast( + account.index, + rateLimitToastDebounceMs, + ) + ) { + await showToast( + `Rate limited. Switching accounts (retry in ${waitLabel}, attempt ${attempt}).`, + "warning", + { duration: toastDurationMs }, + ); + accountManager.markToastShown(account.index); + } + break; + } + if ( + !rateLimit && + !unsupportedModelInfo.isUnsupported && + errorResponse.status !== 403 + ) { + capabilityPolicyStore.recordFailure( + entitlementAccountKey, + capabilityModelKey, + ); + } + runtimeMetrics.failedRequests++; + runtimeMetrics.lastError = `HTTP ${response.status}`; + return errorResponse; + } resetRateLimitBackoff(account.index, quotaKey); runtimeMetrics.cumulativeLatencyMs += fetchLatencyMs; diff --git a/test/index-retry.test.ts b/test/index-retry.test.ts index 3813edc48..f5ee4a063 100644 --- a/test/index-retry.test.ts +++ b/test/index-retry.test.ts @@ -35,6 +35,10 @@ vi.mock("../lib/request/request-transformer.js", () => ({ applyFastSessionDefaults: (config: T) => config, })); +vi.mock("../lib/quota-cache.js", () => ({ + loadQuotaCache: vi.fn(async () => ({ byAccountId: {}, byEmail: {} })), +})); + vi.mock("../lib/accounts.js", () => { class AccountManager { private calls = 0; @@ -57,6 +61,11 @@ vi.mock("../lib/accounts.js", () => { return this.getCurrentOrNextForFamily(); } + getAccountByIndex(index: number) { + if (index !== 0) return null; + return { index: 0, accountId: "account-1", email: "user@example.com" }; + } + recordSuccess() {} recordRateLimit() {} @@ -120,6 +129,7 @@ vi.mock("../lib/accounts.js", () => { formatWaitTime: (ms: number) => `${ms}ms`, sanitizeEmail: (email: string) => email, parseRateLimitReason: () => "unknown", + getQuotaKey: (family: string, model?: string | null) => (model ? `${family}:${model}` : family), lookupCodexCliTokensByEmail: vi.fn(async () => null), isCodexCliSyncEnabled: () => true, }; diff --git a/test/index.test.ts b/test/index.test.ts index d6d95497f..1eda77707 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -197,6 +197,11 @@ vi.mock("../lib/request/rate-limit-backoff.js", () => ({ resetRateLimitBackoff: vi.fn(), })); +const loadQuotaCacheMock = vi.fn(async () => ({ byAccountId: {}, byEmail: {} })); +vi.mock("../lib/quota-cache.js", () => ({ + loadQuotaCache: loadQuotaCacheMock, +})); + vi.mock("../lib/request/fetch-helpers.js", () => ({ extractRequestUrl: (input: unknown) => (typeof input === "string" ? input : String(input)), rewriteUrlForCodex: (url: string) => url, @@ -330,6 +335,10 @@ vi.mock("../lib/accounts.js", () => { return this.accounts[index] ?? null; } + getAccountByIndex(index: number) { + return this.accounts[index] ?? null; + } + getAccountsSnapshot() { return this.accounts; } @@ -348,6 +357,7 @@ vi.mock("../lib/accounts.js", () => { sanitizeEmail: (email: string) => email, shouldUpdateAccountIdFromToken: () => true, parseRateLimitReason: () => "unknown", + getQuotaKey: (family: string, model?: string | null) => (model ? `${family}:${model}` : family), lookupCodexCliTokensByEmail: vi.fn(async () => null), isCodexCliSyncEnabled: () => true, }; @@ -1017,6 +1027,8 @@ describe("OpenAIOAuthPlugin fetch handler", () => { beforeEach(() => { vi.clearAllMocks(); syncCodexCliSelectionMock.mockClear(); + loadQuotaCacheMock.mockReset(); + loadQuotaCacheMock.mockResolvedValue({ byAccountId: {}, byEmail: {} }); mockStorage.accounts = [ { accountId: "acc-1", @@ -1169,6 +1181,263 @@ describe("OpenAIOAuthPlugin fetch handler", () => { consumeSpy.mockRestore(); }); + it("rotates immediately to next account on short-window 429", async () => { + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementationOnce(() => accountOne as never) + .mockImplementationOnce(() => accountTwo as never) + .mockImplementation(() => null as never); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const rateLimitSpy = vi.spyOn(AccountManager.prototype, "markRateLimitedWithReason"); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + vi.mocked(fetchHelpers.handleErrorResponse).mockResolvedValueOnce({ + response: new Response("rate limited", { status: 429 }), + rateLimit: { retryAfterMs: 1_000, code: "rate_limit_exceeded" }, + errorBody: {}, + } as never); + globalThis.fetch = vi + .fn() + .mockResolvedValueOnce(new Response("rate limited", { status: 429 })) + .mockResolvedValueOnce( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(2); + const firstHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + const secondHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[1]?.[1] as RequestInit)?.headers, + ); + expect(firstHeaders.get("x-test-access-token")).toBe("access-acc-1"); + expect(secondHeaders.get("x-test-access-token")).toBe("access-acc-2"); + expect(rateLimitSpy).toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + rateLimitSpy.mockRestore(); + }); + + it("uses quota cache bootstrap to skip an already rate-limited account", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementationOnce(() => accountOne as never) + .mockImplementationOnce(() => accountTwo as never) + .mockImplementation(() => null as never); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-2"); + expect(loadQuotaCacheMock).toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + }); + + it("disables quota cache bootstrap when CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP=0", async () => { + const previous = process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP; + process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP = "0"; + try { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementationOnce(() => accountOne as never) + .mockImplementationOnce(() => accountTwo as never) + .mockImplementation(() => null as never); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-1"); + expect(loadQuotaCacheMock).not.toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + toAuthSpy.mockRestore(); + } finally { + if (previous === undefined) { + delete process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP; + } else { + process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP = previous; + } + } + }); + it("treats timeout-triggered abort as network failure", async () => { const { AccountManager } = await import("../lib/accounts.js"); const configModule = await import("../lib/config.js"); From f9838a3c44b26db1ab3dae82432fea5ebe04ada3 Mon Sep 17 00:00:00 2001 From: ndycode Date: Thu, 5 Mar 2026 23:08:50 +0800 Subject: [PATCH 2/8] fix: harden quota bootstrap rotation behavior - block exhausted accounts until the furthest reset window - make quota bootstrap loading best-effort on cache failures - add regressions for email fallback, concurrency, refresh races, and EPERM/EBUSY Co-authored-by: Codex --- index.ts | 31 +++- test/index.test.ts | 428 ++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 446 insertions(+), 13 deletions(-) diff --git a/index.ts b/index.ts index 0f2860dfa..54c1a480a 100644 --- a/index.ts +++ b/index.ts @@ -350,7 +350,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { typeof value === "number" && Number.isFinite(value) && value > now, ); if (candidates.length === 0) return null; - return Math.min(...candidates); + return Math.max(...candidates); }; const getQuotaCacheBootstrapWaitMs = ( @@ -413,10 +413,19 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { return quotaCacheBootstrapLoadPromise; } quotaCacheBootstrapLoadPromise = (async () => { - const loaded = await loadQuotaCache(); - quotaCacheBootstrapData = loaded; - quotaCacheBootstrapLoadedAt = Date.now(); - return loaded; + try { + const loaded = await loadQuotaCache(); + quotaCacheBootstrapData = loaded; + quotaCacheBootstrapLoadedAt = Date.now(); + return loaded; + } catch (error) { + logWarn( + `[${PLUGIN_NAME}] failed to load quota cache bootstrap data: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return { byAccountId: {}, byEmail: {} }; + } })(); try { return await quotaCacheBootstrapLoadPromise; @@ -433,7 +442,17 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { enabled: boolean, ): Promise => { if (!enabled || accountSnapshots.length === 0) return; - const cache = await loadQuotaCacheForBootstrap(); + let cache: QuotaCacheData; + try { + cache = await loadQuotaCacheForBootstrap(); + } catch (error) { + logWarn( + `[${PLUGIN_NAME}] quota cache bootstrap skipped: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return; + } const now = Date.now(); const modelKey = getQuotaKey(modelFamily, model); const baseKey = getQuotaKey(modelFamily); diff --git a/test/index.test.ts b/test/index.test.ts index 1eda77707..82c37e85a 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -209,7 +209,7 @@ vi.mock("../lib/quota-cache.js", () => ({ updatedInit: init, body: { model: "gpt-5.1" }, })), - shouldRefreshToken: () => false, + shouldRefreshToken: vi.fn(() => false), refreshAndUpdateToken: vi.fn(async (auth: unknown) => auth), createCodexHeaders: vi.fn(() => new Headers()), handleErrorResponse: vi.fn(async (response: Response) => ({ response })), @@ -1063,6 +1063,16 @@ describe("OpenAIOAuthPlugin fetch handler", () => { return { plugin, sdk, mockClient }; }; + const createDeferred = () => { + let resolve!: (value: T) => void; + let reject!: (error?: unknown) => void; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + return { promise, resolve, reject }; + }; + it("returns success response for successful fetch", async () => { globalThis.fetch = vi.fn().mockResolvedValue( new Response(JSON.stringify({ content: "test" }), { status: 200 }), @@ -1264,7 +1274,124 @@ describe("OpenAIOAuthPlugin fetch handler", () => { rateLimitSpy.mockRestore(); }); - it("uses quota cache bootstrap to skip an already rate-limited account", async () => { + it("uses quota cache bootstrap to skip an already rate-limited account via normalized email fallback", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "missing-acc-1", + email: " USER1@Example.com ", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: {}, + byEmail: { + "user1@example.com": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-2"); + expect(loadQuotaCacheMock).toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + + it("uses the furthest quota reset window when bootstrapping rate limits", async () => { const now = Date.now(); const { AccountManager } = await import("../lib/accounts.js"); const fetchHelpers = await import("../lib/request/fetch-helpers.js"); @@ -1289,10 +1416,13 @@ describe("OpenAIOAuthPlugin fetch handler", () => { status: 429, model: "gpt-5.1", primary: { + usedPercent: 100, + resetAtMs: now + 60_000, + }, + secondary: { usedPercent: 100, resetAtMs: now + 5 * 60_000, }, - secondary: {}, }, }, byEmail: {}, @@ -1302,9 +1432,22 @@ describe("OpenAIOAuthPlugin fetch handler", () => { .mockReturnValue(2); const selectionSpy = vi .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") - .mockImplementationOnce(() => accountOne as never) - .mockImplementationOnce(() => accountTwo as never) - .mockImplementation(() => null as never); + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); const snapshotSpy = vi .spyOn(AccountManager.prototype, "getAccountsSnapshot") .mockReturnValue([accountOne, accountTwo] as never); @@ -1321,6 +1464,24 @@ describe("OpenAIOAuthPlugin fetch handler", () => { refresh: `refresh-${account.accountId ?? "unknown"}`, expires: Date.now() + 60_000, })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( (_init, _accountId, accessToken) => new Headers({ "x-test-access-token": String(accessToken ?? "") }), @@ -1341,12 +1502,265 @@ describe("OpenAIOAuthPlugin fetch handler", () => { (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, ); expect(headers.get("x-test-access-token")).toBe("access-acc-2"); - expect(loadQuotaCacheMock).toHaveBeenCalled(); + const bootstrapWaitMs = markRateLimitSpy.mock.calls + .map((call) => call[1]) + .find((value): value is number => typeof value === "number"); + expect(bootstrapWaitMs).toBeGreaterThanOrEqual(4 * 60_000); countSpy.mockRestore(); selectionSpy.mockRestore(); snapshotSpy.mockRestore(); getByIndexSpy.mockRestore(); toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + + it("loads quota bootstrap cache once for concurrent fetch calls", async () => { + const testNow = Date.now() + 31 * 60_000; + const nowSpy = vi.spyOn(Date, "now").mockReturnValue(testNow); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + try { + const deferredLoad = createDeferred<{ byAccountId: Record; byEmail: Record }>(); + loadQuotaCacheMock.mockImplementationOnce(() => deferredLoad.promise); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const firstRequest = sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + const secondRequest = sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + for (let attempt = 0; attempt < 20 && loadQuotaCacheMock.mock.calls.length === 0; attempt += 1) { + await Promise.resolve(); + } + + expect(loadQuotaCacheMock).toHaveBeenCalledTimes(1); + deferredLoad.resolve({ byAccountId: {}, byEmail: {} }); + const [firstResponse, secondResponse] = await Promise.all([firstRequest, secondRequest]); + expect(firstResponse.status).toBe(200); + expect(secondResponse.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(2); + } finally { + nowSpy.mockRestore(); + } + }); + + it("handles token refresh races while rotating away from a bootstrap-rate-limited account", async () => { + const now = Date.now() + 31 * 60_000; + const nowSpy = vi.spyOn(Date, "now").mockReturnValue(now); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + try { + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 2 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => { + const accountId = account.accountId ?? "unknown"; + return { + type: "oauth" as const, + access: `stale-${accountId}`, + refresh: `refresh-${accountId}`, + expires: Date.now() - 60_000, + }; + }); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + const blockedRefresh = createDeferred<{ + type: "oauth"; + access: string; + refresh: string; + expires: number; + }>(); + let blockedRefreshUsed = false; + let refreshCallCount = 0; + vi.mocked(fetchHelpers.shouldRefreshToken).mockImplementation( + (auth: { access?: string }) => auth.access?.startsWith("stale-") ?? false, + ); + vi.mocked(fetchHelpers.refreshAndUpdateToken).mockImplementation( + async (auth: { refresh: string }) => { + refreshCallCount += 1; + if (!blockedRefreshUsed) { + blockedRefreshUsed = true; + return blockedRefresh.promise; + } + return { + type: "oauth" as const, + access: `refreshed-${refreshCallCount}`, + refresh: auth.refresh, + expires: Date.now() + 60_000, + }; + }, + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const firstRequest = sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + const secondRequest = sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + await Promise.resolve(); + blockedRefresh.resolve({ + type: "oauth", + access: "refreshed-blocked", + refresh: "refresh-acc-2", + expires: Date.now() + 60_000, + }); + + const [firstResponse, secondResponse] = await Promise.all([firstRequest, secondRequest]); + expect(firstResponse.status).toBe(200); + expect(secondResponse.status).toBe(200); + expect(refreshCallCount).toBeGreaterThanOrEqual(2); + expect(selectionSpy).toHaveBeenCalled(); + expect(getByIndexSpy).toHaveBeenCalled(); + const accessTokens = vi + .mocked(globalThis.fetch) + .mock.calls + .map((call) => { + const headers = new Headers((call[1] as RequestInit | undefined)?.headers); + return headers.get("x-test-access-token"); + }) + .filter((token): token is string => typeof token === "string"); + expect(accessTokens.length).toBe(2); + expect(accessTokens.every((token) => token.startsWith("refreshed-"))).toBe(true); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + } finally { + nowSpy.mockRestore(); + } + }); + + it("continues request flow when quota bootstrap cache reads hit EPERM/EBUSY", async () => { + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const epermError = Object.assign(new Error("locked"), { code: "EPERM" }); + const ebusyError = Object.assign(new Error("busy"), { code: "EBUSY" }); + loadQuotaCacheMock + .mockRejectedValueOnce(epermError) + .mockRejectedValueOnce(ebusyError); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const firstResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + const secondResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(firstResponse.status).toBe(200); + expect(secondResponse.status).toBe(200); + expect(loadQuotaCacheMock).toHaveBeenCalledTimes(2); + expect(globalThis.fetch).toHaveBeenCalledTimes(2); + const firstHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + const secondHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[1]?.[1] as RequestInit)?.headers, + ); + expect(firstHeaders.get("x-test-access-token")).toBe("access-token"); + expect(secondHeaders.get("x-test-access-token")).toBe("access-token"); }); it("disables quota cache bootstrap when CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP=0", async () => { From a17553b984a5bd5825c710af05dbb6e799e03279 Mon Sep 17 00:00:00 2001 From: ndycode Date: Thu, 5 Mar 2026 23:25:26 +0800 Subject: [PATCH 3/8] fix: align quota bootstrap to model-specific cache entries - avoid cross-model bootstrap blocking within the same family - align scheduler keys with the model that produced cached quota state - make concurrency/race tests deterministic with explicit barriers Co-authored-by: Codex --- index.ts | 17 +++++-- test/index.test.ts | 120 +++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 127 insertions(+), 10 deletions(-) diff --git a/index.ts b/index.ts index 54c1a480a..8047aa82c 100644 --- a/index.ts +++ b/index.ts @@ -454,18 +454,25 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { return; } const now = Date.now(); - const modelKey = getQuotaKey(modelFamily, model); + const requestedModel = + typeof model === "string" && model.trim().length > 0 ? model.trim() : null; const baseKey = getQuotaKey(modelFamily); for (const snapshotCandidate of accountSnapshots) { const entry = getQuotaCacheEntryForCandidate(cache, snapshotCandidate); if (!entry) continue; - if (getModelFamily(entry.model) !== modelFamily) { + const entryModel = entry.model.trim(); + if (getModelFamily(entryModel) !== modelFamily) { + continue; + } + if (requestedModel && entryModel !== requestedModel) { continue; } if (!shouldApplyQuotaCacheEntry(entry, now)) { continue; } + const appliedModel = requestedModel ?? entryModel; + const appliedModelKey = getQuotaKey(modelFamily, appliedModel); const waitMs = Math.max( QUOTA_CACHE_BOOTSTRAP_MIN_WAIT_MS, @@ -476,7 +483,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; const existingBaseResetAt = account.rateLimitResetTimes[baseKey] ?? 0; - const existingModelResetAt = account.rateLimitResetTimes[modelKey] ?? 0; + const existingModelResetAt = account.rateLimitResetTimes[appliedModelKey] ?? 0; const existingResetAt = Math.max(existingBaseResetAt, existingModelResetAt); const nextResetAt = now + waitMs; if (existingResetAt >= nextResetAt) { @@ -488,9 +495,9 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { waitMs, modelFamily, "quota", - model, + appliedModel, ); - const quotaScheduleKey = `${resolveEntitlementAccountKey(snapshotCandidate)}:${model ?? modelFamily}`; + const quotaScheduleKey = `${resolveEntitlementAccountKey(snapshotCandidate)}:${appliedModel}`; preemptiveQuotaScheduler.update( quotaScheduleKey, toQuotaSchedulerSnapshot(entry), diff --git a/test/index.test.ts b/test/index.test.ts index 82c37e85a..d73944ba0 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -1391,6 +1391,112 @@ describe("OpenAIOAuthPlugin fetch handler", () => { markRateLimitSpy.mockRestore(); }); + it("ignores quota bootstrap entries for a different model in the same family", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi.spyOn(AccountManager.prototype, "markRateLimitedWithReason"); + vi.mocked(fetchHelpers.transformRequestForCodex).mockImplementationOnce( + async (init: unknown) => ({ + updatedInit: init, + body: { model: "gpt-5.1-mini" }, + }), + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1-mini" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-1"); + expect(markRateLimitSpy).not.toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + it("uses the furthest quota reset window when bootstrapping rate limits", async () => { const now = Date.now(); const { AccountManager } = await import("../lib/accounts.js"); @@ -1520,7 +1626,11 @@ describe("OpenAIOAuthPlugin fetch handler", () => { const fetchHelpers = await import("../lib/request/fetch-helpers.js"); try { const deferredLoad = createDeferred<{ byAccountId: Record; byEmail: Record }>(); - loadQuotaCacheMock.mockImplementationOnce(() => deferredLoad.promise); + const bootstrapLoadStarted = createDeferred(); + loadQuotaCacheMock.mockImplementationOnce(async () => { + bootstrapLoadStarted.resolve(); + return deferredLoad.promise; + }); vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( (_init, _accountId, accessToken) => new Headers({ "x-test-access-token": String(accessToken ?? "") }), @@ -1538,9 +1648,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { method: "POST", body: JSON.stringify({ model: "gpt-5.1" }), }); - for (let attempt = 0; attempt < 20 && loadQuotaCacheMock.mock.calls.length === 0; attempt += 1) { - await Promise.resolve(); - } + await bootstrapLoadStarted.promise; expect(loadQuotaCacheMock).toHaveBeenCalledTimes(1); deferredLoad.resolve({ byAccountId: {}, byEmail: {} }); @@ -1652,6 +1760,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { refresh: string; expires: number; }>(); + const firstRefreshStarted = createDeferred(); let blockedRefreshUsed = false; let refreshCallCount = 0; vi.mocked(fetchHelpers.shouldRefreshToken).mockImplementation( @@ -1662,6 +1771,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { refreshCallCount += 1; if (!blockedRefreshUsed) { blockedRefreshUsed = true; + firstRefreshStarted.resolve(); return blockedRefresh.promise; } return { @@ -1689,7 +1799,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { method: "POST", body: JSON.stringify({ model: "gpt-5.1" }), }); - await Promise.resolve(); + await firstRefreshStarted.promise; blockedRefresh.resolve({ type: "oauth", access: "refreshed-blocked", From 5088063da07d4321fe7621d2349c49a84931c55f Mon Sep 17 00:00:00 2001 From: ndycode Date: Thu, 5 Mar 2026 23:45:51 +0800 Subject: [PATCH 4/8] fix: remove unreachable quota bootstrap catch - simplify bootstrap apply path to use the loader's best-effort fallback - keep error handling centralized in loadQuotaCacheForBootstrap Co-authored-by: Codex --- index.ts | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/index.ts b/index.ts index 8047aa82c..7b8f4c81d 100644 --- a/index.ts +++ b/index.ts @@ -442,17 +442,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { enabled: boolean, ): Promise => { if (!enabled || accountSnapshots.length === 0) return; - let cache: QuotaCacheData; - try { - cache = await loadQuotaCacheForBootstrap(); - } catch (error) { - logWarn( - `[${PLUGIN_NAME}] quota cache bootstrap skipped: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - return; - } + const cache = await loadQuotaCacheForBootstrap(); const now = Date.now(); const requestedModel = typeof model === "string" && model.trim().length > 0 ? model.trim() : null; From ce0d061392d5ebe83993b90a58fcd8b737a72533 Mon Sep 17 00:00:00 2001 From: ndycode Date: Fri, 6 Mar 2026 00:24:45 +0800 Subject: [PATCH 5/8] fix: harden quota bootstrap reload and model-id matching - avoid over-blocking when quota reset timestamps are known but elapsed - add failure cooldown with retry interval to prevent load hammering while allowing recovery - normalize model ids for bootstrap exact matching using provider-prefix stripping + lowercase - add regression coverage for model-id variants, elapsed resets, and cooldown recovery behavior Co-authored-by: Codex --- index.ts | 59 +++- test/index.test.ts | 664 ++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 703 insertions(+), 20 deletions(-) diff --git a/index.ts b/index.ts index 7b8f4c81d..c267e25b0 100644 --- a/index.ts +++ b/index.ts @@ -308,8 +308,12 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { const QUOTA_CACHE_BOOTSTRAP_TTL_MS = 30 * 60_000; const QUOTA_CACHE_BOOTSTRAP_MIN_WAIT_MS = 1_000; + const QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS = 60_000; + const QUOTA_CACHE_BOOTSTRAP_FAILURE_RETRY_INTERVAL_MS = 5_000; let quotaCacheBootstrapData: QuotaCacheData | null = null; let quotaCacheBootstrapLoadedAt = 0; + let quotaCacheBootstrapFailureUntil = 0; + let quotaCacheBootstrapFailureRetryAt = 0; let quotaCacheBootstrapLoadPromise: Promise | null = null; type QuotaBootstrapAccountCandidate = { @@ -328,6 +332,20 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { return fallback.length > 0 ? fallback : null; }; + const normalizeQuotaCacheModelId = (value: string | undefined): string | null => { + if (typeof value !== "string") return null; + const trimmed = value.trim(); + if (trimmed.length === 0) return null; + const withoutProvider = trimmed.includes("/") ? (trimmed.split("/").pop() ?? trimmed) : trimmed; + const normalized = withoutProvider.trim().toLowerCase(); + return normalized.length > 0 ? normalized : null; + }; + + const hasQuotaCacheEntries = (cache: QuotaCacheData): boolean => + Object.keys(cache.byAccountId).length > 0 || Object.keys(cache.byEmail).length > 0; + + const createEmptyQuotaCache = (): QuotaCacheData => ({ byAccountId: {}, byEmail: {} }); + const getQuotaCacheEntryForCandidate = ( cache: QuotaCacheData, candidate: QuotaBootstrapAccountCandidate, @@ -361,6 +379,12 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (typeof futureResetAtMs === "number") { return Math.max(0, Math.floor(futureResetAtMs - now)); } + const hasKnownResetTime = + typeof entry.primary.resetAtMs === "number" || + typeof entry.secondary.resetAtMs === "number"; + if (hasKnownResetTime) { + return 0; + } const maxAgeUntil = entry.updatedAt + QUOTA_CACHE_BOOTSTRAP_TTL_MS; if (!Number.isFinite(maxAgeUntil) || maxAgeUntil <= now) { return 0; @@ -409,22 +433,45 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (quotaCacheBootstrapData && now - quotaCacheBootstrapLoadedAt < QUOTA_CACHE_BOOTSTRAP_TTL_MS) { return quotaCacheBootstrapData; } + if ( + now < quotaCacheBootstrapFailureUntil && + now < quotaCacheBootstrapFailureRetryAt + ) { + return createEmptyQuotaCache(); + } if (quotaCacheBootstrapLoadPromise) { return quotaCacheBootstrapLoadPromise; } quotaCacheBootstrapLoadPromise = (async () => { try { const loaded = await loadQuotaCache(); - quotaCacheBootstrapData = loaded; - quotaCacheBootstrapLoadedAt = Date.now(); + if (hasQuotaCacheEntries(loaded)) { + quotaCacheBootstrapData = loaded; + quotaCacheBootstrapLoadedAt = Date.now(); + quotaCacheBootstrapFailureUntil = 0; + quotaCacheBootstrapFailureRetryAt = 0; + } else { + const failureAt = Date.now(); + quotaCacheBootstrapData = null; + quotaCacheBootstrapLoadedAt = 0; + quotaCacheBootstrapFailureUntil = failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS; + quotaCacheBootstrapFailureRetryAt = + failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_RETRY_INTERVAL_MS; + } return loaded; } catch (error) { + const failureAt = Date.now(); + quotaCacheBootstrapData = null; + quotaCacheBootstrapLoadedAt = 0; + quotaCacheBootstrapFailureUntil = failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS; + quotaCacheBootstrapFailureRetryAt = + failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_RETRY_INTERVAL_MS; logWarn( `[${PLUGIN_NAME}] failed to load quota cache bootstrap data: ${ error instanceof Error ? error.message : String(error) }`, ); - return { byAccountId: {}, byEmail: {} }; + return createEmptyQuotaCache(); } })(); try { @@ -444,14 +491,14 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (!enabled || accountSnapshots.length === 0) return; const cache = await loadQuotaCacheForBootstrap(); const now = Date.now(); - const requestedModel = - typeof model === "string" && model.trim().length > 0 ? model.trim() : null; + const requestedModel = normalizeQuotaCacheModelId(model); const baseKey = getQuotaKey(modelFamily); for (const snapshotCandidate of accountSnapshots) { const entry = getQuotaCacheEntryForCandidate(cache, snapshotCandidate); if (!entry) continue; - const entryModel = entry.model.trim(); + const entryModel = normalizeQuotaCacheModelId(entry.model); + if (!entryModel) continue; if (getModelFamily(entryModel) !== modelFamily) { continue; } diff --git a/test/index.test.ts b/test/index.test.ts index d73944ba0..24ac06a76 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -1391,6 +1391,129 @@ describe("OpenAIOAuthPlugin fetch handler", () => { markRateLimitSpy.mockRestore(); }); + it("matches quota bootstrap entries when requested model ids differ by provider prefix or casing", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + vi.mocked(fetchHelpers.transformRequestForCodex).mockImplementationOnce( + async (init: unknown) => ({ + updatedInit: init, + body: { model: "OpenAI/GPT-5.1" }, + }), + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "OpenAI/GPT-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-2"); + expect(markRateLimitSpy).toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + it("ignores quota bootstrap entries for a different model in the same family", async () => { const now = Date.now(); const { AccountManager } = await import("../lib/accounts.js"); @@ -1417,9 +1540,378 @@ describe("OpenAIOAuthPlugin fetch handler", () => { model: "gpt-5.1", primary: { usedPercent: 100, - resetAtMs: now + 5 * 60_000, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + vi.mocked(fetchHelpers.transformRequestForCodex).mockImplementationOnce( + async (init: unknown) => ({ + updatedInit: init, + body: { model: "gpt-5.1-mini" }, + }), + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1-mini" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-1"); + expect(markRateLimitSpy).not.toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + + it("matches semantically equivalent model ids while bootstrapping quota entries", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "OPENAI/GPT-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + vi.mocked(fetchHelpers.transformRequestForCodex).mockImplementationOnce( + async (init: unknown) => ({ + updatedInit: init, + body: { model: "openai/gpt-5.1" }, + }), + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "openai/gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-2"); + expect(markRateLimitSpy).toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + + it("does not apply ttl fallback when quota reset timestamps are elapsed", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now - 60_000, + }, + secondary: { + usedPercent: 100, + resetAtMs: now - 30_000, + }, + }, + }, + byEmail: {}, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-1"); + expect(markRateLimitSpy).not.toHaveBeenCalled(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + }); + + it("does not block accounts when quota reset times are known but already elapsed", async () => { + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "acc-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now - 5 * 60_000, + }, + secondary: { + usedPercent: 100, + resetAtMs: now - 60_000, }, - secondary: {}, }, }, byEmail: {}, @@ -1462,12 +1954,6 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expires: Date.now() + 60_000, })); const markRateLimitSpy = vi.spyOn(AccountManager.prototype, "markRateLimitedWithReason"); - vi.mocked(fetchHelpers.transformRequestForCodex).mockImplementationOnce( - async (init: unknown) => ({ - updatedInit: init, - body: { model: "gpt-5.1-mini" }, - }), - ); vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( (_init, _accountId, accessToken) => new Headers({ "x-test-access-token": String(accessToken ?? "") }), @@ -1479,7 +1965,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { const { sdk } = await setupPlugin(); const response = await sdk.fetch!("https://api.openai.com/v1/chat", { method: "POST", - body: JSON.stringify({ model: "gpt-5.1-mini" }), + body: JSON.stringify({ model: "gpt-5.1" }), }); expect(response.status).toBe(200); @@ -1837,10 +2323,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { it("continues request flow when quota bootstrap cache reads hit EPERM/EBUSY", async () => { const fetchHelpers = await import("../lib/request/fetch-helpers.js"); const epermError = Object.assign(new Error("locked"), { code: "EPERM" }); - const ebusyError = Object.assign(new Error("busy"), { code: "EBUSY" }); - loadQuotaCacheMock - .mockRejectedValueOnce(epermError) - .mockRejectedValueOnce(ebusyError); + loadQuotaCacheMock.mockRejectedValueOnce(epermError); vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( (_init, _accountId, accessToken) => new Headers({ "x-test-access-token": String(accessToken ?? "") }), @@ -1861,7 +2344,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(firstResponse.status).toBe(200); expect(secondResponse.status).toBe(200); - expect(loadQuotaCacheMock).toHaveBeenCalledTimes(2); + expect(loadQuotaCacheMock).toHaveBeenCalledTimes(1); expect(globalThis.fetch).toHaveBeenCalledTimes(2); const firstHeaders = new Headers( (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, @@ -1873,6 +2356,159 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(secondHeaders.get("x-test-access-token")).toBe("access-token"); }); + it("retries quota bootstrap loads after the failure retry interval and recovers within cooldown", async () => { + vi.useFakeTimers(); + const now = Date.now(); + const { AccountManager } = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const accountOne = { + index: 0, + accountId: "acc-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "acc-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + const epermError = Object.assign(new Error("locked"), { code: "EPERM" }); + loadQuotaCacheMock + .mockRejectedValueOnce(epermError) + .mockResolvedValueOnce({ + byAccountId: {}, + byEmail: { + "user1@example.com": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + "user@example.com": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + }); + const countSpy = vi + .spyOn(AccountManager.prototype, "getAccountCount") + .mockReturnValue(2); + const selectionSpy = vi + .spyOn(AccountManager.prototype, "getCurrentOrNextForFamilyHybrid") + .mockImplementation((family: string, model?: string | null) => { + const nowMs = Date.now(); + const accounts = [accountOne, accountTwo]; + const modelKey = model ? `${family}:${model}` : family; + for (const candidate of accounts) { + const resets = candidate.rateLimitResetTimes ?? {}; + const blockedUntil = Math.max( + resets[family] ?? 0, + resets[modelKey] ?? 0, + ); + if (blockedUntil <= nowMs) { + return candidate as never; + } + } + return null as never; + }); + const snapshotSpy = vi + .spyOn(AccountManager.prototype, "getAccountsSnapshot") + .mockReturnValue([accountOne, accountTwo] as never); + const getByIndexSpy = vi + .spyOn(AccountManager.prototype, "getAccountByIndex") + .mockImplementation((index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + const toAuthSpy = vi + .spyOn(AccountManager.prototype, "toAuthDetails") + .mockImplementation((account: { accountId?: string }) => ({ + type: "oauth" as const, + access: `access-${account.accountId ?? "unknown"}`, + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + })); + const markRateLimitSpy = vi + .spyOn(AccountManager.prototype, "markRateLimitedWithReason") + .mockImplementation( + ( + account: { rateLimitResetTimes?: Record }, + waitMs: number, + family: string, + _reason?: string, + model?: string | null, + ) => { + const resetAt = Date.now() + waitMs; + account.rateLimitResetTimes = account.rateLimitResetTimes ?? {}; + account.rateLimitResetTimes[family] = resetAt; + if (model) { + account.rateLimitResetTimes[`${family}:${model}`] = resetAt; + } + }, + ); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + try { + const { sdk } = await setupPlugin(); + const firstResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + const secondResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + await vi.advanceTimersByTimeAsync(5_100); + const thirdResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(firstResponse.status).toBe(200); + expect(secondResponse.status).toBe(200); + expect(thirdResponse.status).toBe(200); + expect(loadQuotaCacheMock).toHaveBeenCalledTimes(2); + expect(globalThis.fetch).toHaveBeenCalledTimes(3); + const firstHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit)?.headers, + ); + const secondHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[1]?.[1] as RequestInit)?.headers, + ); + const thirdHeaders = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[2]?.[1] as RequestInit)?.headers, + ); + expect(firstHeaders.get("x-test-access-token")).not.toBe("access-acc-2"); + expect(secondHeaders.get("x-test-access-token")).not.toBe("access-acc-2"); + expect(thirdHeaders.get("x-test-access-token")).toBeTruthy(); + } finally { + vi.useRealTimers(); + countSpy.mockRestore(); + selectionSpy.mockRestore(); + snapshotSpy.mockRestore(); + getByIndexSpy.mockRestore(); + toAuthSpy.mockRestore(); + markRateLimitSpy.mockRestore(); + } + }); + it("disables quota cache bootstrap when CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP=0", async () => { const previous = process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP; process.env.CODEX_AUTH_QUOTA_CACHE_BOOTSTRAP = "0"; From 193e6123438b6df777d6c6e3e91841acc2e80747 Mon Sep 17 00:00:00 2001 From: ndycode Date: Fri, 6 Mar 2026 00:34:22 +0800 Subject: [PATCH 6/8] fix: unify preemptive quota scheduler key normalization - use shared quota schedule key builder in bootstrap, runtime deferral, and stream fallback updates - add regression in index-retry for provider-prefixed model-id deferral matching - tighten recovery assertion to verify rotation target account deterministically Co-authored-by: Codex --- index.ts | 29 +++++++++-- test/index-retry.test.ts | 106 +++++++++++++++++++++++++++++++++++++++ test/index.test.ts | 2 +- 3 files changed, 133 insertions(+), 4 deletions(-) diff --git a/index.ts b/index.ts index c267e25b0..b7e438f03 100644 --- a/index.ts +++ b/index.ts @@ -341,6 +341,17 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { return normalized.length > 0 ? normalized : null; }; + const buildQuotaScheduleKey = ( + account: { accountId?: string; email?: string } | string, + model: string | undefined, + modelFamily: ModelFamily, + ): string => { + const accountKey = + typeof account === "string" ? account : resolveEntitlementAccountKey(account); + const normalizedModel = normalizeQuotaCacheModelId(model) ?? modelFamily; + return `${accountKey}:${normalizedModel}`; + }; + const hasQuotaCacheEntries = (cache: QuotaCacheData): boolean => Object.keys(cache.byAccountId).length > 0 || Object.keys(cache.byEmail).length > 0; @@ -534,7 +545,11 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { "quota", appliedModel, ); - const quotaScheduleKey = `${resolveEntitlementAccountKey(snapshotCandidate)}:${appliedModel}`; + const quotaScheduleKey = buildQuotaScheduleKey( + snapshotCandidate, + appliedModel, + modelFamily, + ); preemptiveQuotaScheduler.update( quotaScheduleKey, toQuotaSchedulerSnapshot(entry), @@ -1855,7 +1870,11 @@ while (attempted.size < Math.max(1, accountCount)) { promptCacheKey: effectivePromptCacheKey, }, ); - const quotaScheduleKey = `${entitlementAccountKey}:${model ?? modelFamily}`; + const quotaScheduleKey = buildQuotaScheduleKey( + entitlementAccountKey, + model, + modelFamily, + ); const capabilityModelKey = model ?? modelFamily; const quotaDeferral = preemptiveQuotaScheduler.getDeferral(quotaScheduleKey); if (quotaDeferral.defer && quotaDeferral.waitMs > 0) { @@ -2428,7 +2447,11 @@ while (attempted.size < Math.max(1, accountCount)) { ); if (fallbackSnapshot) { preemptiveQuotaScheduler.update( - `${resolveEntitlementAccountKey(fallbackAccount)}:${model ?? modelFamily}`, + buildQuotaScheduleKey( + fallbackAccount, + model, + modelFamily, + ), fallbackSnapshot, ); } diff --git a/test/index-retry.test.ts b/test/index-retry.test.ts index f5ee4a063..738478f6a 100644 --- a/test/index-retry.test.ts +++ b/test/index-retry.test.ts @@ -223,5 +223,111 @@ describe("OpenAIAuthPlugin rate-limit retry", () => { expect(globalThis.fetch).toHaveBeenCalledTimes(1); expect(response.status).toBe(200); }); + + it("uses normalized quota scheduler keys for provider-prefixed model ids", async () => { + const { OpenAIAuthPlugin } = await import("../index.js"); + const accountsModule = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const quotaCacheModule = await import("../lib/quota-cache.js"); + + const loadQuotaCacheMock = vi.mocked(quotaCacheModule.loadQuotaCache); + const now = Date.now(); + const accountOne = { + index: 0, + accountId: "account-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "account-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "account-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + + vi.spyOn(accountsModule.AccountManager.prototype, "getAccountCount").mockReturnValue(2); + vi.spyOn(accountsModule.AccountManager.prototype, "getAccountByIndex").mockImplementation( + (index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + let selectionCallCount = 0; + vi.spyOn( + accountsModule.AccountManager.prototype, + "getCurrentOrNextForFamilyHybrid", + ).mockImplementation( + () => { + selectionCallCount += 1; + return (selectionCallCount === 1 ? accountOne : accountTwo) as never; + }, + ); + vi.spyOn(accountsModule.AccountManager.prototype, "toAuthDetails").mockImplementation( + (account: { index?: number; accountId?: string }) => ({ + type: "oauth", + access: account.index === 0 ? "access-acc-1" : "access-acc-2", + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + }), + ); + const markRateLimitSpy = vi.spyOn( + accountsModule.AccountManager.prototype, + "markRateLimitedWithReason", + ); + vi.spyOn(fetchHelpers, "transformRequestForCodex").mockImplementationOnce( + async (init: unknown) => ({ + updatedInit: init, + body: { model: "OPENAI/GPT-5.1" }, + }), + ); + vi.spyOn(fetchHelpers, "createCodexHeaders").mockImplementation( + (_init: unknown, _accountId: unknown, accessToken: unknown) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + + const client = { + tui: { showToast: vi.fn() }, + auth: { set: vi.fn() }, + } as any; + const plugin = await OpenAIAuthPlugin({ client }); + + const getAuth = async () => ({ + type: "oauth" as const, + access: "a", + refresh: "r", + expires: Date.now() + 60_000, + multiAccount: true, + }); + + const sdk = (await plugin.auth.loader(getAuth, { options: {}, models: {} })) as any; + const response = await sdk.fetch("https://example.com", { + method: "POST", + body: JSON.stringify({ model: "OPENAI/GPT-5.1" }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit | undefined)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-2"); + expect(markRateLimitSpy).toHaveBeenCalled(); + }); }); diff --git a/test/index.test.ts b/test/index.test.ts index 24ac06a76..bd0e7dacc 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -2497,7 +2497,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { ); expect(firstHeaders.get("x-test-access-token")).not.toBe("access-acc-2"); expect(secondHeaders.get("x-test-access-token")).not.toBe("access-acc-2"); - expect(thirdHeaders.get("x-test-access-token")).toBeTruthy(); + expect(thirdHeaders.get("x-test-access-token")).toBe("access-acc-2"); } finally { vi.useRealTimers(); countSpy.mockRestore(); From 2e02fca6c462238aa90d1ad7f29a17a117467d66 Mon Sep 17 00:00:00 2001 From: ndycode Date: Fri, 6 Mar 2026 00:45:52 +0800 Subject: [PATCH 7/8] fix: enforce quota bootstrap cooldown and cache empty loads Co-authored-by: Codex --- index.ts | 28 ++++------------------------ test/index.test.ts | 31 +++++++++++++++++++++++++++++-- 2 files changed, 33 insertions(+), 26 deletions(-) diff --git a/index.ts b/index.ts index b7e438f03..6454a2ef1 100644 --- a/index.ts +++ b/index.ts @@ -309,11 +309,9 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { const QUOTA_CACHE_BOOTSTRAP_TTL_MS = 30 * 60_000; const QUOTA_CACHE_BOOTSTRAP_MIN_WAIT_MS = 1_000; const QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS = 60_000; - const QUOTA_CACHE_BOOTSTRAP_FAILURE_RETRY_INTERVAL_MS = 5_000; let quotaCacheBootstrapData: QuotaCacheData | null = null; let quotaCacheBootstrapLoadedAt = 0; let quotaCacheBootstrapFailureUntil = 0; - let quotaCacheBootstrapFailureRetryAt = 0; let quotaCacheBootstrapLoadPromise: Promise | null = null; type QuotaBootstrapAccountCandidate = { @@ -352,9 +350,6 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { return `${accountKey}:${normalizedModel}`; }; - const hasQuotaCacheEntries = (cache: QuotaCacheData): boolean => - Object.keys(cache.byAccountId).length > 0 || Object.keys(cache.byEmail).length > 0; - const createEmptyQuotaCache = (): QuotaCacheData => ({ byAccountId: {}, byEmail: {} }); const getQuotaCacheEntryForCandidate = ( @@ -444,10 +439,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (quotaCacheBootstrapData && now - quotaCacheBootstrapLoadedAt < QUOTA_CACHE_BOOTSTRAP_TTL_MS) { return quotaCacheBootstrapData; } - if ( - now < quotaCacheBootstrapFailureUntil && - now < quotaCacheBootstrapFailureRetryAt - ) { + if (now < quotaCacheBootstrapFailureUntil) { return createEmptyQuotaCache(); } if (quotaCacheBootstrapLoadPromise) { @@ -456,27 +448,15 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { quotaCacheBootstrapLoadPromise = (async () => { try { const loaded = await loadQuotaCache(); - if (hasQuotaCacheEntries(loaded)) { - quotaCacheBootstrapData = loaded; - quotaCacheBootstrapLoadedAt = Date.now(); - quotaCacheBootstrapFailureUntil = 0; - quotaCacheBootstrapFailureRetryAt = 0; - } else { - const failureAt = Date.now(); - quotaCacheBootstrapData = null; - quotaCacheBootstrapLoadedAt = 0; - quotaCacheBootstrapFailureUntil = failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS; - quotaCacheBootstrapFailureRetryAt = - failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_RETRY_INTERVAL_MS; - } + quotaCacheBootstrapData = loaded; + quotaCacheBootstrapLoadedAt = Date.now(); + quotaCacheBootstrapFailureUntil = 0; return loaded; } catch (error) { const failureAt = Date.now(); quotaCacheBootstrapData = null; quotaCacheBootstrapLoadedAt = 0; quotaCacheBootstrapFailureUntil = failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS; - quotaCacheBootstrapFailureRetryAt = - failureAt + QUOTA_CACHE_BOOTSTRAP_FAILURE_RETRY_INTERVAL_MS; logWarn( `[${PLUGIN_NAME}] failed to load quota cache bootstrap data: ${ error instanceof Error ? error.message : String(error) diff --git a/test/index.test.ts b/test/index.test.ts index bd0e7dacc..a6ec6a16c 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -2147,6 +2147,33 @@ describe("OpenAIOAuthPlugin fetch handler", () => { } }); + it("caches valid empty quota bootstrap loads within ttl", async () => { + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + loadQuotaCacheMock.mockResolvedValueOnce({ byAccountId: {}, byEmail: {} }); + vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation( + (_init, _accountId, accessToken) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "ok" }), { status: 200 }), + ); + + const { sdk } = await setupPlugin(); + const firstResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + const secondResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(firstResponse.status).toBe(200); + expect(secondResponse.status).toBe(200); + expect(loadQuotaCacheMock).toHaveBeenCalledTimes(1); + expect(globalThis.fetch).toHaveBeenCalledTimes(2); + }); + it("handles token refresh races while rotating away from a bootstrap-rate-limited account", async () => { const now = Date.now() + 31 * 60_000; const nowSpy = vi.spyOn(Date, "now").mockReturnValue(now); @@ -2356,7 +2383,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(secondHeaders.get("x-test-access-token")).toBe("access-token"); }); - it("retries quota bootstrap loads after the failure retry interval and recovers within cooldown", async () => { + it("retries quota bootstrap loads after the failure cooldown window and recovers", async () => { vi.useFakeTimers(); const now = Date.now(); const { AccountManager } = await import("../lib/accounts.js"); @@ -2475,7 +2502,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => { method: "POST", body: JSON.stringify({ model: "gpt-5.1" }), }); - await vi.advanceTimersByTimeAsync(5_100); + await vi.advanceTimersByTimeAsync(60_100); const thirdResponse = await sdk.fetch!("https://api.openai.com/v1/chat", { method: "POST", body: JSON.stringify({ model: "gpt-5.1" }), From 2d204ca7ebb230241def68e19528b3b69f547fdc Mon Sep 17 00:00:00 2001 From: ndycode Date: Fri, 6 Mar 2026 01:30:52 +0800 Subject: [PATCH 8/8] fix: align quota bootstrap key for model-less requests Co-authored-by: Codex --- index.ts | 2 +- test/index-retry.test.ts | 104 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 1 deletion(-) diff --git a/index.ts b/index.ts index 6454a2ef1..c82da853e 100644 --- a/index.ts +++ b/index.ts @@ -527,7 +527,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { ); const quotaScheduleKey = buildQuotaScheduleKey( snapshotCandidate, - appliedModel, + requestedModel ?? undefined, modelFamily, ); preemptiveQuotaScheduler.update( diff --git a/test/index-retry.test.ts b/test/index-retry.test.ts index 738478f6a..6cedaaf56 100644 --- a/test/index-retry.test.ts +++ b/test/index-retry.test.ts @@ -329,5 +329,109 @@ describe("OpenAIAuthPlugin rate-limit retry", () => { expect(headers.get("x-test-access-token")).toBe("access-acc-2"); expect(markRateLimitSpy).toHaveBeenCalled(); }); + + it("uses family-level scheduler keys when request model is omitted", async () => { + const { OpenAIAuthPlugin } = await import("../index.js"); + const accountsModule = await import("../lib/accounts.js"); + const fetchHelpers = await import("../lib/request/fetch-helpers.js"); + const quotaCacheModule = await import("../lib/quota-cache.js"); + + const loadQuotaCacheMock = vi.mocked(quotaCacheModule.loadQuotaCache); + const now = Date.now(); + const accountOne = { + index: 0, + accountId: "account-1", + email: "user1@example.com", + refreshToken: "refresh-1", + rateLimitResetTimes: {}, + }; + const accountTwo = { + index: 1, + accountId: "account-2", + email: "user2@example.com", + refreshToken: "refresh-2", + rateLimitResetTimes: {}, + }; + + loadQuotaCacheMock.mockResolvedValueOnce({ + byAccountId: { + "account-1": { + updatedAt: now, + status: 429, + model: "gpt-5.1", + primary: { + usedPercent: 100, + resetAtMs: now + 5 * 60_000, + }, + secondary: {}, + }, + }, + byEmail: {}, + }); + + vi.spyOn(accountsModule.AccountManager.prototype, "getAccountCount").mockReturnValue(2); + vi.spyOn(accountsModule.AccountManager.prototype, "getAccountByIndex").mockImplementation( + (index: number) => + (index === 0 ? accountOne : index === 1 ? accountTwo : null) as never, + ); + let selectionCallCount = 0; + vi.spyOn( + accountsModule.AccountManager.prototype, + "getCurrentOrNextForFamilyHybrid", + ).mockImplementation(() => { + selectionCallCount += 1; + return (selectionCallCount === 1 ? accountOne : accountTwo) as never; + }); + vi.spyOn(accountsModule.AccountManager.prototype, "toAuthDetails").mockImplementation( + (account: { index?: number; accountId?: string }) => ({ + type: "oauth", + access: account.index === 0 ? "access-acc-1" : "access-acc-2", + refresh: `refresh-${account.accountId ?? "unknown"}`, + expires: Date.now() + 60_000, + }), + ); + const markRateLimitSpy = vi.spyOn( + accountsModule.AccountManager.prototype, + "markRateLimitedWithReason", + ); + vi.spyOn(fetchHelpers, "transformRequestForCodex").mockImplementationOnce( + async (init: unknown) => ({ + updatedInit: init, + body: {}, + }), + ); + vi.spyOn(fetchHelpers, "createCodexHeaders").mockImplementation( + (_init: unknown, _accountId: unknown, accessToken: unknown) => + new Headers({ "x-test-access-token": String(accessToken ?? "") }), + ); + + const client = { + tui: { showToast: vi.fn() }, + auth: { set: vi.fn() }, + } as any; + const plugin = await OpenAIAuthPlugin({ client }); + + const getAuth = async () => ({ + type: "oauth" as const, + access: "a", + refresh: "r", + expires: Date.now() + 60_000, + multiAccount: true, + }); + + const sdk = (await plugin.auth.loader(getAuth, { options: {}, models: {} })) as any; + const response = await sdk.fetch("https://example.com", { + method: "POST", + body: JSON.stringify({ messages: [{ role: "user", content: "hello" }] }), + }); + + expect(response.status).toBe(200); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + const headers = new Headers( + (vi.mocked(globalThis.fetch).mock.calls[0]?.[1] as RequestInit | undefined)?.headers, + ); + expect(headers.get("x-test-access-token")).toBe("access-acc-2"); + expect(markRateLimitSpy).toHaveBeenCalled(); + }); });