usage-meter-test 0.1.0

package/dist/providers/prismaProvider.js

@@ -0,0 +1,323 @@
+function toAccount(row) {
+    return {
+        userId: row.userId,
+        balance: BigInt(row.balance),
+        frozen: BigInt(row.frozen),
+        updatedAt: new Date(row.updatedAt),
+    };
+}
+function toHold(row) {
+    const base = {
+        userId: row.userId,
+        requestId: row.requestId,
+        idempotencyKey: row.idempotencyKey,
+        amount: BigInt(row.amount),
+        status: row.status,
+        expiresAt: new Date(row.expiresAt),
+        createdAt: new Date(row.createdAt),
+        updatedAt: new Date(row.updatedAt),
+    };
+    if (row.capturedIdempotencyKey != null)
+        base.capturedIdempotencyKey = row.capturedIdempotencyKey;
+    return base;
+}
+function toLedger(row) {
+    const base = {
+        userId: row.userId,
+        type: row.type,
+        amount: BigInt(row.amount),
+        idempotencyKey: row.idempotencyKey,
+        createdAt: new Date(row.createdAt),
+    };
+    if (row.referenceId != null)
+        base.referenceId = row.referenceId;
+    if (row.meta != null)
+        base.meta = row.meta;
+    return base;
+}
+function toPolicy(row) {
+    return {
+        policyId: row.id,
+        period: row.period,
+        limitCredits: BigInt(row.limitCredits),
+        limitRequests: row.limitRequests ?? undefined,
+    };
+}
+function toBucket(row) {
+    return {
+        userId: row.userId,
+        policyId: row.policyId,
+        periodStart: new Date(row.periodStart),
+        periodEnd: new Date(row.periodEnd),
+        usedCredits: BigInt(row.usedCredits),
+        reservedCredits: BigInt(row.reservedCredits),
+        usedRequests: row.usedRequests,
+        reservedRequests: row.reservedRequests,
+        updatedAt: new Date(row.updatedAt),
+    };
+}
+function computePeriod(period, now) {
+    const d = new Date(now);
+    if (period === "daily") {
+        const start = new Date(Date.UTC(d.getUTCFullYear(), d.getUTCMonth(), d.getUTCDate(), 0, 0, 0, 0));
+        const end = new Date(Date.UTC(d.getUTCFullYear(), d.getUTCMonth(), d.getUTCDate() + 1, 0, 0, 0, 0));
+        return { start, end };
+    }
+    const start = new Date(Date.UTC(d.getUTCFullYear(), d.getUTCMonth(), 1, 0, 0, 0, 0));
+    const end = new Date(Date.UTC(d.getUTCFullYear(), d.getUTCMonth() + 1, 1, 0, 0, 0, 0));
+    return { start, end };
+}
+async function resolvePlanKeyForUserId(userId, tx, options) {
+    if (options?.resolvePlanKey)
+        return options.resolvePlanKey(userId, tx);
+    const userDelegate = tx.user;
+    if (!userDelegate?.findUnique)
+        return "default";
+    const isNumericId = /^[0-9]+$/.test(userId);
+    const row = isNumericId
+        ? await userDelegate.findUnique({ where: { id: Number(userId) } })
+        : await userDelegate.findUnique({ where: { uuid: userId } });
+    const planKey = row?.planKey;
+    return typeof planKey === "string" && planKey.trim() ? planKey.trim() : "default";
+}
+function isRetryableTransactionError(err) {
+    const anyErr = err;
+    // Prisma errors commonly expose `code` (e.g., P2034 for deadlock / write conflict).
+    // Different DBs map to different codes; keep this conservative.
+    return anyErr?.code === "P2034";
+}
+export function createPrismaPersistence(prisma, options) {
+    const clock = options?.now ?? (() => new Date());
+    const maxRetries = options?.maxRetries ?? 3;
+    const isolationLevel = options?.isolationLevel ?? "Serializable";
+    return {
+        now: () => clock(),
+        withTransaction: async (fn) => {
+            let attempt = 0;
+            // eslint-disable-next-line no-constant-condition
+            while (true) {
+                attempt += 1;
+                try {
+                    return await prisma.$transaction(async (clientTx) => {
+                        const accounts = {
+                            getOrCreate: async (userId) => {
+                                const row = await clientTx.creditAccount.upsert({
+                                    where: { userId },
+                                    create: { userId, balance: 0n, frozen: 0n },
+                                    update: {},
+                                });
+                                return toAccount(row);
+                            },
+                            update: async (userId, patch) => {
+                                // Atomic compare-and-set using `version` to avoid lost updates under concurrency.
+                                // Requires `CreditAccount.version` in the Prisma schema.
+                                const maxCasRetries = 5;
+                                for (let i = 0; i < maxCasRetries; i += 1) {
+                                    const current = await clientTx.creditAccount.findUnique({ where: { userId } });
+                                    if (!current) {
+                                        const created = await clientTx.creditAccount.create({
+                                            data: {
+                                                userId,
+                                                balance: patch.balance ?? 0n,
+                                                frozen: patch.frozen ?? 0n,
+                                                version: 0,
+                                            },
+                                        });
+                                        return toAccount(created);
+                                    }
+                                    const nextBalance = patch.balance ?? current.balance;
+                                    const nextFrozen = patch.frozen ?? current.frozen;
+                                    const res = await clientTx.creditAccount.updateMany({
+                                        where: { userId, version: current.version },
+                                        data: { balance: nextBalance, frozen: nextFrozen, version: { increment: 1 } },
+                                    });
+                                    if (res.count === 1) {
+                                        const updated = await clientTx.creditAccount.findUnique({ where: { userId } });
+                                        if (!updated)
+                                            throw new Error("CreditAccount disappeared after update");
+                                        return toAccount(updated);
+                                    }
+                                }
+                                throw new Error("CreditAccount CAS update failed after retries");
+                            },
+                        };
+                        const holds = {
+                            getByRequestId: async (userId, requestId) => {
+                                const row = await clientTx.creditHold.findUnique({
+                                    where: { userId_requestId: { userId, requestId } },
+                                });
+                                return row ? toHold(row) : null;
+                            },
+                            create: async (hold) => {
+                                const row = await clientTx.creditHold.create({
+                                    data: {
+                                        userId: hold.userId,
+                                        requestId: hold.requestId,
+                                        idempotencyKey: hold.idempotencyKey,
+                                        amount: hold.amount,
+                                        status: hold.status,
+                                        expiresAt: hold.expiresAt,
+                                    },
+                                });
+                                return toHold(row);
+                            },
+                            update: async (userId, requestId, patch) => {
+                                const row = await clientTx.creditHold.update({
+                                    where: { userId_requestId: { userId, requestId } },
+                                    data: {
+                                        ...(patch.status ? { status: patch.status } : {}),
+                                        ...(patch.expiresAt ? { expiresAt: patch.expiresAt } : {}),
+                                    },
+                                });
+                                return toHold(row);
+                            },
+                            lockCapture: async (userId, requestId, capturedIdempotencyKey, now) => {
+                                await clientTx.creditHold.updateMany({
+                                    where: {
+                                        userId,
+                                        requestId,
+                                        status: "held",
+                                        expiresAt: { gt: now },
+                                        capturedIdempotencyKey: null,
+                                    },
+                                    data: { capturedIdempotencyKey },
+                                });
+                                const row = await clientTx.creditHold.findUnique({ where: { userId_requestId: { userId, requestId } } });
+                                if (!row)
+                                    throw new Error("Hold not found");
+                                return toHold(row);
+                            },
+                            finalizeCapture: async (userId, requestId, capturedIdempotencyKey) => {
+                                await clientTx.creditHold.updateMany({
+                                    where: {
+                                        userId,
+                                        requestId,
+                                        capturedIdempotencyKey,
+                                    },
+                                    data: { status: "captured" },
+                                });
+                                const row = await clientTx.creditHold.findUnique({ where: { userId_requestId: { userId, requestId } } });
+                                if (!row)
+                                    throw new Error("Hold not found");
+                                return toHold(row);
+                            },
+                        };
+                        const ledger = {
+                            getByIdempotencyKey: async (userId, idempotencyKey) => {
+                                const row = await clientTx.creditLedger.findUnique({
+                                    where: { userId_idempotencyKey: { userId, idempotencyKey } },
+                                });
+                                return row ? toLedger(row) : null;
+                            },
+                            append: async (entry) => {
+                                try {
+                                    const row = await clientTx.creditLedger.create({
+                                        data: {
+                                            userId: entry.userId,
+                                            idempotencyKey: entry.idempotencyKey,
+                                            type: entry.type,
+                                            amount: entry.amount,
+                                            referenceId: entry.referenceId ?? null,
+                                            meta: entry.meta ?? null,
+                                        },
+                                    });
+                                    return toLedger(row);
+                                }
+                                catch (e) {
+                                    // If a concurrent request created the same idempotency key, read it back.
+                                    const existing = await clientTx.creditLedger.findUnique({
+                                        where: { userId_idempotencyKey: { userId: entry.userId, idempotencyKey: entry.idempotencyKey } },
+                                    });
+                                    if (existing)
+                                        return toLedger(existing);
+                                    throw e;
+                                }
+                            },
+                        };
+                        const quota = {
+                            getActivePolicy: async (userId) => {
+                                const planKey = await resolvePlanKeyForUserId(userId, clientTx, options);
+                                const row = await clientTx.planQuotaPolicy.findFirst({
+                                    where: { planKey, active: true },
+                                    orderBy: { updatedAt: "desc" },
+                                });
+                                return row ? toPolicy(row) : null;
+                            },
+                            getOrCreateBucket: async (userId, policy, now) => {
+                                const { start, end } = computePeriod(policy.period, now);
+                                const row = await clientTx.quotaBucket.upsert({
+                                    where: { userId_policyId_periodStart: { userId, policyId: policy.policyId, periodStart: start } },
+                                    create: {
+                                        userId,
+                                        policyId: policy.policyId,
+                                        periodStart: start,
+                                        periodEnd: end,
+                                        usedCredits: 0n,
+                                        reservedCredits: 0n,
+                                        usedRequests: 0,
+                                        reservedRequests: 0,
+                                    },
+                                    update: {},
+                                });
+                                return toBucket(row);
+                            },
+                            updateBucket: async (userId, policyId, periodStart, patch) => {
+                                // Atomic compare-and-set using `version`.
+                                const maxCasRetries = 5;
+                                for (let i = 0; i < maxCasRetries; i += 1) {
+                                    const current = await clientTx.quotaBucket.findUnique({
+                                        where: { userId_policyId_periodStart: { userId, policyId, periodStart } },
+                                    });
+                                    if (!current)
+                                        throw new Error("Quota bucket not found");
+                                    const nextUsedCredits = patch.usedCredits ?? current.usedCredits;
+                                    const nextReservedCredits = patch.reservedCredits ?? current.reservedCredits;
+                                    const nextUsedRequests = patch.usedRequests ?? current.usedRequests;
+                                    const nextReservedRequests = patch.reservedRequests ?? current.reservedRequests;
+                                    const res = await clientTx.quotaBucket.updateMany({
+                                        where: { userId, policyId, periodStart, version: current.version },
+                                        data: {
+                                            usedCredits: nextUsedCredits,
+                                            reservedCredits: nextReservedCredits,
+                                            usedRequests: nextUsedRequests,
+                                            reservedRequests: nextReservedRequests,
+                                            version: { increment: 1 },
+                                        },
+                                    });
+                                    if (res.count === 1) {
+                                        const updated = await clientTx.quotaBucket.findUnique({
+                                            where: { userId_policyId_periodStart: { userId, policyId, periodStart } },
+                                        });
+                                        if (!updated)
+                                            throw new Error("Quota bucket disappeared after update");
+                                        return toBucket(updated);
+                                    }
+                                }
+                                throw new Error("QuotaBucket CAS update failed after retries");
+                            },
+                        };
+                        const usage = {
+                            recordUsage: async (row) => {
+                                // Optional: integrate with your app's logging/audit table.
+                                void row;
+                            },
+                        };
+                        const meteringTx = {
+                            accounts,
+                            holds,
+                            ledger,
+                            quota,
+                            usage,
+                        };
+                        return fn(meteringTx);
+                    }, { isolationLevel });
+                }
+                catch (err) {
+                    if (attempt <= maxRetries && isRetryableTransactionError(err))
+                        continue;
+                    throw err;
+                }
+            }
+        },
+    };
+}

package/dist/types.d.ts

@@ -0,0 +1,168 @@
+export type UserId = string;
+export type RequestId = string;
+export type Credits = bigint;
+export type IdempotencyKey = string;
+export type Bps = bigint;
+export type Period = "daily" | "monthly";
+export type LedgerEntryType = "grant" | "purchase" | "charge" | "refund" | "adjust" | "expire";
+export type HoldStatus = "held" | "captured" | "released" | "expired";
+export interface CreditAccount {
+    userId: UserId;
+    balance: Credits;
+    frozen: Credits;
+    updatedAt: Date;
+}
+export interface CreditHold {
+    userId: UserId;
+    requestId: RequestId;
+    idempotencyKey: IdempotencyKey;
+    capturedIdempotencyKey?: IdempotencyKey;
+    amount: Credits;
+    status: HoldStatus;
+    expiresAt: Date;
+    createdAt: Date;
+    updatedAt: Date;
+}
+export interface LedgerEntry {
+    userId: UserId;
+    type: LedgerEntryType;
+    amount: Credits;
+    referenceId?: string;
+    idempotencyKey: IdempotencyKey;
+    meta?: Record<string, unknown>;
+    createdAt: Date;
+}
+export interface QuotaPolicy {
+    policyId: string;
+    period: Period;
+    limitCredits: Credits;
+    limitRequests?: number;
+}
+export interface QuotaBucket {
+    userId: UserId;
+    policyId: string;
+    periodStart: Date;
+    periodEnd: Date;
+    usedCredits: Credits;
+    reservedCredits: Credits;
+    usedRequests: number;
+    reservedRequests: number;
+    updatedAt: Date;
+}
+export interface UsageTokens {
+    inputTokens: number;
+    outputTokens: number;
+}
+export interface ToolUsage {
+    tool: string;
+    units: number;
+}
+export interface UsageFacts {
+    requestId: RequestId;
+    model: string;
+    tokens: UsageTokens;
+    tools?: ToolUsage[];
+    meta?: Record<string, unknown>;
+}
+export interface EstimateRequest {
+    userId: UserId;
+    requestId: RequestId;
+    model: string;
+    tokens: UsageTokens;
+    maxOutputTokens: number;
+    tools?: ToolUsage[];
+    riskBps?: Bps;
+}
+export interface EstimateResult {
+    estimatedCost: Credits;
+    maxPossibleCost: Credits;
+    breakdown: PriceBreakdown;
+}
+export interface HoldRequest {
+    userId: UserId;
+    requestId: RequestId;
+    amount: Credits;
+    idempotencyKey: IdempotencyKey;
+    expiresAt: Date;
+    meta?: Record<string, unknown>;
+}
+export interface HoldResult {
+    hold: CreditHold;
+    account: CreditAccount;
+    quotaBucket?: QuotaBucket;
+}
+export interface CaptureRequest {
+    userId: UserId;
+    requestId: RequestId;
+    actualCost: Credits;
+    idempotencyKey: IdempotencyKey;
+    usage: UsageFacts;
+}
+export interface CaptureResult {
+    charged: Credits;
+    hold: CreditHold;
+    account: CreditAccount;
+    ledgerEntry: LedgerEntry;
+    quotaBucket?: QuotaBucket;
+}
+export interface ReleaseRequest {
+    userId: UserId;
+    requestId: RequestId;
+    reason: string;
+}
+export interface ReleaseResult {
+    hold: CreditHold;
+    account: CreditAccount;
+    quotaBucket?: QuotaBucket;
+}
+export interface PriceBreakdown {
+    model: string;
+    inputCost: Credits;
+    outputCost: Credits;
+    toolCost: Credits;
+    minCostApplied: boolean;
+    modelMultiplierBps: Bps;
+    riskMultiplierBps: Bps;
+    totalBeforeMultipliers: Credits;
+    totalAfterMultipliers: Credits;
+    totalFinal: Credits;
+}
+export interface MeteringTransaction {
+    accounts: CreditAccountStore;
+    holds: CreditHoldStore;
+    ledger: CreditLedgerStore;
+    quota?: QuotaStore;
+    usage?: UsageStore;
+}
+export interface MeteringPersistence {
+    withTransaction<T>(fn: (tx: MeteringTransaction) => Promise<T>): Promise<T>;
+    now(): Date;
+}
+export interface CreditAccountStore {
+    getOrCreate(userId: UserId): Promise<CreditAccount>;
+    update(userId: UserId, patch: Partial<Pick<CreditAccount, "balance" | "frozen">>): Promise<CreditAccount>;
+}
+export interface CreditHoldStore {
+    getByRequestId(userId: UserId, requestId: RequestId): Promise<CreditHold | null>;
+    create(hold: CreditHold): Promise<CreditHold>;
+    update(userId: UserId, requestId: RequestId, patch: Partial<Pick<CreditHold, "status" | "expiresAt">>): Promise<CreditHold>;
+    lockCapture(userId: UserId, requestId: RequestId, capturedIdempotencyKey: IdempotencyKey, now: Date): Promise<CreditHold>;
+    finalizeCapture(userId: UserId, requestId: RequestId, capturedIdempotencyKey: IdempotencyKey): Promise<CreditHold>;
+}
+export interface CreditLedgerStore {
+    getByIdempotencyKey(userId: UserId, idempotencyKey: IdempotencyKey): Promise<LedgerEntry | null>;
+    append(entry: LedgerEntry): Promise<LedgerEntry>;
+}
+export interface QuotaStore {
+    getActivePolicy(userId: UserId): Promise<QuotaPolicy | null>;
+    getOrCreateBucket(userId: UserId, policy: QuotaPolicy, now: Date): Promise<QuotaBucket>;
+    updateBucket(userId: UserId, policyId: string, periodStart: Date, patch: Partial<Pick<QuotaBucket, "usedCredits" | "reservedCredits" | "usedRequests" | "reservedRequests">>): Promise<QuotaBucket>;
+}
+export interface UsageStore {
+    recordUsage(result: {
+        userId: UserId;
+        usage: UsageFacts;
+        estimatedCost: Credits;
+        actualCost?: Credits;
+    }): Promise<void>;
+}

package/dist/types.js

@@ -0,0 +1 @@
+export {};

package/dist/usageMeter.d.ts

@@ -0,0 +1,18 @@
+import { type PricingConfig } from "./pricing.js";
+import type { CaptureRequest, CaptureResult, EstimateRequest, EstimateResult, HoldRequest, HoldResult, MeteringPersistence, ReleaseRequest, ReleaseResult } from "./types.js";
+export interface UsageMeterOptions {
+    pricing: PricingConfig;
+}
+export declare class UsageMeter {
+    readonly pricing: PricingConfig;
+    readonly persistence: MeteringPersistence;
+    constructor(persistence: MeteringPersistence, options: UsageMeterOptions);
+    estimate(req: EstimateRequest): EstimateResult;
+    hold(req: HoldRequest): Promise<HoldResult>;
+    capture(req: CaptureRequest): Promise<CaptureResult>;
+    release(req: ReleaseRequest): Promise<ReleaseResult>;
+    private tryGetBucket;
+    private reserveQuota;
+    private unreserveQuota;
+    private settleQuota;
+}