Fix stale Codex usage after reset (#21748)

* fix(rate-limits): refresh Codex usage after reset

* fix(rate-limits): converge weekly Codex reset usage

---------

Co-authored-by: m4air <m4air@Mac.localdomain>

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'

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) {

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<void> {
+  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<RateLimitState> {
-    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)