diff --git a/src/main/rate-limits/service-account-target-selection.test.ts b/src/main/rate-limits/service-account-target-selection.test.ts index 9241831bdf6..7f9174f5d1f 100644 --- a/src/main/rate-limits/service-account-target-selection.test.ts +++ b/src/main/rate-limits/service-account-target-selection.test.ts @@ -146,6 +146,50 @@ describe('RateLimitService', () => { }) }) + it('retries the post-reset usage read when the first response still shows the old quota', async () => { + const service = new RateLimitService() + service.setCodexHomePathResolver(() => ({ kind: 'ready', codexHomePath: '/tmp/codex-home' })) + vi.mocked(consumeCodexRateLimitResetCredit).mockResolvedValueOnce('reset') + vi.mocked(fetchCodexRateLimits) + .mockResolvedValueOnce(okProvider('codex', 100, Date.now())) + .mockResolvedValueOnce(okProvider('codex', 0, Date.now())) + + const result = await service.consumeCodexRateLimitResetCredit({ + idempotencyKey: '55555555-5555-4555-8555-555555555555', + target: { runtime: 'host', wslDistro: null }, + codexHomePath: '/tmp/codex-home' + }) + + expect(result.state.codex?.session?.usedPercent).toBe(0) + expect(fetchCodexRateLimits).toHaveBeenCalledTimes(2) + }) + + it('retries when a weekly-only quota is still stale after reset', async () => { + const service = new RateLimitService() + service.setCodexHomePathResolver(() => ({ kind: 'ready', codexHomePath: '/tmp/codex-home' })) + vi.mocked(consumeCodexRateLimitResetCredit).mockResolvedValueOnce('reset') + vi.mocked(fetchCodexRateLimits) + .mockResolvedValueOnce({ + ...okProvider('codex', 0, Date.now()), + session: null, + weekly: { ...okProvider('codex', 100, Date.now()).session! } + }) + .mockResolvedValueOnce({ + ...okProvider('codex', 0, Date.now()), + session: null, + weekly: { ...okProvider('codex', 0, Date.now()).session! } + }) + + const result = await service.consumeCodexRateLimitResetCredit({ + idempotencyKey: '66666666-6666-4666-8666-666666666666', + target: { runtime: 'host', wslDistro: null }, + codexHomePath: '/tmp/codex-home' + }) + + expect(result.state.codex?.weekly?.usedPercent).toBe(0) + expect(fetchCodexRateLimits).toHaveBeenCalledTimes(2) + }) + it('returns a refreshed scoped state without overwriting a target selected during reset', async () => { const service = new RateLimitService() const idempotencyKey = '22222222-2222-4222-8222-222222222222' diff --git a/src/main/rate-limits/service/service-account-refresh.ts b/src/main/rate-limits/service/service-account-refresh.ts index 9e758a98518..ed8578e3694 100644 --- a/src/main/rate-limits/service/service-account-refresh.ts +++ b/src/main/rate-limits/service/service-account-refresh.ts @@ -110,7 +110,8 @@ export abstract class RateLimitServiceAccountRefresh extends RateLimitServiceIna const state = await this.fetchCodexResetResultState( codexTarget, codexHomePath, - scopedStateBeforeReset + scopedStateBeforeReset, + outcome ) return { outcome, state } } catch (error) { diff --git a/src/main/rate-limits/service/service-fetch-targets.ts b/src/main/rate-limits/service/service-fetch-targets.ts index 09274ad3d82..51ce2db3d0b 100644 --- a/src/main/rate-limits/service/service-fetch-targets.ts +++ b/src/main/rate-limits/service/service-fetch-targets.ts @@ -12,6 +12,33 @@ import { type RateLimitState, toErrorMessage } from './service-types' +import type { CodexRateLimitResetOutcome } from '../../../shared/rate-limit-types' + +const CODEX_RESET_REFRESH_RETRIES = 3 +const CODEX_RESET_REFRESH_DELAY_MS = 250 + +function waitForCodexResetRefresh(): Promise { + return new Promise((resolve) => setTimeout(resolve, CODEX_RESET_REFRESH_DELAY_MS)) +} + +function codexResetUsageVisible( + fresh: ProviderRateLimits, + previous: ProviderRateLimits | null +): boolean { + if (fresh.status !== 'ok') { + return false + } + if (!previous) { + return (fresh.session?.usedPercent ?? 0) <= 0 && (fresh.weekly?.usedPercent ?? 0) <= 0 + } + const sessionImproved = + fresh.session !== null && + (previous.session === null || fresh.session.usedPercent < previous.session.usedPercent) + const weeklyImproved = + fresh.weekly !== null && + (previous.weekly === null || fresh.weekly.usedPercent < previous.weekly.usedPercent) + return sessionImproved || weeklyImproved +} export abstract class RateLimitServiceFetchTargets extends RateLimitServiceResultPolicy { protected resolveCodexHome(target?: CodexAccountSelectionTarget): { @@ -77,27 +104,47 @@ export abstract class RateLimitServiceFetchTargets extends RateLimitServiceResul protected async fetchCodexResetResultState( target: NormalizedCodexAccountSelectionTarget, codexHomePath: string | null, - stateBeforeReset: RateLimitState + stateBeforeReset: RateLimitState, + outcome: CodexRateLimitResetOutcome ): Promise { - const controller = this.beginFetchCycle() - let fresh: ProviderRateLimits - try { - fresh = await fetchCodexRateLimits({ - codexHomePath, - allowPtyFallback: this.shouldAllowCodexPtyFallback(), - signal: controller.signal - }) - } catch (error) { - fresh = { - provider: 'codex', - session: null, - weekly: null, - updatedAt: Date.now(), - error: toErrorMessage(error), - status: 'error' + let fresh: ProviderRateLimits = stateBeforeReset.codex ?? { + provider: 'codex', + session: null, + weekly: null, + updatedAt: 0, + error: null, + status: 'fetching' + } + for ( + let attempt = 0; + attempt <= (outcome === 'reset' ? CODEX_RESET_REFRESH_RETRIES : 0); + attempt += 1 + ) { + if (attempt > 0) { + await waitForCodexResetRefresh() + } + const controller = this.beginFetchCycle() + try { + fresh = await fetchCodexRateLimits({ + codexHomePath, + allowPtyFallback: this.shouldAllowCodexPtyFallback(), + signal: controller.signal + }) + } catch (error) { + fresh = { + provider: 'codex', + session: null, + weekly: null, + updatedAt: Date.now(), + error: toErrorMessage(error), + status: 'error' + } + } finally { + this.finishFetchCycle(controller) + } + if (outcome !== 'reset' || codexResetUsageVisible(fresh, stateBeforeReset.codex)) { + break } - } finally { - this.finishFetchCycle(controller) } const scopedCodex = this.applyStalePolicy(fresh, stateBeforeReset.codex)