import { db } from "@epicure/db"; import { tierDefinitions, userUsage } from "@epicure/db"; import { eq, sql } from "@epicure/db"; function currentMonth() { const now = new Date(); return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`; } export type LimitKey = "recipe" | "aiCall" | "storage"; /** Sentinel stored in tier_definitions numeric columns to mean "no cap". */ export const UNLIMITED = -1; export class TierLimitError extends Error { constructor(public readonly limit: LimitKey, public readonly tier: string) { super(`Tier limit reached: ${limit} (tier: ${tier})`); this.name = "TierLimitError"; } } /** * Atomically checks the tier limit and increments the usage counter in a * single SQL statement, eliminating the TOCTOU race that existed when * checkTierLimit and incrementUsage were called separately. * * Throws TierLimitError if the limit has already been reached. * Use this instead of the separate checkTierLimit + incrementUsage pair * for "recipe" and "aiCall" keys. */ export async function checkAndIncrementTierLimit( userId: string, userTier: "free" | "pro", key: "recipe" | "aiCall" ): Promise { const [tierDef] = await db .select() .from(tierDefinitions) .where(eq(tierDefinitions.tier, userTier)); if (!tierDef) return; const month = currentMonth(); const id = `${userId}-${month}`; if (key === "aiCall") { const limit = tierDef.aiCallsPerMonth; const cap = limit === UNLIMITED ? sql`true` : sql`user_usage.ai_calls_used < ${limit}`; const result = await db.execute(sql` INSERT INTO user_usage (id, user_id, month, ai_calls_used, recipe_count, storage_used_mb) VALUES (${id}, ${userId}, ${month}, 1, 0, 0) ON CONFLICT (user_id, month) DO UPDATE SET ai_calls_used = user_usage.ai_calls_used + 1 WHERE ${cap} RETURNING ai_calls_used `); if (result.length === 0) { throw new TierLimitError("aiCall", userTier); } } else { const limit = tierDef.maxRecipes; const cap = limit === UNLIMITED ? sql`true` : sql`user_usage.recipe_count < ${limit}`; const result = await db.execute(sql` INSERT INTO user_usage (id, user_id, month, ai_calls_used, recipe_count, storage_used_mb) VALUES (${id}, ${userId}, ${month}, 0, 1, 0) ON CONFLICT (user_id, month) DO UPDATE SET recipe_count = user_usage.recipe_count + 1 WHERE ${cap} RETURNING recipe_count `); if (result.length === 0) { throw new TierLimitError("recipe", userTier); } } } /** * Refunds one aiCall credit for the current month. Call this when an AI * request failed after the quota was already charged (e.g. provider error) * so users aren't billed against their limit for a call that never succeeded. */ export async function refundAiCall(userId: string): Promise { const month = currentMonth(); await db.execute(sql` UPDATE user_usage SET ai_calls_used = GREATEST(ai_calls_used - 1, 0) WHERE user_id = ${userId} AND month = ${month} `); } export async function incrementUsage( userId: string, key: LimitKey, amount = 1 ): Promise { const month = currentMonth(); const id = `${userId}-${month}`; const initialValues = { id, userId, month, aiCallsUsed: key === "aiCall" ? amount : 0, recipeCount: key === "recipe" ? amount : 0, storageUsedMb: key === "storage" ? amount : 0, }; const incrementSet = key === "aiCall" ? { aiCallsUsed: sql`${userUsage.aiCallsUsed} + ${amount}` } : key === "recipe" ? { recipeCount: sql`${userUsage.recipeCount} + ${amount}` } : { storageUsedMb: sql`${userUsage.storageUsedMb} + ${amount}` }; await db .insert(userUsage) .values(initialValues) .onConflictDoUpdate({ target: [userUsage.userId, userUsage.month], set: incrementSet, }); }