diff --git a/index.ts b/index.ts index 7db88088a..c82da853e 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,237 @@ 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; + const QUOTA_CACHE_BOOTSTRAP_FAILURE_COOLDOWN_MS = 60_000; + let quotaCacheBootstrapData: QuotaCacheData | null = null; + let quotaCacheBootstrapLoadedAt = 0; + let quotaCacheBootstrapFailureUntil = 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 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 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 createEmptyQuotaCache = (): QuotaCacheData => ({ byAccountId: {}, byEmail: {} }); + + 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.max(...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 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; + } + 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 (now < quotaCacheBootstrapFailureUntil) { + return createEmptyQuotaCache(); + } + if (quotaCacheBootstrapLoadPromise) { + return quotaCacheBootstrapLoadPromise; + } + quotaCacheBootstrapLoadPromise = (async () => { + try { + const loaded = await loadQuotaCache(); + 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; + logWarn( + `[${PLUGIN_NAME}] failed to load quota cache bootstrap data: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return createEmptyQuotaCache(); + } + })(); + 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 requestedModel = normalizeQuotaCacheModelId(model); + const baseKey = getQuotaKey(modelFamily); + + for (const snapshotCandidate of accountSnapshots) { + const entry = getQuotaCacheEntryForCandidate(cache, snapshotCandidate); + if (!entry) continue; + const entryModel = normalizeQuotaCacheModelId(entry.model); + if (!entryModel) continue; + 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, + 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[appliedModelKey] ?? 0; + const existingResetAt = Math.max(existingBaseResetAt, existingModelResetAt); + const nextResetAt = now + waitMs; + if (existingResetAt >= nextResetAt) { + continue; + } + + accountManager.markRateLimitedWithReason( + account, + waitMs, + modelFamily, + "quota", + appliedModel, + ); + const quotaScheduleKey = buildQuotaScheduleKey( + snapshotCandidate, + requestedModel ?? undefined, + modelFamily, + ); + preemptiveQuotaScheduler.update( + quotaScheduleKey, + toQuotaSchedulerSnapshot(entry), + ); + } + }; + const sanitizeResponseHeadersForLog = (headers: Headers): Record => { const allowed = new Set([ "content-type", @@ -1135,6 +1372,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 +1678,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( @@ -1604,7 +1850,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) { @@ -2005,83 +2255,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; @@ -2197,7 +2427,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 3813edc48..6cedaaf56 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, }; @@ -213,5 +223,215 @@ 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(); + }); + + 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(); + }); }); diff --git a/test/index.test.ts b/test/index.test.ts index d6d95497f..a6ec6a16c 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, @@ -204,7 +209,7 @@ vi.mock("../lib/request/rate-limit-backoff.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 })), @@ -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", @@ -1051,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 }), @@ -1169,6 +1191,1440 @@ 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 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("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"); + 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: "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, + }, + }, + }, + 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.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("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"); + 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 + 5 * 60_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-2"); + 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 }>(); + 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 ?? "") }), + ); + 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 bootstrapLoadStarted.promise; + + 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("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); + 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; + }>(); + const firstRefreshStarted = createDeferred(); + 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; + firstRefreshStarted.resolve(); + 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 firstRefreshStarted.promise; + 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" }); + loadQuotaCacheMock.mockRejectedValueOnce(epermError); + 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); + 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("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"); + 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(60_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")).toBe("access-acc-2"); + } 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"; + 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");