import { redis } from './redis'; import { prisma } from './prisma'; import { hasAnyActiveByok } from './byok'; import { getCurrentPeriodKey, getRedisKey, parseRedisInt, isValidFeature, } from './quota-utils'; type SubscriptionTier = 'BASIC' | 'PRO' | 'BUSINESS' | 'ENTERPRISE'; export interface EntitlementResult { allowed: boolean; remaining: number; limit: number; tier: SubscriptionTier; reason?: 'QUOTA_EXCEEDED' | 'TIER_LIMITED' | 'FEATURE_NOT_AVAILABLE'; message?: string; upgradeTier?: 'PRO' | 'BUSINESS'; byokConfigured?: boolean; } 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 TIER_LIMITS: Record> = { BASIC: { semantic_search: 30, auto_tag: 20, auto_title: 10, brainstorm_create: 1, brainstorm_expand: 10, brainstorm_enrich: 20, }, PRO: { semantic_search: 100, auto_tag: 200, auto_title: 200, reformulate: 50, chat: 100, brainstorm_create: 30, brainstorm_expand: 100, brainstorm_enrich: 200, }, BUSINESS: { semantic_search: 1000, auto_tag: 1000, auto_title: 1000, reformulate: 500, chat: 1000, brainstorm_create: 200, brainstorm_expand: 500, brainstorm_enrich: 1000, }, ENTERPRISE: { semantic_search: 'unlimited', auto_tag: 'unlimited', auto_title: 'unlimited', reformulate: 'unlimited', chat: 'unlimited', brainstorm_create: 'unlimited', brainstorm_expand: 'unlimited', brainstorm_enrich: 'unlimited', }, }; const TTL_SECONDS = 90 * 24 * 60 * 60; const INCREMENT_LUA = ` local current = tonumber(redis.call('GET', KEYS[1]) or '0') local ttl = tonumber(ARGV[1]) 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 `; 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 `; function getLimit(tier: SubscriptionTier, feature: string): number | undefined { const tierLimits = TIER_LIMITS[tier]; const limit = tierLimits?.[feature]; if (limit === 'unlimited') return Infinity; if (limit === undefined) return undefined; return limit; } 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 = getLimit(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 current = parseRedisInt(currentStr); 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, allowing request (fail-open):', err); 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 = getLimit(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 newCount = await redis.eval( RESERVE_LUA, 1, key, String(limit), String(TTL_SECONDS), ) as number; if (newCount === -1) { const byokConfigured = await hasAnyActiveByok(userId); throw new QuotaExceededError( tier === 'BASIC' ? 'PRO' : 'BUSINESS', feature, limit, limit, byokConfigured, { currentTier: tier }, ); } } catch (err) { if (err instanceof QuotaExceededError) throw err; console.error('[entitlements] Redis unavailable, allowing request (fail-open):', err); } } /** 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 period = getCurrentPeriodKey(); const features = Object.keys(TIER_LIMITS[tier]); if (features.length === 0) return {}; const keys = features.map((f) => `usage:${userId}:${f}:${period}`); try { const values = await redis.mget(...keys); const result: Record = {}; for (let i = 0; i < features.length; i++) { const feature = features[i]; const limit = getLimit(tier, feature) ?? 0; const current = parseRedisInt(values[i]); 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 = getLimit(tier, feature) ?? 0; result[feature] = { remaining: limit, limit, used: 0 }; } return result; } } export { TIER_LIMITS, getLimit }; export type { SubscriptionTier };