Skip to content
15 changes: 14 additions & 1 deletion src/providers/quota.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1377,6 +1377,7 @@ async function fetchAnthropicQuota(provider: string): Promise<ProviderQuotaRepor
const stillOwnsToken = getAccountCredential("anthropic", probedAccountId)?.access === accessToken;
if (stillOwnsToken && mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) {
accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota });
persistAccountQuotaCache();
}
}
return report(provider, "anthropic:oauth-usage", quota);
Expand Down Expand Up @@ -1404,6 +1405,7 @@ async function fetchKiroQuota(provider: string): Promise<ProviderQuotaReport | n
if (mayCommitAccountQuotaKey(probedAccountKey, writerGeneration)) {
accountQuotaCache.set(probedAccountKey, { ts: Date.now(), quota: snapshot.quota });
commitKiroAccountUsageState(probedAccountKey, snapshot);
persistAccountQuotaCache();
}
return report(provider, "kiro:usage-limits", snapshot.quota);
}
Expand Down Expand Up @@ -1486,6 +1488,7 @@ function accountCacheKey(provider: string, accountId: string): string {
* Returns null when nothing is cached (or the cached row has no bars).
*/
export function getCachedProviderAccountQuota(provider: string, accountId: string): ProviderQuota | null {
hydrateAccountQuotaCache();
const entry = accountQuotaCache.get(accountCacheKey(provider, accountId));
return entry?.quota ?? null;
}
Expand Down Expand Up @@ -1516,6 +1519,7 @@ export function sweepExpiredProviderAccountQuotaRows(now = Date.now()): number {

export function reconcileProviderAccountQuotaRows(context: GenerationContext): number {
if (context.generation <= lastReconciledGeneration) return 0;
hydrateAccountQuotaCache();
let removed = 0;
for (const key of accountQuotaCache.keys()) {
if (context.oauthAccountKeys.has(key)) continue;
Expand Down Expand Up @@ -1544,6 +1548,13 @@ export function resetProviderQuotaReconcileStateForTests(): void {
liveProviderQuotaKeys = new Set();
}

/** Test-only reset that models a fresh process loading the disk snapshot. */
export function resetProviderAccountQuotaDiskStateForTests(): void {
accountQuotaCache.clear();
diskHydrated = false;
cancelPendingAccountQuotaPersist();
}

/** Drop cached per-account rows (all, or just one provider's). */
export function clearAccountQuotaCache(provider?: string): void {
if (!provider) {
Expand All @@ -1552,7 +1563,7 @@ export function clearAccountQuotaCache(provider?: string): void {
clearKiroAccountUsageState();
// A cleared cache must not be re-seeded from the file it was just cleared of, and any
// pending write of the old rows is abandoned.
diskHydrated = false;
diskHydrated = true;
cancelPendingAccountQuotaPersist();
return;
}
Expand Down Expand Up @@ -1595,6 +1606,7 @@ async function fetchAccountQuota(
accountId: string,
forceRefresh: boolean,
): Promise<AccountQuotaCacheEntry> {
hydrateAccountQuotaCache();
const key = accountCacheKey(provider, accountId);
const writerGeneration = captureConfigGeneration();
const cached = accountQuotaCache.get(key);
Expand Down Expand Up @@ -1641,6 +1653,7 @@ async function fetchAccountQuota(
// superseded config generation must not publish either half.
if (provider === "kiro") commitKiroAccountUsageState(key, kiroSnapshot);
sweepExpiredOnWrite(entry.ts);
persistAccountQuotaCache();
}
return entry;
} catch {
Expand Down
13 changes: 13 additions & 0 deletions tests/provider-account-quota-persistence.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ import {
readPersistedAccountQuotas,
schedulePersistAccountQuotas,
} from "../src/providers/account-quota-disk";
import {
getCachedProviderAccountQuota,
resetProviderAccountQuotaDiskStateForTests,
} from "../src/providers/quota";
import type { ProviderQuota } from "../src/providers/quota-types";

const previousHome = process.env.OPENCODEX_HOME;
Expand All @@ -22,6 +26,7 @@ beforeEach(() => {
});

afterEach(() => {
resetProviderAccountQuotaDiskStateForTests();
cancelPendingAccountQuotaPersist();
if (previousHome === undefined) delete process.env.OPENCODEX_HOME;
else process.env.OPENCODEX_HOME = previousHome;
Expand All @@ -40,6 +45,14 @@ describe("provider account quota persistence", () => {
expect(readPersistedAccountQuotas().get(KEY)?.monthlyPercent).toBe(15);
});

test("routing hydrates the runtime cache from a persisted snapshot", async () => {
schedulePersistAccountQuotas(() => [[KEY, quota(95)]]);
await settle();

resetProviderAccountQuotaDiskStateForTests();
expect(getCachedProviderAccountQuota("kiro", "acct-a")?.monthlyPercent).toBe(95);
});

test("a stale snapshot is discarded rather than ordering the pool on old data", async () => {
schedulePersistAccountQuotas(() => [[KEY, quota(15, Date.now() - 7 * 60 * 60_000)]]);
await settle();
Expand Down
5 changes: 5 additions & 0 deletions tests/provider-account-quota.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { saveCredential } from "../src/oauth/store";
import { readPersistedAccountQuotas } from "../src/providers/account-quota-disk";
import type { OcxConfig } from "../src/types";
import {
clearAccountQuotaCache,
Expand Down Expand Up @@ -76,6 +77,10 @@ describe("fetchProviderAccountQuotas", () => {
expect(seenTokens.sort()).toEqual(["Bearer token-first", "Bearer token-second"]);
// The 5-hour window lands in the canonical fields, not in customWindows.
for (const row of rows) expect(row.quota?.customWindows).toBeUndefined();

await new Promise(resolve => setTimeout(resolve, 400));
const persisted = readPersistedAccountQuotas();
expect(rows.map(row => persisted.get(`anthropic\u0000${row.accountId}`)?.fiveHourPercent).sort()).toEqual([3, 70]);
});

test("a cached row is reused instead of re-probing upstream", async () => {
Expand Down
Loading