fix(quota): keep Kiro active while any _freetrial pool has quota (#13088) (#13324)

* fix(quota): keep Kiro active while any _freetrial pool has quota (#13088)

* fix(quota): rename changelog and remove blank line

---------

Co-authored-by: giauphan <giauphan@users.noreply.github.com>
Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

changelog.d/fixes/13088-kiro-quota-freetrial.md

@@ -0,0 +1 @@
+- fix(quota): keep Kiro active while any _freetrial pool has quota (#13088)

config/quality/eslint-suppressions.json

@@ -535,11 +535,6 @@
       "count": 1
     }
   },
-  "open-sse/services/genericQuotaFetcher.ts": {
-    "no-restricted-syntax": {
-      "count": 1
-    }
-  },
   "open-sse/services/grokQuotaFetcher.ts": {
     "@typescript-eslint/no-unused-vars": {
       "count": 1

open-sse/services/genericQuotaFetcher.ts

@@ -26,6 +26,7 @@ import {
 } from "./quotaPreflight.ts";
 import { getAntigravityQuotaFamily, getQuotaFetchScope } from "./antigravityQuotaFamily.ts";
 import { boundedMap } from "../../src/lib/quota/boundedMap.ts";
+import { toNumberOrNull } from "@/shared/utils/numeric";
 
 type UsageFetcher = (
   connection: Parameters<typeof getUsageForProvider>[0],
@@ -150,15 +151,6 @@ if (typeof _cacheCleanup === "object" && "unref" in _cacheCleanup) {
   (_cacheCleanup as { unref?: () => void }).unref?.();
 }
 
-function toNumber(value: unknown): number | null {
-  if (typeof value === "number" && Number.isFinite(value)) return value;
-  if (typeof value === "string") {
-    const parsed = parseFloat(value);
-    if (Number.isFinite(parsed)) return parsed;
-  }
-  return null;
-}
-
 /**
  * Compute percentUsed (0-1) for a single quota entry. Prefers the explicit
  * remainingPercentage / used / total fields surfaced by per-provider
@@ -176,15 +168,15 @@ function percentUsedForQuota(entry: unknown): number | null {
   // otherwise one unreported model falsely exhausts the whole connection.
   if (q.fractionReported === false) return null;
 
-  const remainingPercentage = toNumber(q.remainingPercentage);
+  const remainingPercentage = toNumberOrNull(q.remainingPercentage);
   if (remainingPercentage !== null) {
     // remainingPercentage is 0-100 in the usage.ts contract.
     const used = (100 - Math.max(0, Math.min(100, remainingPercentage))) / 100;
     return used;
   }
 
-  const used = toNumber(q.used);
-  const total = toNumber(q.total);
+  const used = toNumberOrNull(q.used);
+  const total = toNumberOrNull(q.total);
   if (used !== null && total !== null && total > 0) {
     return Math.max(0, Math.min(1, used / total));
   }
@@ -219,6 +211,28 @@ type UsageToQuotaContext = {
   provider?: string | null;
 };
 
+function aggregateGroupedQuotaValues(
+  windows: Record<string, { percentUsed: number; resetAt: string | null }>
+): { percentUsed: number; resetAt: string | null } {
+  const effectiveByBase = new Map<string, { percentUsed: number; resetAt: string | null }>();
+  for (const [key, entry] of Object.entries(windows)) {
+    const base = key.endsWith("_freetrial") ? key.slice(0, -10) : key;
+    const cur = effectiveByBase.get(base);
+    if (!cur || entry.percentUsed < cur.percentUsed) {
+      effectiveByBase.set(base, { percentUsed: entry.percentUsed, resetAt: entry.resetAt ?? null });
+    }
+  }
+  let percentUsed = 0;
+  let resetAt: string | null = null;
+  for (const eff of effectiveByBase.values()) {
+    if (eff.percentUsed > percentUsed) {
+      percentUsed = eff.percentUsed;
+      resetAt = eff.resetAt;
+    }
+  }
+  return { percentUsed, resetAt };
+}
+
 export function convertUsageToQuotaInfo(
   usage: unknown,
   context: UsageToQuotaContext = {}
@@ -267,13 +281,7 @@ export function convertUsageToQuotaInfo(
   if (Object.keys(providerScopedWindows).length === 0) return null;
 
   const normalized = normalizeQuotaWindows(providerScopedWindows, context);
-  const scopedEntries = Object.values(providerScopedWindows);
-  const percentUsed = scopedEntries.reduce((worst, entry) => Math.max(worst, entry.percentUsed), 0);
-  const resetAt =
-    scopedEntries.reduce<{ percentUsed: number; resetAt: string | null } | null>(
-      (worst, entry) => (!worst || entry.percentUsed > worst.percentUsed ? entry : worst),
-      null
-    )?.resetAt ?? null;
+  const { percentUsed, resetAt } = aggregateGroupedQuotaValues(providerScopedWindows);
 
   return {
     used: 0,

open-sse/services/quotaPreflight.ts

@@ -205,38 +205,89 @@ function limitReachedResult(quota: QuotaInfo): PreflightQuotaResult {
   );
 }
 
+function isEntryExhausted(
+  windowName: string,
+  percentUsed: number,
+  thresholds?: PreflightQuotaThresholds
+): boolean {
+  const minRemainingPercent = resolveOrDefault(
+    thresholds?.resolveMinRemainingPercent,
+    windowName,
+    DEFAULT_MIN_REMAINING_PERCENT
+  );
+  return isRemainingAtOrBelowThreshold(remainingPercentFrom(percentUsed), minRemainingPercent);
+}
+
+function evaluateQuotaGroup(
+  entries: Array<[string, QuotaWindowInfo]>,
+  thresholds?: PreflightQuotaThresholds
+): {
+  exhausted: boolean;
+  worstPercent: number;
+  worstWindow: string | null;
+  worstResetAt: string | null;
+} {
+  let exhausted = true;
+  let worstPercent = -1;
+  let worstWindow: string | null = null;
+  let worstResetAt: string | null = null;
+  for (const [windowName, windowInfo] of entries) {
+    if (!isEntryExhausted(windowName, windowInfo.percentUsed, thresholds)) {
+      exhausted = false;
+    }
+    if (windowInfo.percentUsed > worstPercent) {
+      worstPercent = windowInfo.percentUsed;
+      worstWindow = windowName;
+      worstResetAt = windowInfo.resetAt ?? null;
+    }
+  }
+  return { exhausted, worstPercent: Math.max(0, worstPercent), worstWindow, worstResetAt };
+}
+
+function groupQuotaWindowsByBase(
+  windows: NonNullable<QuotaInfo["windows"]>
+): Map<string, Array<[string, QuotaWindowInfo]>> {
+  const groups = new Map<string, Array<[string, QuotaWindowInfo]>>();
+  for (const [windowName, windowInfo] of Object.entries(windows)) {
+    if (!Number.isFinite(windowInfo.percentUsed)) continue;
+    const base = windowName.endsWith("_freetrial") ? windowName.slice(0, -10) : windowName;
+    const list = groups.get(base);
+    if (list) list.push([windowName, windowInfo]);
+    else groups.set(base, [[windowName, windowInfo]]);
+  }
+  return groups;
+}
+
 function quotaWindowCutoffResult(
   windows: NonNullable<QuotaInfo["windows"]>,
   thresholds?: PreflightQuotaThresholds
 ): PreflightQuotaResult | null {
-  let worstUsedPercent = 0;
-  let worstWindow: string | null = null;
-  let worstResetAt: string | null = null;
+  const groups = groupQuotaWindowsByBase(windows);
+  if (groups.size === 0) return null;
 
-  for (const [windowName, windowInfo] of Object.entries(windows)) {
-    if (!Number.isFinite(windowInfo.percentUsed)) continue;
-    const minRemainingPercent = resolveOrDefault(
-      thresholds?.resolveMinRemainingPercent,
-      windowName,
-      DEFAULT_MIN_REMAINING_PERCENT
+  let worstExhaustedPercent = 0;
+  let worstExhaustedWindow: string | null = null;
+  let worstExhaustedResetAt: string | null = null;
+  let hasExhaustedGroup = false;
+
+  for (const entries of groups.values()) {
+    const { exhausted, worstPercent, worstWindow, worstResetAt } = evaluateQuotaGroup(
+      entries,
+      thresholds
     );
-    if (
-      !isRemainingAtOrBelowThreshold(
-        remainingPercentFrom(windowInfo.percentUsed),
-        minRemainingPercent
-      )
-    ) {
-      continue;
+    if (exhausted) {
+      hasExhaustedGroup = true;
+      if (worstPercent > worstExhaustedPercent || worstExhaustedWindow === null) {
+        worstExhaustedPercent = worstPercent;
+        worstExhaustedWindow = worstWindow;
+        worstExhaustedResetAt = worstResetAt;
+      }
     }
-    if (windowInfo.percentUsed <= worstUsedPercent && worstWindow !== null) continue;
-    worstUsedPercent = windowInfo.percentUsed;
-    worstWindow = windowName;
-    worstResetAt = windowInfo.resetAt ?? null;
   }
 
-  return worstWindow === null
-    ? null
-    : exhaustedResult(worstUsedPercent, worstResetAt, worstWindow);
+  return hasExhaustedGroup
+    ? exhaustedResult(worstExhaustedPercent, worstExhaustedResetAt, worstExhaustedWindow)
+    : null;
 }
 
 function quotaPercentCutoffResult(

tests/unit/generic-quota-fetcher.test.ts

@@ -401,3 +401,28 @@ test("stamp expiry during in-flight fetch still writes the wrapper cache", async
   assert.equal(calls.length, 2, "expired stamp during await is not a 429; cache the result");
   invalidateGenericQuotaCache("agy", connectionId);
 });
+
+test("convertUsageToQuotaInfo aggregates _freetrial and non-freetrial windows using best remaining", async () => {
+  const usage = {
+    quotas: {
+      credit: { used: 50, total: 50 },
+      credit_freetrial: { used: 0, total: 500 },
+    },
+  };
+  const result = convertUsageToQuotaInfo(usage, { provider: "kiro" });
+  // Used: min of 100% and 0% = 0%
+  assert.equal(result?.percentUsed, 0);
+  assert.equal(result?.limitReached, false);
+});
+
+test("convertUsageToQuotaInfo blocks when both base and freetrial are exhausted", async () => {
+  const usage = {
+    quotas: {
+      credit: { used: 50, total: 50 },
+      credit_freetrial: { used: 500, total: 500 },
+    },
+  };
+  const result = convertUsageToQuotaInfo(usage, { provider: "kiro" });
+  assert.equal(result?.percentUsed, 1);
+  assert.equal(result?.limitReached, true);
+});

tests/unit/quota-preflight.test.ts

@@ -346,3 +346,32 @@ test("registerQuotaWindows / getQuotaWindows round-trips", () => {
   // Unknown provider returns an empty list rather than undefined.
   assert.deepEqual([...getQuotaWindows("provider-with-no-registration-anywhere")], []);
 });
+
+test("evaluateQuotaCutoff permits connection when base is exhausted but _freetrial is available", () => {
+  const quota = {
+    used: 50,
+    total: 50,
+    percentUsed: 0,
+    windows: {
+      credit: { percentUsed: 1.0, resetAt: null },
+      credit_freetrial: { percentUsed: 0.0, resetAt: null },
+    },
+  };
+  const result = evaluateQuotaCutoff(quota);
+  assert.equal(result.proceed, true);
+});
+
+test("evaluateQuotaCutoff blocks connection when both base and _freetrial are exhausted", () => {
+  const quota = {
+    used: 50,
+    total: 50,
+    percentUsed: 1.0,
+    windows: {
+      credit: { percentUsed: 1.0, resetAt: null },
+      credit_freetrial: { percentUsed: 1.0, resetAt: null },
+    },
+  };
+  const result = evaluateQuotaCutoff(quota);
+  assert.equal(result.proceed, false);
+  assert.equal(result.reason, "quota_exhausted");
+});