import { dailyCounterKey, PRO_DAILY_QUOTA_LIMIT, PRO_DAILY_QUOTA_TTL_SECONDS, } from '../../server/_shared/pro-mcp-token'; import type { PipelineFn, QuotaRejected, QuotaReserved } from './types'; // --------------------------------------------------------------------------- // Daily quota helpers (Pro-only). INCR-first reservation runs synchronously // on the critical path BEFORE tool dispatch — never inside `waitUntil`. // On pre-dispatch cap rejection we best-effort DECR. Once dispatch begins, // callers keep the slot charged even if execution later errors or exceeds // budget. // // The cap itself is plan-driven (plan 2026-07-25-001 U3): the caller passes the // allowance resolved from the entitlement, and `PRO_DAILY_QUOTA_LIMIT` is the // fallback for anyone who can't supply one. // --------------------------------------------------------------------------- /** * Normalise a plan-resolved allowance into the value this module enforces. * * `null` (unlimited) passes through; a finite non-negative number is honoured * verbatim — including `0`, which is a real "no allowance" and must not be * mistaken for a missing one. EVERYTHING else — undefined, a legacy row with no * `planLimits`, NaN/Infinity, a negative, a stringified number — resolves to * `PRO_DAILY_QUOTA_LIMIT`. That direction is deliberate: an unreadable limit * must never buy a caller a HIGHER cap than the plan default. * * Exported because the settings-UI reader (`api/user/mcp-quota.ts`) must DISPLAY * exactly the limit this module ENFORCES. A second copy of this normalisation * would be the drift the endpoint's whole reason for existing is to prevent. */ export function resolveDailyLimit(planDailyLimit?: number | null): number | null { if (planDailyLimit === null) return null; if (typeof planDailyLimit === 'number' && Number.isFinite(planDailyLimit) && planDailyLimit >= 0) { return planDailyLimit; } return PRO_DAILY_QUOTA_LIMIT; } /** * Plans whose catalog `mcpCallsPerDay` must NOT drive the daily cap on the * pro (OAuth) MCP context. The KTD6 boundary is a PLAN boundary, not a * credential boundary: API-tier subscribers can mint pro OAuth tokens too * (tier>=1 + mcpAccess), and without this gate their catalog allowance * (1000/10000) would leak through the OAuth door while their `user_key` * stays hardcoded at 50. Raising API-tier MCP allowances is a deliberate * follow-up; until then both credential classes must agree on the cap. */ const API_TIER_MCP_CAPPED_PLAN_KEYS = new Set([ 'api_starter', 'api_starter_annual', 'api_business', 'api_business_annual', ]); /** * Gate a plan-resolved MCP allowance on plan family: API-tier plans report * `undefined` (→ the 50/day default via `resolveDailyLimit`); every other * plan's allowance passes through verbatim — pro/pro_business plan-driven * numbers, enterprise's `null` (unlimited), free's `0`. * * Shared by the enforcement path (`checkMcpEntitlementGate`) and the * settings display (`api/user/mcp-quota.ts`) so the number a user reads is * the number the reservation applies. */ export function resolvePlanDrivenMcpAllowance( planKey: string | undefined, mcpCallsPerDay: number | null | undefined, ): number | null | undefined { if (planKey && API_TIER_MCP_CAPPED_PLAN_KEYS.has(planKey)) return undefined; return mcpCallsPerDay; } export async function reserveQuota( userId: string, pipeline: PipelineFn, planDailyLimit?: number | null, ): Promise { // `null` = unlimited: the counter still moves (metering is not optional) but // the rejection branch below is skipped entirely. const limit = resolveDailyLimit(planDailyLimit); const key = dailyCounterKey(userId); if (!key) return { ok: false, reason: 'redis-unavailable' }; let pipeResult: Array<{ result: unknown }> | null; try { pipeResult = await pipeline([ ['INCR', key], ['EXPIRE', key, PRO_DAILY_QUOTA_TTL_SECONDS], ]); } catch { pipeResult = null; } if (!pipeResult || !Array.isArray(pipeResult) || pipeResult.length === 0) { // Hard cap correctness: NEVER dispatch on reservation failure. return { ok: false, reason: 'redis-unavailable' }; } const incrRaw = pipeResult[0]?.result; const newCount = typeof incrRaw === 'number' ? incrRaw : Number(incrRaw); if (!Number.isFinite(newCount) || newCount < 1) { return { ok: false, reason: 'redis-unavailable' }; } // Build idempotent rollback. `await rollback()` runs DECR once; subsequent // calls are no-ops. let rolledBack = false; const rollback = async (): Promise => { if (rolledBack) return; rolledBack = true; try { await pipeline([['DECR', key]]); } catch { // Best-effort: a transient Redis failure means the counter overshoots // by 1, which is the cost-protection-correct direction. } }; if (limit !== null && newCount > limit) { // Reject and roll back immediately so the floor stays at the limit // (or wherever concurrent rollbacks land it). await rollback(); // Counter-clamp (F4): if multiple DECR rollbacks have failed during // a Redis hiccup, the counter can overshoot indefinitely (e.g. land // at 2x the limit). Without clamping, every subsequent INCR for the // rest of the UTC day yields >limit → the user is locked out until // the 48h key TTL expires. The clamp target is the RESOLVED limit, // not the plan default — clamping a 250/day caller down to 50 would // hand them 200 free calls on the next Redis hiccup. // // After the rollback, peek at the post-DECR count via a single // best-effort INCR-then-DECR pair — if it's STILL above the limit, // we know the rollback didn't land. Force a defensive // `SET key KEEPTTL` so the next legitimate INCR (next UTC // day OR next request after the hiccup) starts at limit+1 → 429, // not limit+N → 429-forever. // // Why use INCR-then-DECR instead of GET? Keeps the helper to the // same pipeline contract (the tests' makePipelineMock supports // INCR/DECR/EXPIRE only) and avoids adding a new verb. The probe // costs one round-trip but only on the rejection path. if (newCount > limit + 1) { try { const probe = await pipeline([['INCR', key], ['DECR', key]]); const probeIncrRaw = probe?.[0]?.result; const postRollbackCount = typeof probeIncrRaw === 'number' ? probeIncrRaw - 1 : Number.NaN; if (Number.isFinite(postRollbackCount) && postRollbackCount > limit) { // Rollback chain has overshot — force the counter back to the // limit via SET KEEPTTL. This is fail-soft: a concurrent INCR // immediately after this SET will land at limit+1 and 429 // normally, which is the desired behavior. // // Use DECR repeatedly as the pipeline-supported clamp (avoids // adding a new verb to test mocks). DECR N times where N is // the overshoot delta. Cap at 100 DECRs to bound the worst- // case round-trip cost. const overshoot = postRollbackCount - limit; const decrs = Math.min(overshoot, 100); const clamp = Array.from({ length: decrs }, () => ['DECR', key] as Array); // Best-effort: failure here is the cost-protection-correct // direction (counter stays high → users 429, no DoS exposure). await pipeline(clamp).catch(() => {}); } } catch { // Probe failed — leave counter as-is. Worst case the user 429s // until UTC midnight; never under-cap, never DoS exposure. } } return { ok: false, reason: 'cap-exceeded', floor: limit }; } return { ok: true, newCount, rollback }; }