fix(quota): coalesce provider usage refreshes

2026-08-18 19:13:58 +03:00
parent 84e940a9e4
commit 3326ad0ea8
7 changed files with 267 additions and 39 deletions
@@ -20,6 +20,7 @@ const AUTH = JSON.stringify({
  'command-code': { type: 'oauth', access: 'test-token' },
  'zai-coding-plan': { key: 'test-token' },
  deepseek: { key: 'test-token' },
+  anthropic: { access: 'test-token', refresh: 'test-refresh' },
 });
 ((fs as unknown) as { existsSync: () => boolean }).existsSync = () => true;
 ((fs as unknown) as { readFileSync: () => string }).readFileSync = () => AUTH;
@@ -202,6 +203,27 @@ describe('Crof quota provider (VS Code parity)', () => {
 });

 describe('Codex quota provider (VS Code parity)', () => {
+  test('coalesces concurrent refreshes for the same provider', async () => {
+    let resolveResponse: ((response: Response) => void) | undefined;
+    let requestCount = 0;
+    globalThis.fetch = (() => {
+      requestCount += 1;
+      return new Promise<Response>((resolve) => {
+        resolveResponse = resolve;
+      });
+    }) as typeof fetch;
+
+    const first = fetchQuotaForProvider('codex');
+    const second = fetchQuotaForProvider('codex');
+    resolveResponse?.(mockResponse({ rate_limit: null }));
+
+    const [firstResult, secondResult] = await Promise.all([first, second]);
+
+    assert.equal(firstResult.ok, true);
+    assert.equal(secondResult.ok, true);
+    assert.equal(requestCount, 1);
+  });
+
  test('surfaces spend_control individual limit for business accounts', async () => {
    stubFetchReturning(() => Promise.resolve(mockResponse({
      plan_type: 'business',
@@ -226,6 +248,60 @@ describe('Codex quota provider (VS Code parity)', () => {
  });
 });

+describe('Claude quota provider (VS Code parity)', () => {
+  test('parses current limits, model-scoped limits, and extra usage', async () => {
+    stubFetchReturning(() => Promise.resolve(mockResponse({
+      limits: [
+        { kind: 'session', percent: 12, resets_at: '2026-08-20T12:00:00Z', scope: null },
+        { kind: 'weekly_all', percent: 34, resets_at: '2026-08-24T12:00:00Z', scope: null },
+        { kind: 'weekly_scoped', percent: 56, resets_at: '2026-08-24T12:00:00Z', scope: { model: { display_name: 'Sonnet' } } },
+      ],
+      spend: {
+        enabled: true,
+        percent: 25,
+        used: { amount_minor: 2500, exponent: 2, currency: 'USD' },
+        limit: { amount_minor: 10000, exponent: 2, currency: 'USD' },
+      },
+    })));
+
+    const result = await fetchQuotaForProvider('claude');
+
+    assert.equal(result.ok, true);
+    assert.equal(result.usage?.windows['5h']?.usedPercent, 12);
+    assert.equal(result.usage?.windows['7d']?.usedPercent, 34);
+    assert.equal(result.usage?.models?.Sonnet?.windows['7d']?.usedPercent, 56);
+    assert.equal(result.usage?.windows.extra_usage?.valueLabel, '$25.00 / $100.00');
+  });
+
+  test('keeps serving the last good values while Anthropic rate limits', async () => {
+    const responses = [
+      mockResponse({ five_hour: { utilization: 12, resets_at: '2026-08-20T12:00:00Z' } }),
+      {
+        ok: false,
+        status: 429,
+        headers: new Headers({ 'retry-after': '120' }),
+        json: async () => ({}),
+      } as Response,
+    ];
+    let requestCount = 0;
+    globalThis.fetch = (async () => {
+      const response = responses[requestCount];
+      requestCount += 1;
+      return response;
+    }) as typeof fetch;
+
+    const initial = await fetchQuotaForProvider('claude');
+    const rateLimited = await fetchQuotaForProvider('claude');
+    const duringCooldown = await fetchQuotaForProvider('claude');
+
+    assert.equal(initial.ok, true);
+    assert.equal(rateLimited.ok, true);
+    assert.equal(duringCooldown.ok, true);
+    assert.equal(duringCooldown.usage?.windows['5h']?.usedPercent, 12);
+    assert.equal(requestCount, 2);
+  });
+});
+
 describe('Z.ai quota provider (VS Code parity)', () => {
  test('surfaces 5-hour, weekly, and MCP quota windows', async () => {
    stubFetchReturning(() => Promise.resolve(mockResponse({
@@ -170,6 +170,7 @@ export type ProviderResult = {
  usage: ProviderUsage | null;
  fetchedAt: number;
  error?: string;
+  planLabel?: string | null;
 };

 const OPENCODE_CONFIG_DIR = path.join(os.homedir(), '.config', 'opencode');
@@ -1255,6 +1256,112 @@ const fetchGoogleQuota = async (): Promise<ProviderResult> => {
  });
 };

+const CLAUDE_DEFAULT_COOLDOWN_MS = 5 * 60 * 1000;
+const CLAUDE_MAX_COOLDOWN_MS = 60 * 60 * 1000;
+let claudeCredentialFingerprint: string | null = null;
+let claudeCachedUsage: ProviderUsage | null = null;
+let claudeCooldownUntil = 0;
+
+const claudeCooldownFromResponse = (response: Response): number => {
+  const raw = response.headers.get('retry-after');
+  const seconds = raw ? Number(raw) : Number.NaN;
+  if (Number.isFinite(seconds) && seconds > 0) {
+    return Math.min(seconds * 1000, CLAUDE_MAX_COOLDOWN_MS);
+  }
+  if (raw) {
+    const retryAt = Date.parse(raw);
+    if (Number.isFinite(retryAt) && retryAt > Date.now()) {
+      return Math.min(retryAt - Date.now(), CLAUDE_MAX_COOLDOWN_MS);
+    }
+  }
+  return CLAUDE_DEFAULT_COOLDOWN_MS;
+};
+
+const buildClaudeRateLimitResult = (): ProviderResult => (
+  claudeCachedUsage
+    ? buildResult({
+        providerId: 'claude',
+        providerName: 'Claude',
+        ok: true,
+        configured: true,
+        usage: claudeCachedUsage,
+      })
+    : buildResult({
+        providerId: 'claude',
+        providerName: 'Claude',
+        ok: false,
+        configured: true,
+        error: 'Rate limited by Anthropic. Retrying shortly.',
+      })
+);
+
+const buildClaudeUsage = (payload: Record<string, unknown>): ProviderUsage => {
+  const windows: Record<string, UsageWindow> = {};
+  const models: Record<string, ProviderUsage> = {};
+  const limits = Array.isArray(payload.limits) ? payload.limits : [];
+
+  for (const entry of limits) {
+    const limit = asObject(entry);
+    if (!limit) continue;
+    const usedPercent = toNumber(limit.percent);
+    const resetAt = toTimestamp(limit.resets_at);
+    if (limit.kind === 'session') {
+      windows['5h'] = toUsageWindow({ usedPercent, windowSeconds: 5 * 60 * 60, resetAt });
+    } else if (limit.kind === 'weekly_all') {
+      windows['7d'] = toUsageWindow({ usedPercent, windowSeconds: 7 * 24 * 60 * 60, resetAt });
+    } else if (limit.kind === 'weekly_scoped') {
+      const modelName = asNonEmptyString(asObject(asObject(limit.scope)?.model)?.display_name);
+      if (modelName) {
+        models[modelName] = {
+          windows: {
+            '7d': toUsageWindow({ usedPercent, windowSeconds: 7 * 24 * 60 * 60, resetAt }),
+          },
+        };
+      }
+    }
+  }
+
+  if (!limits.length) {
+    const fiveHour = asObject(payload.five_hour);
+    const sevenDay = asObject(payload.seven_day);
+    if (fiveHour) {
+      windows['5h'] = toUsageWindow({
+        usedPercent: toNumber(fiveHour.utilization),
+        windowSeconds: 5 * 60 * 60,
+        resetAt: toTimestamp(fiveHour.resets_at),
+      });
+    }
+    if (sevenDay) {
+      windows['7d'] = toUsageWindow({
+        usedPercent: toNumber(sevenDay.utilization),
+        windowSeconds: 7 * 24 * 60 * 60,
+        resetAt: toTimestamp(sevenDay.resets_at),
+      });
+    }
+  }
+
+  const spend = asObject(payload.spend);
+  if (spend?.enabled === true) {
+    const usedMoney = asObject(spend.used);
+    const limitMoney = asObject(spend.limit);
+    const usedMinor = toNumber(usedMoney?.amount_minor);
+    const limitMinor = toNumber(limitMoney?.amount_minor);
+    const exponent = toNumber(usedMoney?.exponent) ?? 2;
+    const currency = asNonEmptyString(usedMoney?.currency);
+    const prefix = currency === 'USD' || !currency ? '$' : `${currency} `;
+    const used = usedMinor === null ? null : usedMinor / 10 ** exponent;
+    const limit = limitMinor === null ? null : limitMinor / 10 ** (toNumber(limitMoney?.exponent) ?? 2);
+    windows.extra_usage = toUsageWindow({
+      usedPercent: toNumber(spend.percent),
+      windowSeconds: null,
+      resetAt: null,
+      valueLabel: used === null ? null : `${prefix}${formatMoney(used)}${limit === null ? '' : ` / ${prefix}${formatMoney(limit)}`}`,
+    });
+  }
+
+  return Object.keys(models).length ? { windows, models } : { windows };
+};
+
 const fetchClaudeQuota = async (): Promise<ProviderResult> => {
  const auth = readAuthFile();
  const entry = normalizeAuthEntry(getAuthEntry(auth, ['anthropic', 'claude'])) as Record<string, unknown> | null;
@@ -1270,6 +1377,15 @@ const fetchClaudeQuota = async (): Promise<ProviderResult> => {
    });
  }

+  const refreshToken = typeof entry?.refresh === 'string' ? entry.refresh : '';
+  const fingerprint = `${accessToken}\0${refreshToken}`;
+  if (claudeCredentialFingerprint !== fingerprint) {
+    claudeCredentialFingerprint = fingerprint;
+    claudeCachedUsage = null;
+    claudeCooldownUntil = 0;
+  }
+  if (Date.now() < claudeCooldownUntil) return buildClaudeRateLimitResult();
+
  try {
    const response = await fetch('https://api.anthropic.com/api/oauth/usage', {
      method: 'GET',
@@ -1279,6 +1395,21 @@ const fetchClaudeQuota = async (): Promise<ProviderResult> => {
      },
    });

+    if (response.status === 429) {
+      claudeCooldownUntil = Date.now() + claudeCooldownFromResponse(response);
+      return buildClaudeRateLimitResult();
+    }
+
+    if (response.status === 401 || response.status === 403) {
+      return buildResult({
+        providerId: 'claude',
+        providerName: 'Claude',
+        ok: false,
+        configured: true,
+        error: 'Claude session expired. Open Claude Code to sign in again.',
+      });
+    }
+
    if (!response.ok) {
      return buildResult({
        providerId: 'claude',
@@ -1290,47 +1421,14 @@ const fetchClaudeQuota = async (): Promise<ProviderResult> => {
    }

    const payload = await response.json() as Record<string, unknown>;
-    const windows: Record<string, UsageWindow> = {};
-    const fiveHour = (payload as Record<string, unknown>).five_hour as Record<string, unknown> | undefined;
-    const sevenDay = (payload as Record<string, unknown>).seven_day as Record<string, unknown> | undefined;
-    const sevenDaySonnet = (payload as Record<string, unknown>).seven_day_sonnet as Record<string, unknown> | undefined;
-    const sevenDayOpus = (payload as Record<string, unknown>).seven_day_opus as Record<string, unknown> | undefined;
-
-    if (fiveHour) {
-      windows['5h'] = toUsageWindow({
-        usedPercent: toNumber(fiveHour.utilization),
-        windowSeconds: null,
-        resetAt: toTimestamp(fiveHour.resets_at),
-      });
-    }
-    if (sevenDay) {
-      windows['7d'] = toUsageWindow({
-        usedPercent: toNumber(sevenDay.utilization),
-        windowSeconds: null,
-        resetAt: toTimestamp(sevenDay.resets_at),
-      });
-    }
-    if (sevenDaySonnet) {
-      windows['7d-sonnet'] = toUsageWindow({
-        usedPercent: toNumber(sevenDaySonnet.utilization),
-        windowSeconds: null,
-        resetAt: toTimestamp(sevenDaySonnet.resets_at),
-      });
-    }
-    if (sevenDayOpus) {
-      windows['7d-opus'] = toUsageWindow({
-        usedPercent: toNumber(sevenDayOpus.utilization),
-        windowSeconds: null,
-        resetAt: toTimestamp(sevenDayOpus.resets_at),
-      });
-    }
-
+    const usage = buildClaudeUsage(payload);
+    claudeCachedUsage = usage;
    return buildResult({
      providerId: 'claude',
      providerName: 'Claude',
      ok: true,
      configured: true,
-      usage: { windows },
+      usage,
    });
  } catch (error) {
    return buildResult({
@@ -2709,7 +2807,7 @@ const fetchXaiQuota = async (): Promise<ProviderResult> => {
  }
 };

-export const fetchQuotaForProvider = async (providerId: string): Promise<ProviderResult> => {
+const fetchQuotaForProviderUncoalesced = async (providerId: string): Promise<ProviderResult> => {
  switch (providerId) {
    case 'claude':
      return fetchClaudeQuota();
@@ -2782,3 +2880,16 @@ export const fetchQuotaForProvider = async (providerId: string): Promise<Provide
      });
  }
 };
+
+const pendingQuotaFetches = new Map<string, Promise<ProviderResult>>();
+
+export const fetchQuotaForProvider = (providerId: string): Promise<ProviderResult> => {
+  const existing = pendingQuotaFetches.get(providerId);
+  if (existing) return existing;
+
+  const pending = fetchQuotaForProviderUncoalesced(providerId).finally(() => {
+    if (pendingQuotaFetches.get(providerId) === pending) pendingQuotaFetches.delete(providerId);
+  });
+  pendingQuotaFetches.set(providerId, pending);
+  return pending;
+};