podcastdistributiona/lib/usage/meter.ts

import { prisma } from "@/lib/db";
import { periodKey } from "@/lib/utils";
import { UNLIMITED, type UsageMetric } from "@/lib/billing/plans";

export type OwnerType = "user" | "organization";

/**
 * METERING INVARIANT (reserve-then-verify; see security finding M2)
 * ----------------------------------------------------------------
 * Usage is RESERVED atomically up front by the enqueuing caller (server
 * actions / API route) via `reserveUsage` (through `reserveLimit`), BEFORE the
 * episode is created/enqueued. Reservation atomically increments the counter
 * and rejects (refunding itself) when the post-increment count exceeds the cap,
 * which closes the time-of-check/time-of-use race that let concurrent requests
 * all pass a plain read-check and then blow past the monthly cap.
 *
 * Consequences for the rest of the system:
 *   - The worker / generation pipeline DOES NOT increment usage. The metric was
 *     already counted at enqueue time. Double counting is therefore impossible.
 *   - On TERMINAL generation failure the worker REFUNDS the reserved metrics via
 *     `refundUsage`, so a failed job does not permanently consume quota.
 *   - Inline (synchronous) generations that don't go through the worker
 *     (e.g. regenerate-section, repurpose) likewise RESERVE once up front and
 *     never call `incrementUsage` afterwards; they refund on failure.
 *
 * In short: "the enqueuing caller reserves; the worker refunds on terminal
 * failure; the worker does not increment."
 */

/** Increment a monthly usage counter for a billing subject. */
export async function incrementUsage(
  ownerId: string,
  ownerType: OwnerType,
  metric: UsageMetric,
  by = 1
): Promise<void> {
  const key = periodKey(new Date());
  await prisma.usageRecord.upsert({
    where: { ownerId_periodKey_metric: { ownerId, periodKey: key, metric } },
    create: { ownerId, ownerType, periodKey: key, metric, count: by },
    update: { count: { increment: by } },
  });
}

/**
 * Atomically reserve `by` units of `metric` and verify the post-increment count
 * against `limit`. The upsert's `increment` is serialized by the
 * (ownerId, periodKey, metric) unique constraint and returns the row's
 * POST-increment count, so concurrent reservations can't both slip under the
 * cap. If the reservation would exceed `limit`, it is REFUNDED (decremented)
 * and `false` is returned. UNLIMITED (-1) is always allowed.
 *
 * Returns `true` when the reservation succeeded (quota consumed), `false` when
 * it was rejected (and already refunded — caller need not undo anything).
 */
export async function reserveUsage(
  ownerId: string,
  ownerType: OwnerType,
  metric: UsageMetric,
  limit: number,
  by = 1
): Promise<boolean> {
  const key = periodKey(new Date());
  const row = await prisma.usageRecord.upsert({
    where: { ownerId_periodKey_metric: { ownerId, periodKey: key, metric } },
    create: { ownerId, ownerType, periodKey: key, metric, count: by },
    update: { count: { increment: by } },
  });

  if (limit !== UNLIMITED && row.count > limit) {
    // Over cap — undo our own increment and reject.
    await refundUsage(ownerId, ownerType, metric, by);
    return false;
  }
  return true;
}

/**
 * Refund `by` units of a previously reserved metric (floored at 0). Used when a
 * reserved generation fails downstream, or by `reserveUsage` to undo a
 * rejected reservation.
 */
export async function refundUsage(
  ownerId: string,
  ownerType: OwnerType,
  metric: UsageMetric,
  by = 1
): Promise<void> {
  const key = periodKey(new Date());
  const existing = await prisma.usageRecord.findUnique({
    where: { ownerId_periodKey_metric: { ownerId, periodKey: key, metric } },
  });
  if (!existing) return;
  const next = Math.max(0, existing.count - by);
  await prisma.usageRecord.update({
    where: { ownerId_periodKey_metric: { ownerId, periodKey: key, metric } },
    data: { count: next },
  });
}

/** Current-period count for a single metric. */
export async function getUsage(
  ownerId: string,
  metric: UsageMetric,
  date = new Date()
): Promise<number> {
  const rec = await prisma.usageRecord.findUnique({
    where: { ownerId_periodKey_metric: { ownerId, periodKey: periodKey(date), metric } },
  });
  return rec?.count ?? 0;
}

/** Current-period counts for all metrics. */
export async function getUsageSummary(
  ownerId: string,
  date = new Date()
): Promise<Record<UsageMetric, number>> {
  const rows = await prisma.usageRecord.findMany({
    where: { ownerId, periodKey: periodKey(date) },
  });
  const summary: Record<UsageMetric, number> = { script: 0, audio: 0, art: 0, repurpose: 0 };
  for (const row of rows) {
    if (row.metric in summary) summary[row.metric as UsageMetric] = row.count;
  }
  return summary;
}