Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 65 additions & 20 deletions src/providers/quota.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import {
sweepExpiredOnWrite,
type GenerationContext,
} from "../lib/state-store-sweeper";
import { readBoundedResponseBody } from "../lib/bounded-body";
import {
aggregateCodexPoolCapacity,
CODEX_CAPACITY_MAX_QUOTA_AGE_MS,
Expand All @@ -33,6 +34,8 @@ const ACCOUNT_TOKEN_SKEW_MS = 60_000;

const CACHE_TTL_MS = 5 * 60_000;
const REQUEST_TIMEOUT_MS = 8_000;
/** Successful provider quota payloads are small; reject oversized or stalled JSON before parsing. */
export const QUOTA_RESPONSE_MAX_BYTES = 512 * 1024;
const KIMI_CODE_BASE_URL = "https://api.kimi.com/coding/v1";
const KIMI_CODE_USAGE_URL = `${KIMI_CODE_BASE_URL}/usages`;
const A6API_BASE_URL = "https://api.a6api.com";
Expand Down Expand Up @@ -260,6 +263,43 @@ function asRecord(value: unknown): Record<string, unknown> | null {
return value && typeof value === "object" && !Array.isArray(value) ? value as Record<string, unknown> : null;
}

const QUOTA_JSON_READ_FAILURE = Symbol("quota-json-read-failure");

async function readQuotaJson(
response: Response,
timeoutMs = REQUEST_TIMEOUT_MS,
): Promise<unknown | typeof QUOTA_JSON_READ_FAILURE> {
const declaredLength = Number(response.headers.get("content-length"));
if (Number.isFinite(declaredLength) && declaredLength > QUOTA_RESPONSE_MAX_BYTES) {
try {
void response.body?.cancel(
new DOMException("Provider quota response is too large", "QuotaExceededError"),
).catch(() => undefined);
} catch {
// Best-effort cancellation only.
}
return QUOTA_JSON_READ_FAILURE;
}

try {
const bounded = await readBoundedResponseBody(response, {
maxBytes: QUOTA_RESPONSE_MAX_BYTES,
totalTimeoutMs: timeoutMs,
inactivityTimeoutMs: timeoutMs,
});
if (bounded.oversized || bounded.truncated || !bounded.displaySafe) return QUOTA_JSON_READ_FAILURE;
return JSON.parse(bounded.text) as unknown;
} catch {
return QUOTA_JSON_READ_FAILURE;
}
}

/** Test-only access to the quota reader's deadline and cancellation contract. */
export async function readProviderQuotaJsonForTests(response: Response, timeoutMs: number): Promise<unknown> {
const result = await readQuotaJson(response, timeoutMs);
return result === QUOTA_JSON_READ_FAILURE ? null : result;
}

function isBuiltInChatGptForwardProvider(name: string, provider: OcxProviderConfig): boolean {
return name === OPENAI_CODEX_PROVIDER_ID && isCanonicalOpenAiForwardProvider(provider);
}
Expand Down Expand Up @@ -354,8 +394,13 @@ async function fetchA6apiQuota(provider: string, config: OcxProviderConfig): Pro
? TERMINAL_QUOTA_FAILURE
: null;
}
const subscription = a6apiPayload(await subscriptionResponse.json().catch(() => null));
const token = a6apiPayload(await tokenResponse.json().catch(() => null));
const [subscriptionBody, tokenBody] = await Promise.all([
readQuotaJson(subscriptionResponse),
readQuotaJson(tokenResponse),
]);
if (subscriptionBody === QUOTA_JSON_READ_FAILURE || tokenBody === QUOTA_JSON_READ_FAILURE) return null;
const subscription = a6apiPayload(subscriptionBody);
const token = a6apiPayload(tokenBody);
const unlimited = token?.unlimited_quota === true
|| token?.unlimited_quota === 1
|| token?.unlimited_quota === "true";
Expand Down Expand Up @@ -433,7 +478,7 @@ async function fetchOpenRouterQuota(provider: string, config: OcxProviderConfig)
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
if (!data) return null;
const limit = toFiniteNumber(data.limit);
Expand Down Expand Up @@ -480,7 +525,7 @@ async function fetchDeepSeekQuota(provider: string, config: OcxProviderConfig):
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
// The payload nests balances under `balance_infos` rows keyed by currency;
// prefer a USD row, then CNY, then the first row that parses.
const infos = Array.isArray(body?.balance_infos) ? body.balance_infos as unknown[] : null;
Expand Down Expand Up @@ -528,7 +573,7 @@ async function fetchClineQuota(provider: string, config: OcxProviderConfig): Pro
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
const limits = Array.isArray(data?.limits) ? data.limits : null;
if (!limits) return null;
Expand Down Expand Up @@ -576,7 +621,7 @@ async function fetchZaiQuota(provider: string, config: OcxProviderConfig): Promi
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
if (!body || body.success === false) return null;
const data = asRecord(body.data) ?? body;
// The plugin renders a 5h token window, a weekly window, and a monthly MCP
Expand Down Expand Up @@ -632,7 +677,7 @@ async function fetchMinimaxQuota(provider: string, config: OcxProviderConfig): P
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
if (!body || body.success === false) return null;
const data = asRecord(body.data) ?? body;
const remainsMs = toFiniteNumber(data.remains_time ?? data.remainsTime);
Expand Down Expand Up @@ -675,7 +720,7 @@ async function fetchMoonshotQuota(provider: string, config: OcxProviderConfig):
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
if (!data) return null;
const available = toFiniteNumber(data.available_balance);
Expand Down Expand Up @@ -711,7 +756,7 @@ async function fetchVeniceQuota(provider: string, config: OcxProviderConfig): Pr
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
if (!data) return null;
const diemBalance = toFiniteNumber(data.balance);
Expand Down Expand Up @@ -754,7 +799,7 @@ async function fetchSyntheticQuota(provider: string, config: OcxProviderConfig):
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
const quota: ProviderQuota = { updatedAt: Date.now() };
let windows = 0;
Expand Down Expand Up @@ -802,7 +847,7 @@ async function fetchDeepInfraQuota(provider: string, config: OcxProviderConfig):
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
if (!data) return null;
const stripeBalance = toFiniteNumber(data.stripe_balance);
Expand Down Expand Up @@ -844,7 +889,7 @@ async function fetchNeuralwattQuota(provider: string, config: OcxProviderConfig)
? TERMINAL_QUOTA_FAILURE
: null;
}
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const data = asRecord(body?.data) ?? body;
const quota: ProviderQuota = { updatedAt: Date.now() };
let windows = 0;
Expand Down Expand Up @@ -1028,7 +1073,7 @@ async function fetchXaiWeeklyCredits(accessToken: string, userId: string): Promi
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (!response.ok) return null;
const parsed = parseXaiCreditsResponse(await response.json().catch(() => null));
const parsed = parseXaiCreditsResponse(await readQuotaJson(response));
if (!parsed) return null;
return {
weeklyPercent: parsed.percent,
Expand Down Expand Up @@ -1062,7 +1107,7 @@ async function fetchXaiQuota(provider: string): Promise<ProviderQuotaReport | nu
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (!response.ok) return null;
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const config = asRecord(body?.config);
if (!config) return null;
const limitCents = centsValue(config.monthlyLimit);
Expand Down Expand Up @@ -1108,7 +1153,7 @@ async function fetchAnthropicUsageQuota(accessToken: string): Promise<ProviderQu
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (!response.ok) return null;
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
if (!body) return null;
const fiveHour = parseClaudeBucket(body.five_hour);
const sevenDay = parseClaudeBucket(body.seven_day);
Expand Down Expand Up @@ -1521,7 +1566,7 @@ async function fetchKimiQuota(provider: string, config: OcxProviderConfig): Prom
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (!response.ok) return null;
const quota = parseKimiQuotaPayload(await response.json().catch(() => null));
const quota = parseKimiQuotaPayload(await readQuotaJson(response));
return quota ? report(provider, "kimi:usages", quota) : null;
}

Expand Down Expand Up @@ -1554,7 +1599,7 @@ async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport |
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (periodRes.ok) {
const body = asRecord(await periodRes.json().catch(() => null));
const body = asRecord(await readQuotaJson(periodRes));
const planUsage = asRecord(body?.planUsage);
if (planUsage) {
const resetAt = normalizeResetAt(body?.billingCycleEnd ?? planUsage.billingCycleEnd ?? body?.periodEnd);
Expand Down Expand Up @@ -1616,7 +1661,7 @@ async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport |
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (summaryRes.ok) {
const body = asRecord(await summaryRes.json().catch(() => null));
const body = asRecord(await readQuotaJson(summaryRes));
const individual = asRecord(body?.individualUsage);
const plan = asRecord(individual?.plan);
if (plan) {
Expand Down Expand Up @@ -1645,7 +1690,7 @@ async function fetchCursorQuota(provider: string): Promise<ProviderQuotaReport |
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (!response.ok) return null;
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
if (!body) return null;

// Prefer the gpt-4 bucket (historical "fast requests"); else first model with used+limit.
Expand Down Expand Up @@ -1759,7 +1804,7 @@ async function fetchAntigravityQuota(provider: string, config: OcxProviderConfig
signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
});
if (!response.ok) return null;
const body = asRecord(await response.json().catch(() => null));
const body = asRecord(await readQuotaJson(response));
const models = asRecord(body?.models);
if (!models) return null;

Expand Down
Loading
Loading