import { redis } from './redis'; import { prisma } from './prisma'; import { hasAnyActiveByok } from './byok'; import { getCurrentPeriodKey, getRedisKey, parseRedisInt, isValidFeature, } from './quota-utils'; import { getLimitAsync, getTierFeaturesAsync, invalidateEntitlementCache, TIER_LIMITS, type SubscriptionTier, } from './plan-entitlements'; export interface EntitlementResult { allowed: boolean; remaining: number; limit: number; tier: SubscriptionTier; reason?: 'QUOTA_EXCEEDED' | 'TIER_LIMITED' | 'FEATURE_NOT_AVAILABLE' | 'SERVICE_UNAVAILABLE'; message?: string; upgradeTier?: 'PRO' | 'BUSINESS'; byokConfigured?: boolean; } export class QuotaServiceUnavailableError extends Error { code = 'QUOTA_SERVICE_UNAVAILABLE'; constructor(message = 'Quota service temporarily unavailable') { super(message); } } export class QuotaExceededError extends Error { code = 'QUOTA_EXCEEDED'; upgradeTier: 'PRO' | 'BUSINESS'; feature: string; currentQuota: number; usedQuota: number; byokConfigured: boolean; billingOwnerId?: string; triggeredByUserId?: string; isGuestActor?: boolean; currentTier?: string; constructor( upgradeTier: 'PRO' | 'BUSINESS', feature: string, currentQuota: number, usedQuota: number, byokConfigured: boolean = false, sessionMeta?: { billingOwnerId?: string; triggeredByUserId?: string; isGuestActor?: boolean; currentTier?: string; }, ) { super(`Quota exceeded for ${feature}`); this.upgradeTier = upgradeTier; this.feature = feature; this.currentQuota = currentQuota; this.usedQuota = usedQuota; this.byokConfigured = byokConfigured; this.billingOwnerId = sessionMeta?.billingOwnerId; this.triggeredByUserId = sessionMeta?.triggeredByUserId; this.isGuestActor = sessionMeta?.isGuestActor; this.currentTier = sessionMeta?.currentTier; } toJSON() { return { error: this.code, feature: this.feature, upgradeTier: this.upgradeTier, byokConfigured: this.byokConfigured, isGuestActor: this.isGuestActor ?? false, currentTier: this.currentTier, }; } } const TTL_SECONDS = 90 * 24 * 60 * 60; function getPeriodDates(): { period: string; periodStart: Date; periodEnd: Date } { const period = getCurrentPeriodKey(); const periodStart = new Date(`${period}-01T00:00:00.000Z`); const periodEnd = new Date(Date.UTC(periodStart.getUTCFullYear(), periodStart.getUTCMonth() + 1, 1)); return { period, periodStart, periodEnd }; } async function getDatabaseUsageCounts( userId: string, features: string[], ): Promise> { if (features.length === 0) return {}; const { periodStart } = getPeriodDates(); const rows = await prisma.usageLog.findMany({ where: { userId, periodStart, feature: { in: features } }, select: { feature: true, requestsCount: true }, }); return Object.fromEntries(rows.map((r) => [r.feature, r.requestsCount])); } async function mirrorUsageCountToDatabase( userId: string, feature: string, requestsCount: number, ): Promise { const { periodStart, periodEnd } = getPeriodDates(); await prisma.usageLog.upsert({ where: { userId_feature_periodStart: { userId, feature, periodStart }, }, create: { userId, feature, periodStart, periodEnd, requestsCount, tokensUsed: 0, syncedAt: new Date(), }, update: { requestsCount, syncedAt: new Date(), }, }).catch((err) => { console.error('[entitlements] Failed to mirror usage to DB:', err); }); } /** Fallback when Redis is unavailable — atomic increment in PostgreSQL. */ async function reserveUsageInDatabase( userId: string, feature: string, limit: number, ): Promise { const { periodStart, periodEnd } = getPeriodDates(); return prisma.$transaction(async (tx) => { const existing = await tx.usageLog.findUnique({ where: { userId_feature_periodStart: { userId, feature, periodStart } }, }); const current = existing?.requestsCount ?? 0; if (current >= limit) return -1; const updated = await tx.usageLog.upsert({ where: { userId_feature_periodStart: { userId, feature, periodStart } }, create: { userId, feature, periodStart, periodEnd, requestsCount: 1, tokensUsed: 0, syncedAt: new Date(), }, update: { requestsCount: { increment: 1 }, syncedAt: new Date(), }, }); return updated.requestsCount; }); } function parseReserveResult(raw: unknown): number { const n = typeof raw === 'number' ? raw : Number(raw); return Number.isFinite(n) ? n : NaN; } function shouldFailClosedOnRedisError(): boolean { return process.env.NODE_ENV === 'production'; } const INCREMENT_BY_LUA = ` local count = tonumber(ARGV[1]) or 1 local ttl = tonumber(ARGV[2]) redis.call('INCRBY', KEYS[1], count) local ttlResult = redis.call('TTL', KEYS[1]) if ttlResult == -1 then redis.call('EXPIRE', KEYS[1], ttl) end local newCount = tonumber(redis.call('GET', KEYS[1])) return newCount `; const RESERVE_LUA = ` local limit = tonumber(ARGV[1]) local ttl = tonumber(ARGV[2]) local current = tonumber(redis.call('GET', KEYS[1]) or '0') if current >= limit then return -1 end redis.call('INCRBY', KEYS[1], 1) local ttlResult = redis.call('TTL', KEYS[1]) if ttlResult == -1 then redis.call('EXPIRE', KEYS[1], ttl) end local newCount = tonumber(redis.call('GET', KEYS[1])) return newCount `; export async function getUserInfo( userId: string, ): Promise<{ tier: SubscriptionTier; status: string; currentPeriodEnd?: Date }> { const subscription = await prisma.subscription.findUnique({ where: { userId }, }); if (!subscription) return { tier: 'BASIC', status: 'INACTIVE' }; return { tier: subscription.tier as SubscriptionTier, status: subscription.status, currentPeriodEnd: subscription.currentPeriodEnd, }; } export async function getEffectiveTier( userId: string, ): Promise { const { tier, status, currentPeriodEnd } = await getUserInfo(userId); if (status === 'ACTIVE' || status === 'TRIALING') return tier; if ((status === 'PAST_DUE' || status === 'CANCELED') && currentPeriodEnd) { if (new Date() < new Date(currentPeriodEnd)) return tier; } return 'BASIC'; } export async function canUseFeature( userId: string, feature: string, ): Promise { if (!isValidFeature(feature)) { return { allowed: false, remaining: 0, limit: 0, tier: 'BASIC', reason: 'TIER_LIMITED', message: `Unknown feature: ${feature}`, }; } const tier = await getEffectiveTier(userId); const limit = await getLimitAsync(tier, feature); if (limit === undefined) { return { allowed: false, remaining: 0, limit: 0, tier, reason: 'FEATURE_NOT_AVAILABLE', message: `Feature "${feature}" is not available on the ${tier} plan. Upgrade to unlock it.`, upgradeTier: tier === 'BASIC' ? 'PRO' : 'BUSINESS', }; } if (limit === Infinity) { return { allowed: true, remaining: Infinity, limit: Infinity, tier }; } try { const key = getRedisKey(userId, feature); const currentStr = await redis.get(key); const redisCurrent = parseRedisInt(currentStr); const dbCounts = await getDatabaseUsageCounts(userId, [feature]); const current = Math.max(redisCurrent, dbCounts[feature] ?? 0); const allowed = current < limit; if (!allowed) { const byokConfigured = await hasAnyActiveByok(userId); return { allowed: false, remaining: 0, limit, tier, reason: 'QUOTA_EXCEEDED', upgradeTier: tier === 'BASIC' ? 'PRO' : 'BUSINESS', byokConfigured, }; } return { allowed: true, remaining: Math.max(0, limit - current), limit, tier, byokConfigured: await hasAnyActiveByok(userId), }; } catch (err) { console.error('[entitlements] Redis unavailable:', err); if (shouldFailClosedOnRedisError()) { return { allowed: false, remaining: 0, limit, tier, reason: 'SERVICE_UNAVAILABLE', message: 'Quota service temporarily unavailable. Please try again later.', }; } return { allowed: true, remaining: limit, limit, tier }; } } export async function checkEntitlementOrThrow( userId: string, feature: string, ): Promise { const result = await canUseFeature(userId, feature); if (!result.allowed) { throw new QuotaExceededError( result.upgradeTier ?? (result.tier === 'BASIC' ? 'PRO' : 'BUSINESS'), feature, result.limit, result.limit > 0 ? result.limit - result.remaining : 0, result.byokConfigured ?? false, { currentTier: result.tier }, ); } } export async function reserveUsageOrThrow( userId: string, feature: string, ): Promise { if (!isValidFeature(feature)) { throw new QuotaExceededError('PRO', feature, 0, 0, false, { currentTier: 'BASIC' }); } const tier = await getEffectiveTier(userId); const limit = await getLimitAsync(tier, feature); if (limit === undefined) { throw new QuotaExceededError( tier === 'BASIC' ? 'PRO' : 'BUSINESS', feature, 0, 0, false, { currentTier: tier }, ); } if (limit === Infinity) return; try { const key = getRedisKey(userId, feature); const raw = await redis.eval( RESERVE_LUA, 1, key, String(limit), String(TTL_SECONDS), ); const newCount = parseReserveResult(raw); if (newCount === -1) { const byokConfigured = await hasAnyActiveByok(userId); throw new QuotaExceededError( tier === 'BASIC' ? 'PRO' : 'BUSINESS', feature, limit, limit, byokConfigured, { currentTier: tier }, ); } if (!Number.isFinite(newCount) || newCount < 1) { throw new Error(`Invalid Redis reserve result: ${String(raw)}`); } await mirrorUsageCountToDatabase(userId, feature, newCount); } catch (err) { if (err instanceof QuotaExceededError) throw err; console.error('[entitlements] Redis reserve failed, trying DB fallback:', err); try { const dbCount = await reserveUsageInDatabase(userId, feature, limit); if (dbCount === -1) { const byokConfigured = await hasAnyActiveByok(userId); throw new QuotaExceededError( tier === 'BASIC' ? 'PRO' : 'BUSINESS', feature, limit, limit, byokConfigured, { currentTier: tier }, ); } return; } catch (dbErr) { if (dbErr instanceof QuotaExceededError) throw dbErr; console.error('[entitlements] DB reserve fallback failed:', dbErr); } if (shouldFailClosedOnRedisError()) { throw new QuotaServiceUnavailableError(); } } } /** Host-pays: bill session owner, attach actor metadata for 402 responses (Story 3.4). */ export async function checkSessionEntitlementOrThrow( billingOwnerId: string, triggeredByUserId: string, isGuestActor: boolean, feature: string, ): Promise { try { await reserveUsageOrThrow(billingOwnerId, feature); } catch (err) { if (err instanceof QuotaExceededError) { err.billingOwnerId = billingOwnerId; err.triggeredByUserId = triggeredByUserId; err.isGuestActor = isGuestActor; } throw err; } } export function incrementUsageAsync(userId: string, feature: string, count: number = 1): Promise { if (!isValidFeature(feature)) return Promise.resolve(); const key = getRedisKey(userId, feature); return redis.eval(INCREMENT_BY_LUA, 1, key, String(count), String(TTL_SECONDS)).then(() => {}).catch((err) => { console.error('[entitlements] Async increment failed:', err); }); } export async function getUserQuotas( userId: string, ): Promise> { const tier = await getEffectiveTier(userId); const features = await getTierFeaturesAsync(tier); const period = getCurrentPeriodKey(); if (features.length === 0) return {}; const keys = features.map((f) => `usage:${userId}:${f}:${period}`); try { const values = await redis.mget(...keys); const dbCounts = await getDatabaseUsageCounts(userId, features); const result: Record = {}; for (let i = 0; i < features.length; i++) { const feature = features[i]; const limit = (await getLimitAsync(tier, feature)) ?? 0; const redisCurrent = parseRedisInt(values[i]); const dbCurrent = dbCounts[feature] ?? 0; const current = Math.max(redisCurrent, dbCurrent); result[feature] = { remaining: limit === Infinity ? Infinity : Math.max(0, limit - current), limit, used: current, }; } return result; } catch (err) { console.error('[entitlements] getUserQuotas Redis error:', err); const result: Record = {}; for (const feature of features) { const limit = (await getLimitAsync(tier, feature)) ?? 0; result[feature] = { remaining: limit, limit, used: 0 }; } return result; } } export { TIER_LIMITS, invalidateEntitlementCache }; export type { SubscriptionTier }; /** @deprecated Use getLimitAsync — sync helper for tests only */ export async function getLimit(tier: SubscriptionTier, feature: string): Promise { return getLimitAsync(tier, feature); }