property-management-network/lib/ai/usage.ts

import { and, eq, gte, sql } from "drizzle-orm"
import { db } from "@/lib/db"
import { usage_events, profiles } from "@/lib/db/schema"
import { PLAN_LIMITS } from "@/lib/stripe/plans"
import type { Plan } from "@/types"

type QuotaResult =
  | { ok: true; used: number; limit: number }
  | { ok: false; status: number; error: string }

// Sentinel for "effectively unlimited" plans. PLAN_LIMITS uses Infinity for
// some limits; treat Infinity (or anything absurdly large) as uncapped so we
// skip the count but still record the event.
const UNLIMITED_THRESHOLD = 1_000_000_000

/**
 * Atomic, aggregate monthly AI quota gate.
 *
 * The cap is a true total across ALL AI features (we do NOT filter usage by
 * event_type when counting), and the check + insert run inside a single
 * transaction guarded by a per-user advisory lock, so concurrent requests
 * serialize and cannot race past the limit (kills the count-then-insert TOCTOU).
 */
export async function enforceAiQuota(
  userId: string,
  eventType: string
): Promise<QuotaResult> {
  const profile = await db.query.profiles.findFirst({
    where: eq(profiles.id, userId),
    columns: { plan: true },
  })

  const plan = (profile?.plan ?? "starter") as Plan
  const limit = PLAN_LIMITS[plan].maxAiCalls

  if (limit <= 0) {
    return { ok: false, status: 403, error: "AI features require a Pro plan or higher." }
  }

  const monthStart = new Date()
  monthStart.setUTCDate(1)
  monthStart.setUTCHours(0, 0, 0, 0)

  // Unlimited plans: skip counting, but still record the event for analytics.
  if (!Number.isFinite(limit) || limit >= UNLIMITED_THRESHOLD) {
    await db.insert(usage_events).values({ user_id: userId, event_type: eventType })
    const [{ count }] = await db
      .select({ count: sql<number>`count(*)::int` })
      .from(usage_events)
      .where(and(eq(usage_events.user_id, userId), gte(usage_events.created_at, monthStart.toISOString())))
    return { ok: true, used: count, limit }
  }

  // Sentinel thrown from inside the transaction to signal an over-limit
  // condition (returning a non-ok object) without committing the insert.
  const overLimit: QuotaResult = {
    ok: false,
    status: 429,
    error: "Monthly AI limit reached. Upgrade your plan for more.",
  }

  try {
    return await db.transaction(async (tx) => {
      // Serialize concurrent requests for this user so the count + insert is atomic.
      await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${userId}))`)

      const [{ count }] = await tx
        .select({ count: sql<number>`count(*)::int` })
        .from(usage_events)
        .where(
          and(
            eq(usage_events.user_id, userId),
            gte(usage_events.created_at, monthStart.toISOString())
          )
        )

      if (count >= limit) {
        throw overLimit
      }

      await tx.insert(usage_events).values({ user_id: userId, event_type: eventType })

      return { ok: true, used: count + 1, limit } as QuotaResult
    })
  } catch (err) {
    if (err === overLimit) return overLimit
    throw err
  }
}