diff --git a/changelog.d/fixes/13088-kiro-quota-freetrial.md b/changelog.d/fixes/13088-kiro-quota-freetrial.md new file mode 100644 index 0000000000..2eb634cc08 --- /dev/null +++ b/changelog.d/fixes/13088-kiro-quota-freetrial.md @@ -0,0 +1 @@ +- fix(quota): keep Kiro active while any _freetrial pool has quota (#13088) diff --git a/config/quality/eslint-suppressions.json b/config/quality/eslint-suppressions.json index 88aa34231e..dd4665884d 100644 --- a/config/quality/eslint-suppressions.json +++ b/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 diff --git a/open-sse/services/genericQuotaFetcher.ts b/open-sse/services/genericQuotaFetcher.ts index 43ed09d68d..f5ed9216a8 100644 --- a/open-sse/services/genericQuotaFetcher.ts +++ b/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[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 +): { percentUsed: number; resetAt: string | null } { + const effectiveByBase = new Map(); + 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, diff --git a/open-sse/services/quotaPreflight.ts b/open-sse/services/quotaPreflight.ts index eac585e3d2..0bc6886944 100644 --- a/open-sse/services/quotaPreflight.ts +++ b/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 +): Map> { + const groups = new Map>(); + 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, 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( diff --git a/tests/unit/generic-quota-fetcher.test.ts b/tests/unit/generic-quota-fetcher.test.ts index 54759557b2..3b53d50a35 100644 --- a/tests/unit/generic-quota-fetcher.test.ts +++ b/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); +}); diff --git a/tests/unit/quota-preflight.test.ts b/tests/unit/quota-preflight.test.ts index 11d1e65386..d33001dd95 100644 --- a/tests/unit/quota-preflight.test.ts +++ b/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"); +});