Leon Serfaty
127 lines
4.7 KiB
TypeScript
127 lines
4.7 KiB
TypeScript
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;
|
|
}
|