@yansigit/opencodex 2.33.0 → 2.35.0

package/src/providers/quota.ts

@@ -3,7 +3,9 @@ import {
   effectiveCodexAuthAccountId,
   fetchMainAccountInfoSnapshot,
   listCodexAuthAccountsSnapshot,
+  withSparkVisibility,
 } from "../codex/auth-api";
+import type { StoredAccountQuota } from "../codex/quota";
 import { isMainAccountIdentityGenerationLive } from "../codex/main-account-cache";
 import { MAIN_CODEX_ACCOUNT_ID } from "../codex/main-account";
 import { codexPlanKey } from "../codex/plan";
@@ -173,6 +175,27 @@ function quotaSignatureValue(quota: CodexCapacityQuota | null): unknown {
   };
 }
 
+function providerQuotaFromCodexQuota(
+  quota: StoredAccountQuota | Omit<StoredAccountQuota, "updatedAt"> | null | undefined,
+): CodexCapacityQuota | null {
+  if (!quota) return null;
+  // Every Codex-sourced provider report funnels through here — the pooled path via
+  // listCodexAuthAccountsSnapshot and the `direct` path via fetchMainAccountInfoSnapshot, which
+  // never touches the Codex Auth DTO. Applying the Spark preference at this one point is what
+  // stops the row surviving on /api/provider-quotas after the operator switched it off.
+  quota = withSparkVisibility(quota ?? null) ?? quota;
+  return {
+    ...(quota.shortPercent !== undefined ? { fiveHourPercent: quota.shortPercent } : {}),
+    ...(quota.shortResetAt !== undefined ? { fiveHourResetAt: quota.shortResetAt } : {}),
+    ...(quota.weeklyPercent !== undefined ? { weeklyPercent: quota.weeklyPercent } : {}),
+    ...(quota.weeklyResetAt !== undefined ? { weeklyResetAt: quota.weeklyResetAt } : {}),
+    ...(quota.monthlyPercent !== undefined ? { monthlyPercent: quota.monthlyPercent } : {}),
+    ...(quota.monthlyResetAt !== undefined ? { monthlyResetAt: quota.monthlyResetAt } : {}),
+    ...(quota.customWindows !== undefined ? { customWindows: quota.customWindows } : {}),
+    updatedAt: "updatedAt" in quota ? quota.updatedAt : Date.now(),
+  };
+}
+
 /** Hash only presentation-relevant state; account ids and email addresses never enter the key. */
 function cacheKeyWithAggregationState(
   config: OcxConfig,
@@ -190,7 +213,7 @@ function cacheKeyWithAggregationState(
         plan: codexPlanKey(account.plan) ?? null,
         paused: account.paused,
         needsReauth: account.needsReauth === true,
-        quota: quotaSignatureValue(account.quota as CodexCapacityQuota | null),
+        quota: quotaSignatureValue(providerQuotaFromCodexQuota(account.quota)),
       }));
       const canonicalRows = rows.map(row => JSON.stringify(row)).sort();
       const digest = createHash("sha256").update(JSON.stringify(canonicalRows)).digest("hex").slice(0, 24);
@@ -1110,9 +1133,8 @@ async function fetchChatGptForwardQuota(
 ): Promise<ProviderQuotaReport | null> {
   if (providerCodexAccountMode(provider, providerConfig) === "direct") {
     const snapshot = await fetchMainAccountInfoSnapshot(forceRefresh);
-    const quota = snapshot.info.quota
-      ? { ...snapshot.info.quota, updatedAt: Date.now() } as ProviderQuota
-      : null;
+    const quota = providerQuotaFromCodexQuota(snapshot.info.quota);
+    if (quota) quota.updatedAt = Date.now();
     return quota
       ? tagNativeMainReport(report(provider, "chatgpt:wham", quota), snapshot.mainIdentityGeneration)
       : null;
@@ -1120,10 +1142,14 @@ async function fetchChatGptForwardQuota(
   const snapshot = await (prefetchedSnapshot ?? listCodexAuthAccountsSnapshot(config, forceRefresh));
   const accounts = snapshot.accounts;
   const activeId = effectiveCodexAuthAccountId(config);
-  const capacityAccounts = accounts.map(account => ({ ...account, active: account.id === activeId }));
+  const capacityAccounts = accounts.map(account => ({
+    ...account,
+    active: account.id === activeId,
+    quota: providerQuotaFromCodexQuota(account.quota),
+  }));
   const active = capacityAccounts.find(account => account.active)
-    ?? accounts.find(account => account.id === MAIN_CODEX_ACCOUNT_ID)
-    ?? accounts[0];
+    ?? capacityAccounts.find(account => account.id === MAIN_CODEX_ACCOUNT_ID)
+    ?? capacityAccounts[0];
   const now = Date.now();
   const capacity = aggregateCodexPoolCapacity(capacityAccounts, now);
   if (capacity.aggregation && capacity.quota) {
@@ -1139,7 +1165,12 @@ async function fetchChatGptForwardQuota(
   }
   const activeUsable = !!active && !active.paused && active.needsReauth !== true;
   const quota = activeUsable && active?.quota
-    ? { ...active.quota, updatedAt: active.quota.updatedAt ?? Date.now() } as CodexCapacityQuota
+    ? {
+        ...active.quota,
+        fiveHourPercent: active.quota.fiveHourPercent ?? active.quota.shortPercent,
+        fiveHourResetAt: active.quota.fiveHourResetAt ?? active.quota.shortResetAt,
+        updatedAt: active.quota.updatedAt ?? Date.now(),
+      } as CodexCapacityQuota
     : null;
   const quotaFresh = !!quota
     && Number.isFinite(quota.updatedAt)
@@ -1285,6 +1316,24 @@ function parseClaudeBucket(value: unknown): { percent?: number; resetAt?: number
   return { percent, resetAt };
 }
 
+function parseClaudeLimit(value: unknown): { label: string; percent: number; resetAt?: number } | null {
+  const rec = asRecord(value);
+  if (!rec) return null;
+  const percent = normalizePercent(rec.percent);
+  if (percent === undefined) return null;
+  const scope = asRecord(rec.scope);
+  const model = asRecord(scope?.model);
+  const rawLabel = String(model?.display_name ?? "").trim();
+  if (!rawLabel) return null;
+  const lowerLabel = rawLabel.toLowerCase();
+  const label = lowerLabel.includes("fable") ? "Fable"
+    : lowerLabel.includes("opus") ? "Opus"
+      : lowerLabel.includes("sonnet") ? "Sonnet"
+        : rawLabel;
+  const resetAt = normalizeResetAt(rec.resets_at);
+  return { label, percent, ...(resetAt !== undefined ? { resetAt } : {}) };
+}
+
 /** Claude's OAuth usage endpoint, probed with ONE account's own bearer token. */
 const anthropicUsageInflight = new Map<string, Promise<ProviderQuota | null>>();
 
@@ -1308,11 +1357,25 @@ async function fetchAnthropicUsageQuota(accessToken: string): Promise<ProviderQu
     if (!body) return null;
     const fiveHour = parseClaudeBucket(body.five_hour);
     const sevenDay = parseClaudeBucket(body.seven_day);
+    const fable = parseClaudeBucket(body.seven_day_fable);
     const opus = parseClaudeBucket(body.seven_day_opus);
     const sonnet = parseClaudeBucket(body.seven_day_sonnet);
     const customWindows: ProviderQuotaWindow[] = [];
+    if (fable?.percent !== undefined) customWindows.push({ label: "Fable", percent: fable.percent, ...(fable.resetAt !== undefined ? { resetAt: fable.resetAt } : {}) });
     if (opus?.percent !== undefined) customWindows.push({ label: "Opus", percent: opus.percent, ...(opus.resetAt !== undefined ? { resetAt: opus.resetAt } : {}) });
     if (sonnet?.percent !== undefined) customWindows.push({ label: "Sonnet", percent: sonnet.percent, ...(sonnet.resetAt !== undefined ? { resetAt: sonnet.resetAt } : {}) });
+    const knownLabels = new Set(customWindows.map(window => window.label.toLowerCase()));
+    const limits = Array.isArray(body.limits) ? body.limits : [];
+    for (const rawLimit of limits) {
+      const limitRecord = asRecord(rawLimit);
+      // `session` and `weekly_all` mirror the canonical five-hour and weekly
+      // buckets above; only model-scoped weekly limits add a third window.
+      if (String(limitRecord?.kind ?? "").trim().toLowerCase() !== "weekly_scoped") continue;
+      const limit = parseClaudeLimit(rawLimit);
+      if (!limit || knownLabels.has(limit.label.toLowerCase())) continue;
+      knownLabels.add(limit.label.toLowerCase());
+      customWindows.push(limit);
+    }
     const quota: ProviderQuota = {
       // Claude's 5-hour window is a first-class rate limit, same as the Codex login 5h/weekly
       // rows: report it in the canonical fields so the dashboard renders it with the standard
@@ -1401,7 +1464,12 @@ export interface ProviderAccountQuota {
 
 /** Providers whose per-account quota can be probed. Extend as other OAuth APIs are covered. */
 export function supportsPerAccountQuota(provider: string): boolean {
-  return provider === "anthropic" || provider === "google-antigravity";
+  return (
+    provider === "anthropic" ||
+    provider === "google-antigravity" ||
+    provider === "command-code" ||
+    provider === "cursor"
+  );
 }
 
 function accountCacheKey(provider: string, accountId: string): string {
@@ -1526,6 +1594,14 @@ async function fetchAccountQuota(
         quota = await fetchAnthropicUsageQuota(token);
       } else if (provider === "google-antigravity") {
         quota = await fetchAntigravityAccountQuota(accountId);
+      } else if (provider === "command-code") {
+        const token = await getTokenForAccountQuotaProbe(provider, accountId);
+        const res = await fetchCommandCodeUsageQuota(token);
+        quota = res === TERMINAL_QUOTA_FAILURE ? null : res;
+      } else if (provider === "cursor") {
+        const token = await getTokenForAccountQuotaProbe(provider, accountId);
+        const res = await fetchCursorUsageQuota(token);
+        quota = res?.quota ?? null;
       }
       if (!quota) {
         // Preserve last-good bars and mark unavailable; advance TTL so failures
@@ -1741,8 +1817,9 @@ function parseCommandCodeWindow(value: unknown): { percent: number; resetAt?: nu
   if (!row) return null;
   const cap = toFiniteNumber(row.cap);
   const used = toFiniteNumber(row.used);
-  if (cap === undefined || used === undefined || cap <= 0 || used < 0) return null;
-  const percent = normalizePercent((used / cap) * 100);
+  const percent = cap !== undefined && used !== undefined && cap > 0 && used >= 0
+    ? normalizePercent((used / cap) * 100)
+    : normalizePercent(row.percent);
   if (percent === undefined) return null;
   const resetAt = quotaResetAt(row);
   return { percent, ...(resetAt !== undefined ? { resetAt } : {}) };
@@ -1829,11 +1906,7 @@ async function resolveCommandCodeQuotaBearer(config: OcxProviderConfig): Promise
  * usage view uses (windowLimits.fiveHour / windowLimits.weekly), plus soft
  * whoami (team orgId scoping) and subscription-scoped spend for creditsUsd.
  */
-async function fetchCommandCodeQuota(provider: string, config: OcxProviderConfig): Promise<ProviderQuotaProbeResult> {
-  // Never release credentials to a user-edited or lookalike provider host.
-  if (!isCanonicalCommandCodeBaseUrl(config.baseUrl)) return null;
-  const bearer = await resolveCommandCodeQuotaBearer(config);
-  if (!bearer) return null;
+async function fetchCommandCodeUsageQuota(bearer: string): Promise<ProviderQuota | null | typeof TERMINAL_QUOTA_FAILURE> {
   const whoamiBody = await fetchCommandCodeJson(COMMAND_CODE_WHOAMI_URL, bearer);
   const whoami = asRecord(whoamiBody?.data) ?? whoamiBody;
   const org = asRecord(whoami?.org);
@@ -1857,7 +1930,7 @@ async function fetchCommandCodeQuota(provider: string, config: OcxProviderConfig
   const fiveHour = parseCommandCodeWindow(limits?.fiveHour);
   const weekly = parseCommandCodeWindow(limits?.weekly);
   const creditsUsd = await fetchCommandCodeSpend(bearer, credits, orgQuery);
-  return report(provider, "command-code:credits", {
+  return {
     ...(fiveHour ? {
       fiveHourPercent: fiveHour.percent,
       ...(fiveHour.resetAt !== undefined ? { fiveHourResetAt: fiveHour.resetAt } : {}),
@@ -1868,18 +1941,21 @@ async function fetchCommandCodeQuota(provider: string, config: OcxProviderConfig
     } : {}),
     ...(creditsUsd ? { creditsUsd } : {}),
     updatedAt: Date.now(),
-  });
+  };
 }
 
-/** Cursor included usage via api2.cursor.sh (Bearer from OAuth) — unofficial, may change. */
-async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport | null> {
-  let accessToken: string;
-  try {
-    accessToken = await getValidAccessToken("cursor");
-  } catch {
-    return null;
-  }
+async function fetchCommandCodeQuota(provider: string, config: OcxProviderConfig): Promise<ProviderQuotaProbeResult> {
+  // Never release credentials to a user-edited or lookalike provider host.
+  if (!isCanonicalCommandCodeBaseUrl(config.baseUrl)) return null;
+  const bearer = await resolveCommandCodeQuotaBearer(config);
+  if (!bearer) return null;
+  const result = await fetchCommandCodeUsageQuota(bearer);
+  if (!result || result === TERMINAL_QUOTA_FAILURE) return result;
+  return report(provider, "command-code:credits", result);
+}
 
+/** Cursor included usage via api2.cursor.sh (Bearer from OAuth) — unofficial, may change. */
+async function fetchCursorUsageQuota(accessToken: string): Promise<{ quota: ProviderQuota; source: string } | null> {
   const authHeaders = {
     Accept: "application/json",
     Authorization: `Bearer ${accessToken}`,
@@ -1939,15 +2015,17 @@ async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport |
         }
 
         if (totalPercent !== undefined || customWindows.length > 0) {
-          const built = report(provider, "cursor:period-usage", {
-            ...(totalPercent !== undefined ? {
-              monthlyPercent: totalPercent,
-              ...(resetAt !== undefined ? { monthlyResetAt: resetAt } : {}),
-            } : {}),
-            ...(customWindows.length > 0 ? { customWindows } : {}),
-            updatedAt: Date.now(),
-          });
-          if (built) return { ...built, reverseEngineered: true };
+          return {
+            source: "cursor:period-usage",
+            quota: {
+              ...(totalPercent !== undefined ? {
+                monthlyPercent: totalPercent,
+                ...(resetAt !== undefined ? { monthlyResetAt: resetAt } : {}),
+              } : {}),
+              ...(customWindows.length > 0 ? { customWindows } : {}),
+              updatedAt: Date.now(),
+            },
+          };
         }
       }
     }
@@ -1973,12 +2051,14 @@ async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport |
             ? normalizePercent((used / limit) * 100)
             : undefined);
         if (percent !== undefined) {
-          const built = report(provider, "cursor:usage-summary", {
-            monthlyPercent: percent,
-            monthlyResetAt: normalizeResetAt(body?.billingCycleEnd),
-            updatedAt: Date.now(),
-          });
-          if (built) return { ...built, reverseEngineered: true };
+          return {
+            source: "cursor:usage-summary",
+            quota: {
+              monthlyPercent: percent,
+              monthlyResetAt: normalizeResetAt(body?.billingCycleEnd),
+              updatedAt: Date.now(),
+            },
+          };
         }
       }
     }
@@ -2027,11 +2107,26 @@ async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport |
         return Date.UTC(start.getUTCFullYear(), start.getUTCMonth() + 1, start.getUTCDate());
       })()
     : undefined;
-  const built = report(provider, "cursor:auth-usage", {
-    monthlyPercent: percent,
-    ...(monthlyResetAt !== undefined ? { monthlyResetAt } : {}),
-    updatedAt: Date.now(),
-  });
+  return {
+    source: "cursor:auth-usage",
+    quota: {
+      monthlyPercent: percent,
+      ...(monthlyResetAt !== undefined ? { monthlyResetAt } : {}),
+      updatedAt: Date.now(),
+    },
+  };
+}
+
+async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport | null> {
+  let accessToken: string;
+  try {
+    accessToken = await getValidAccessToken("cursor");
+  } catch {
+    return null;
+  }
+  const result = await fetchCursorUsageQuota(accessToken);
+  if (!result) return null;
+  const built = report(provider, result.source, result.quota);
   return built ? { ...built, reverseEngineered: true } : null;
 }
 
@@ -2219,6 +2314,19 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig
   });
 }
 
+async function fetchAiStudioQuota(name: string, _provider: OcxProviderConfig): Promise<ProviderQuotaReport | null> {
+  const now = Date.now();
+  return {
+    provider: name,
+    label: "Google AI Studio (Web)",
+    source: "Direct Session",
+    updatedAt: now,
+    quota: {
+      updatedAt: now,
+    },
+  };
+}
+
 async function maybeFetchProviderQuota(
   name: string,
   provider: OcxProviderConfig,
@@ -2294,6 +2402,9 @@ async function maybeFetchProviderQuota(
     if ((provider.authMode ?? "key") === "key" && name === "neuralwatt") {
       return fetchNeuralwattQuota(name, provider);
     }
+    if (provider.googleMode === "ai-studio-web" || name === "google-aistudio") {
+      return fetchAiStudioQuota(name, provider);
+    }
     return null;
   } catch {
     return null;