@rudderjs/ai 1.5.0 → 1.6.1

package/dist/budget/with-budget.js

@@ -0,0 +1,175 @@
+import { BudgetExceededError, ModelPricing, assertKnownModelPricing, } from './pricing.js';
+const BudgetState = Symbol.for('rudderjs.ai.budget.state');
+/**
+ * Per-user spend cap middleware (#A6 Phase 3).
+ *
+ * Pre-debits an input-cost estimate before each provider call (refusing
+ * with {@link BudgetExceededError} if the user would exceed any
+ * configured cap), then debits the actual cost difference once the
+ * `usage` chunk arrives.
+ *
+ * The pre-debit reserves budget so two concurrent requests can't both
+ * pass the check before either is billed — the `BudgetStorage` contract
+ * (#A6 Phase 2) requires `checkAndDebit` to be atomic.
+ *
+ * # Example
+ *
+ * ```ts
+ * import { withBudget, memoryBudgetStorage, ModelPricing } from '@rudderjs/ai'
+ *
+ * const budgeted = withBudget({
+ *   user:    (ctx) => ctx.context as string,        // your app's user-id source
+ *   budget:  () => ({ daily: 0.50, monthly: 10 }),  // USD
+ *   storage: memoryBudgetStorage(),
+ *   pricing: ModelPricing,
+ * })
+ *
+ * class MyAgent extends Agent {
+ *   middleware() { return [budgeted] }
+ * }
+ * ```
+ *
+ * # Caveats
+ *
+ * - **Refunds on errors are not issued.** If the provider call fails
+ *   after the pre-debit, the estimate stays debited. This avoids the
+ *   complexity of distinguishing partial-credit cases (the model may
+ *   have produced output before erroring). Apps that need refund-on-error
+ *   should subscribe via `onError` and call `storage` directly.
+ * - **Cache token deltas not counted.** `TokenUsage` does not yet expose
+ *   `cacheReadInputTokens` / `cacheWriteInputTokens`; cached requests are
+ *   billed at the full `inputPer1k` rate today. Refining this is a phase
+ *   3.x follow-up that needs a `TokenUsage` widening.
+ * - **Tokenizer differences.** The default token estimator is
+ *   `text.length / 4`. Provider-reported `usage.promptTokens` may differ
+ *   by a few percent. Pass `estimateTokens: …` for a tiktoken-accurate
+ *   pre-debit if your caps are tight.
+ */
+export function withBudget(opts) {
+    const pricing = opts.pricing ?? ModelPricing;
+    const tz = opts.timezone;
+    const estimateTokens = opts.estimateTokens ?? defaultEstimateTokens;
+    const onExceeded = opts.onExceeded ?? defaultOnExceeded;
+    return {
+        name: 'budget',
+        async onIteration(ctx) {
+            // Pre-debit fires before each model call (every step), including
+            // step 1. `onIteration` runs after `onStart` but before
+            // `prepareStep`/`onConfig('beforeModel')` — so transforms applied by
+            // those later hooks (e.g. a `prepareStep` model swap) aren't
+            // reflected in this estimate. For v1 that's acceptable; tighten
+            // later if it bites.
+            const userId = await opts.user(ctx);
+            if (userId == null)
+                return; // bypass — unauthenticated path or admin
+            const modelKey = ctx.model;
+            // Fail loud if the model isn't priced — silently zero-costing through
+            // a typo'd model is the worst-of-both for budget enforcement.
+            const rate = assertKnownModelPricing(modelKey, pricing);
+            const caps = await opts.budget({ userId, ctx });
+            const definedCaps = [];
+            if (caps.daily != null)
+                definedCaps.push({ period: 'daily', cap: caps.daily });
+            if (caps.monthly != null)
+                definedCaps.push({ period: 'monthly', cap: caps.monthly });
+            if (definedCaps.length === 0)
+                return; // nothing to enforce
+            // Estimate input cost for THIS step from the live messages array.
+            const estimate = estimateInputCostUsd(ctx.messages, [], rate, estimateTokens);
+            // Pre-debit each defined period; throw on first denial.
+            for (const { period, cap } of definedCaps) {
+                const r = await opts.storage.checkAndDebit({
+                    userId,
+                    period,
+                    cap,
+                    costUsd: estimate,
+                    ...(tz != null ? { timezone: tz } : {}),
+                });
+                if (!r.allowed) {
+                    await onExceeded({ userId, period, spent: r.spent, cap: r.cap, ctx });
+                    // If onExceeded didn't throw, fall back to the default error so
+                    // the run is always aborted on a denied debit.
+                    throw new BudgetExceededError({ userId, period, spent: r.spent, cap: r.cap });
+                }
+            }
+            // Stash for onUsage true-up. If estimate was high vs actual we just
+            // accept the small over-charge; if low, onUsage debits the delta.
+            ;
+            ctx[BudgetState] = {
+                userId,
+                caps: definedCaps,
+                pendingEstimate: estimate,
+            };
+        },
+        async onUsage(ctx, usage) {
+            const state = ctx[BudgetState];
+            if (!state)
+                return; // no user, was bypassed
+            const modelKey = ctx.model;
+            const rate = pricing[modelKey];
+            // If pricing was found at onConfig time it's still found here; this is
+            // belt-and-suspenders for the rare case where the agent loop swapped
+            // models mid-run (failover). Skip silently — the pre-debit covered
+            // estimate, so we under-charge actual on a missing-rate model.
+            if (!rate) {
+                state.pendingEstimate = 0;
+                return;
+            }
+            const actualCost = (usage.promptTokens * rate.inputPer1k +
+                usage.completionTokens * rate.outputPer1k) / 1000;
+            const delta = actualCost - state.pendingEstimate;
+            state.pendingEstimate = 0; // clear so a tool round-trip's next onConfig starts fresh
+            if (delta <= 0)
+                return; // overestimated; accept the small over-charge
+            // Always-apply true-up — the response already streamed, we can't
+            // unspend. Pass MAX_SAFE_INTEGER as cap so the storage just records
+            // the delta. Pre-debit enforcement already happened at onConfig.
+            for (const { period } of state.caps) {
+                await opts.storage.checkAndDebit({
+                    userId: state.userId,
+                    period,
+                    cap: Number.MAX_SAFE_INTEGER,
+                    costUsd: delta,
+                    ...(tz != null ? { timezone: tz } : {}),
+                });
+            }
+        },
+    };
+}
+// ─── Helpers ──────────────────────────────────────────────
+function defaultEstimateTokens(text) {
+    return Math.ceil(text.length / 4);
+}
+function defaultOnExceeded(args) {
+    throw new BudgetExceededError({
+        userId: args.userId,
+        period: args.period,
+        spent: args.spent,
+        cap: args.cap,
+    });
+}
+function estimateInputCostUsd(messages, systemPrompts, rate, estimateTokens) {
+    // Concatenating all input into one string for the estimator means a
+    // single tokenizer pass for tiktoken-backed estimators (cheaper than
+    // tokenizing each message separately and summing).
+    const parts = [...systemPrompts];
+    for (const m of messages)
+        parts.push(messageText(m));
+    const tokens = estimateTokens(parts.join('\n'));
+    return (tokens * rate.inputPer1k) / 1000;
+}
+function messageText(m) {
+    if (typeof m.content === 'string')
+        return m.content;
+    if (Array.isArray(m.content)) {
+        let s = '';
+        for (const part of m.content)
+            s += contentPartText(part);
+        return s;
+    }
+    return '';
+}
+function contentPartText(p) {
+    return p.type === 'text' ? p.text : '';
+}
+//# sourceMappingURL=with-budget.js.map

package/dist/budget/with-budget.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"with-budget.js","sourceRoot":"","sources":["../../src/budget/with-budget.ts"],"names":[],"mappings":"AAMA,OAAO,EACL,mBAAmB,EACnB,YAAY,EACZ,uBAAuB,GAExB,MAAM,cAAc,CAAA;AAiFrB,MAAM,WAAW,GAAG,MAAM,CAAC,GAAG,CAAC,0BAA0B,CAAC,CAAA;AAc1D;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA4CG;AACH,MAAM,UAAU,UAAU,CAAC,IAAuB;IAChD,MAAM,OAAO,GAAS,IAAI,CAAC,OAAO,IAAW,YAAY,CAAA;IACzD,MAAM,EAAE,GAAc,IAAI,CAAC,QAAQ,CAAA;IACnC,MAAM,cAAc,GAAG,IAAI,CAAC,cAAc,IAAI,qBAAqB,CAAA;IACnE,MAAM,UAAU,GAAM,IAAI,CAAC,UAAU,IAAQ,iBAAiB,CAAA;IAE9D,OAAO;QACL,IAAI,EAAE,QAAQ;QAEd,KAAK,CAAC,WAAW,CAAC,GAAG;YACnB,iEAAiE;YACjE,wDAAwD;YACxD,qEAAqE;YACrE,6DAA6D;YAC7D,gEAAgE;YAChE,qBAAqB;YACrB,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,IAAI,CAAC,GAAG,CAAC,CAAA;YACnC,IAAI,MAAM,IAAI,IAAI;gBAAE,OAAM,CAAE,yCAAyC;YAErE,MAAM,QAAQ,GAAG,GAAG,CAAC,KAAK,CAAA;YAC1B,sEAAsE;YACtE,8DAA8D;YAC9D,MAAM,IAAI,GAAG,uBAAuB,CAAC,QAAQ,EAAE,OAAO,CAAC,CAAA;YAEvD,MAAM,IAAI,GAAG,MAAM,IAAI,CAAC,MAAM,CAAC,EAAE,MAAM,EAAE,GAAG,EAAE,CAAC,CAAA;YAC/C,MAAM,WAAW,GAAiD,EAAE,CAAA;YACpE,IAAI,IAAI,CAAC,KAAK,IAAM,IAAI;gBAAE,WAAW,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,OAAO,EAAI,GAAG,EAAE,IAAI,CAAC,KAAK,EAAE,CAAC,CAAA;YAClF,IAAI,IAAI,CAAC,OAAO,IAAI,IAAI;gBAAE,WAAW,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,SAAS,EAAE,GAAG,EAAE,IAAI,CAAC,OAAO,EAAE,CAAC,CAAA;YAEpF,IAAI,WAAW,CAAC,MAAM,KAAK,CAAC;gBAAE,OAAM,CAAE,qBAAqB;YAE3D,kEAAkE;YAClE,MAAM,QAAQ,GAAG,oBAAoB,CAAC,GAAG,CAAC,QAAQ,EAAE,EAAE,EAAE,IAAI,EAAE,cAAc,CAAC,CAAA;YAE7E,wDAAwD;YACxD,KAAK,MAAM,EAAE,MAAM,EAAE,GAAG,EAAE,IAAI,WAAW,EAAE,CAAC;gBAC1C,MAAM,CAAC,GAAG,MAAM,IAAI,CAAC,OAAO,CAAC,aAAa,CAAC;oBACzC,MAAM;oBACN,MAAM;oBACN,GAAG;oBACH,OAAO,EAAG,QAAQ;oBAClB,GAAG,CAAC,EAAE,IAAI,IAAI,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;iBACxC,CAAC,CAAA;gBACF,IAAI,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;oBACf,MAAM,UAAU,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,GAAG,EAAE,CAAC,CAAC,GAAG,EAAE,GAAG,EAAE,CAAC,CAAA;oBACrE,gEAAgE;oBAChE,+CAA+C;oBAC/C,MAAM,IAAI,mBAAmB,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,CAAC,CAAC,KAAK,EAAE,GAAG,EAAE,CAAC,CAAC,GAAG,EAAE,CAAC,CAAA;gBAC/E,CAAC;YACH,CAAC;YAED,oEAAoE;YACpE,kEAAkE;YAClE,CAAC;YAAC,GAA8B,CAAC,WAAW,CAAC,GAAG;gBAC9C,MAAM;gBACN,IAAI,EAAE,WAAW;gBACjB,eAAe,EAAE,QAAQ;aAC1B,CAAA;QACH,CAAC;QAED,KAAK,CAAC,OAAO,CAAC,GAAG,EAAE,KAAK;YACtB,MAAM,KAAK,GAAI,GAA8B,CAAC,WAAW,CAAC,CAAA;YAC1D,IAAI,CAAC,KAAK;gBAAE,OAAM,CAAE,wBAAwB;YAE5C,MAAM,QAAQ,GAAG,GAAG,CAAC,KAAK,CAAA;YAC1B,MAAM,IAAI,GAAG,OAAO,CAAC,QAAQ,CAAC,CAAA;YAC9B,uEAAuE;YACvE,qEAAqE;YACrE,mEAAmE;YACnE,+DAA+D;YAC/D,IAAI,CAAC,IAAI,EAAE,CAAC;gBACV,KAAK,CAAC,eAAe,GAAG,CAAC,CAAA;gBACzB,OAAM;YACR,CAAC;YAED,MAAM,UAAU,GAAG,CACjB,KAAK,CAAC,YAAY,GAAO,IAAI,CAAC,UAAU;gBACxC,KAAK,CAAC,gBAAgB,GAAG,IAAI,CAAC,WAAW,CAC1C,GAAG,IAAI,CAAA;YAER,MAAM,KAAK,GAAG,UAAU,GAAG,KAAK,CAAC,eAAe,CAAA;YAChD,KAAK,CAAC,eAAe,GAAG,CAAC,CAAA,CAAE,0DAA0D;YAErF,IAAI,KAAK,IAAI,CAAC;gBAAE,OAAM,CAAE,8CAA8C;YAEtE,iEAAiE;YACjE,oEAAoE;YACpE,iEAAiE;YACjE,KAAK,MAAM,EAAE,MAAM,EAAE,IAAI,KAAK,CAAC,IAAI,EAAE,CAAC;gBACpC,MAAM,IAAI,CAAC,OAAO,CAAC,aAAa,CAAC;oBAC/B,MAAM,EAAI,KAAK,CAAC,MAAM;oBACtB,MAAM;oBACN,GAAG,EAAO,MAAM,CAAC,gBAAgB;oBACjC,OAAO,EAAG,KAAK;oBACf,GAAG,CAAC,EAAE,IAAI,IAAI,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;iBACxC,CAAC,CAAA;YACJ,CAAC;QACH,CAAC;KACF,CAAA;AACH,CAAC;AAED,6DAA6D;AAE7D,SAAS,qBAAqB,CAAC,IAAY;IACzC,OAAO,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,MAAM,GAAG,CAAC,CAAC,CAAA;AACnC,CAAC;AAED,SAAS,iBAAiB,CAAC,IAAwB;IACjD,MAAM,IAAI,mBAAmB,CAAC;QAC5B,MAAM,EAAE,IAAI,CAAC,MAAM;QACnB,MAAM,EAAE,IAAI,CAAC,MAAM;QACnB,KAAK,EAAG,IAAI,CAAC,KAAK;QAClB,GAAG,EAAK,IAAI,CAAC,GAAG;KACjB,CAAC,CAAA;AACJ,CAAC;AAED,SAAS,oBAAoB,CAC3B,QAAqB,EACrB,aAAuB,EACvB,IAAqB,EACrB,cAAwC;IAExC,oEAAoE;IACpE,qEAAqE;IACrE,mDAAmD;IACnD,MAAM,KAAK,GAAa,CAAC,GAAG,aAAa,CAAC,CAAA;IAC1C,KAAK,MAAM,CAAC,IAAI,QAAQ;QAAE,KAAK,CAAC,IAAI,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC,CAAA;IACpD,MAAM,MAAM,GAAG,cAAc,CAAC,KAAK,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC,CAAA;IAC/C,OAAO,CAAC,MAAM,GAAG,IAAI,CAAC,UAAU,CAAC,GAAG,IAAI,CAAA;AAC1C,CAAC;AAED,SAAS,WAAW,CAAC,CAAY;IAC/B,IAAI,OAAO,CAAC,CAAC,OAAO,KAAK,QAAQ;QAAE,OAAO,CAAC,CAAC,OAAO,CAAA;IACnD,IAAI,KAAK,CAAC,OAAO,CAAC,CAAC,CAAC,OAAO,CAAC,EAAE,CAAC;QAC7B,IAAI,CAAC,GAAG,EAAE,CAAA;QACV,KAAK,MAAM,IAAI,IAAI,CAAC,CAAC,OAAO;YAAE,CAAC,IAAI,eAAe,CAAC,IAAI,CAAC,CAAA;QACxD,OAAO,CAAC,CAAA;IACV,CAAC;IACD,OAAO,EAAE,CAAA;AACX,CAAC;AAED,SAAS,eAAe,CAAC,CAAc;IACrC,OAAO,CAAC,CAAC,IAAI,KAAK,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,EAAE,CAAA;AACxC,CAAC"}

package/dist/budget-orm/index.d.ts

@@ -0,0 +1,96 @@
+/**
+ * `@rudderjs/ai/budget-orm` — ORM-backed {@link BudgetStorage} for #A6 Phase 4.
+ *
+ * Production-grade replacement for `memoryBudgetStorage()` (which is
+ * single-process only). Persists per-user spend counters in a
+ * `BudgetUsage` table via the registered `@rudderjs/orm` adapter — works
+ * across queue workers, web processes, and horizontally-scaled deployments.
+ *
+ * Wire it into your AI middleware:
+ *
+ * ```ts
+ * import { withBudget } from '@rudderjs/ai'
+ * import { ormBudgetStorage } from '@rudderjs/ai/budget-orm'
+ *
+ * const budgeted = withBudget({
+ *   user:    (ctx) => ctx.context as string,
+ *   budget:  () => ({ daily: 0.50, monthly: 10 }),
+ *   storage: ormBudgetStorage(),
+ * })
+ * ```
+ *
+ * The schema lives at {@link budgetUsagePrismaSchema} — copy it into your
+ * Prisma schema (or a new `prisma/schema/<file>.prisma` if you use the
+ * multi-file setup). The `@@unique([userId, period, periodKey])`
+ * constraint is the one load-bearing index — without it, the
+ * find-or-create path can race and produce duplicate rows.
+ *
+ * # Atomicity caveat
+ *
+ * `checkAndDebit` does a read-then-conditional-increment. The increment
+ * itself is atomic (`UPDATE col = col + n`), but the cap check sits
+ * between the read and the write. Under high concurrency for a single
+ * user (more than ~1 in-flight budgeted request at a time), total spend
+ * can briefly exceed `cap` by up to `costUsd × concurrency`. For typical
+ * apps this is a non-issue.
+ *
+ * Strict guarantees require a database transaction with serializable
+ * isolation or a Redis-backed counter — both planned as follow-ups. File
+ * an issue if you hit this in production.
+ */
+import { Model } from '@rudderjs/orm';
+import { type BudgetCheckOptions, type BudgetCheckResult, type BudgetPeriod, type BudgetStorage } from '../budget/storage.js';
+/**
+ * Model row backing {@link OrmBudgetStorage}. Exposed so apps that
+ * want admin views (e.g. "show me top spenders this month") can use
+ * `BudgetUsageRecord.where(...).get()` instead of routing every read
+ * through the {@link BudgetStorage} interface.
+ *
+ * The `@@unique([userId, period, periodKey])` constraint is required —
+ * without it, two concurrent first-writes for the same user/period
+ * create duplicate rows and the cap accounting silently drifts.
+ */
+export declare class BudgetUsageRecord extends Model {
+    static table: string;
+    static fillable: string[];
+    id: string;
+    userId: string;
+    /** `'daily'` or `'monthly'`. */
+    period: string;
+    /** TZ-aware bucket key — `YYYY-MM-DD` (daily) or `YYYY-MM` (monthly). */
+    periodKey: string;
+    /** Cumulative USD spend in this period. */
+    spent: number;
+    createdAt: Date;
+    updatedAt: Date | null;
+}
+/**
+ * Production `BudgetStorage` backed by the registered `@rudderjs/orm`
+ * adapter. See the module JSDoc for setup + the atomicity caveat.
+ */
+export declare class OrmBudgetStorage implements BudgetStorage {
+    checkAndDebit(opts: BudgetCheckOptions): Promise<BudgetCheckResult>;
+    /** Apply the read-then-conditional-increment path on an existing row. */
+    private _applyIncrementPath;
+    reset(userId: string, period: BudgetPeriod, now?: Date, timezone?: string): Promise<void>;
+}
+/**
+ * Convenience factory — returns a fresh {@link OrmBudgetStorage}
+ * instance. Prefer this over `new OrmBudgetStorage()` for symmetry with
+ * `memoryBudgetStorage()`.
+ */
+export declare function ormBudgetStorage(): BudgetStorage;
+/**
+ * Reference Prisma schema for `OrmBudgetStorage`. Copy into your
+ * `prisma/schema/<file>.prisma` (or paste alongside an existing model).
+ *
+ * The `@@unique([userId, period, periodKey])` constraint is required —
+ * without it the find-or-create path can race and produce duplicate
+ * rows, breaking cap accounting.
+ *
+ * SQLite stores `Float` as `REAL`; Postgres / MySQL as `DOUBLE
+ * PRECISION` / `DOUBLE`. All three give 15+ significant digits — more
+ * than enough for sub-cent budget tracking.
+ */
+export declare const budgetUsagePrismaSchema = "model BudgetUsage {\n  id        String   @id @default(cuid())\n  userId    String\n  /// 'daily' | 'monthly'\n  period    String\n  /// YYYY-MM-DD (daily) or YYYY-MM (monthly), in the configured timezone\n  periodKey String\n  /// Cumulative USD spend in this period\n  spent     Float    @default(0)\n  createdAt DateTime @default(now())\n  updatedAt DateTime @updatedAt\n\n  @@unique([userId, period, periodKey])\n  @@index([userId])\n}\n";
+//# sourceMappingURL=index.d.ts.map

package/dist/budget-orm/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/budget-orm/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAuCG;AAEH,OAAO,EAAE,KAAK,EAAE,MAAM,eAAe,CAAA;AACrC,OAAO,EACL,KAAK,kBAAkB,EACvB,KAAK,iBAAiB,EACtB,KAAK,YAAY,EACjB,KAAK,aAAa,EAEnB,MAAM,sBAAsB,CAAA;AAI7B;;;;;;;;;GASG;AACH,qBAAa,iBAAkB,SAAQ,KAAK;IAC1C,OAAgB,KAAK,SAAmB;IACxC,OAAgB,QAAQ,WAA6C;IAE7D,EAAE,EAAS,MAAM,CAAA;IACjB,MAAM,EAAK,MAAM,CAAA;IACzB,gCAAgC;IACxB,MAAM,EAAK,MAAM,CAAA;IACzB,yEAAyE;IACjE,SAAS,EAAE,MAAM,CAAA;IACzB,2CAA2C;IACnC,KAAK,EAAM,MAAM,CAAA;IACjB,SAAS,EAAE,IAAI,CAAA;IACf,SAAS,EAAE,IAAI,GAAG,IAAI,CAAA;CAC/B;AAID;;;GAGG;AACH,qBAAa,gBAAiB,YAAW,aAAa;IAC9C,aAAa,CAAC,IAAI,EAAE,kBAAkB,GAAG,OAAO,CAAC,iBAAiB,CAAC;IAuDzE,yEAAyE;YAC3D,mBAAmB;IAsB3B,KAAK,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,YAAY,EAAE,GAAG,CAAC,EAAE,IAAI,EAAE,QAAQ,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC;CAQhG;AAED;;;;GAIG;AACH,wBAAgB,gBAAgB,IAAI,aAAa,CAEhD;AAID;;;;;;;;;;;GAWG;AACH,eAAO,MAAM,uBAAuB,8bAenC,CAAA"}

package/dist/budget-orm/index.js

@@ -0,0 +1,177 @@
+/**
+ * `@rudderjs/ai/budget-orm` — ORM-backed {@link BudgetStorage} for #A6 Phase 4.
+ *
+ * Production-grade replacement for `memoryBudgetStorage()` (which is
+ * single-process only). Persists per-user spend counters in a
+ * `BudgetUsage` table via the registered `@rudderjs/orm` adapter — works
+ * across queue workers, web processes, and horizontally-scaled deployments.
+ *
+ * Wire it into your AI middleware:
+ *
+ * ```ts
+ * import { withBudget } from '@rudderjs/ai'
+ * import { ormBudgetStorage } from '@rudderjs/ai/budget-orm'
+ *
+ * const budgeted = withBudget({
+ *   user:    (ctx) => ctx.context as string,
+ *   budget:  () => ({ daily: 0.50, monthly: 10 }),
+ *   storage: ormBudgetStorage(),
+ * })
+ * ```
+ *
+ * The schema lives at {@link budgetUsagePrismaSchema} — copy it into your
+ * Prisma schema (or a new `prisma/schema/<file>.prisma` if you use the
+ * multi-file setup). The `@@unique([userId, period, periodKey])`
+ * constraint is the one load-bearing index — without it, the
+ * find-or-create path can race and produce duplicate rows.
+ *
+ * # Atomicity caveat
+ *
+ * `checkAndDebit` does a read-then-conditional-increment. The increment
+ * itself is atomic (`UPDATE col = col + n`), but the cap check sits
+ * between the read and the write. Under high concurrency for a single
+ * user (more than ~1 in-flight budgeted request at a time), total spend
+ * can briefly exceed `cap` by up to `costUsd × concurrency`. For typical
+ * apps this is a non-issue.
+ *
+ * Strict guarantees require a database transaction with serializable
+ * isolation or a Redis-backed counter — both planned as follow-ups. File
+ * an issue if you hit this in production.
+ */
+import { Model } from '@rudderjs/orm';
+import { periodKey as buildPeriodKey, } from '../budget/storage.js';
+// ─── ORM Model ────────────────────────────────────────────
+/**
+ * Model row backing {@link OrmBudgetStorage}. Exposed so apps that
+ * want admin views (e.g. "show me top spenders this month") can use
+ * `BudgetUsageRecord.where(...).get()` instead of routing every read
+ * through the {@link BudgetStorage} interface.
+ *
+ * The `@@unique([userId, period, periodKey])` constraint is required —
+ * without it, two concurrent first-writes for the same user/period
+ * create duplicate rows and the cap accounting silently drifts.
+ */
+export class BudgetUsageRecord extends Model {
+    static table = 'budgetUsage';
+    static fillable = ['userId', 'period', 'periodKey', 'spent'];
+}
+// ─── BudgetStorage adapter ────────────────────────────────
+/**
+ * Production `BudgetStorage` backed by the registered `@rudderjs/orm`
+ * adapter. See the module JSDoc for setup + the atomicity caveat.
+ */
+export class OrmBudgetStorage {
+    async checkAndDebit(opts) {
+        if (!Number.isFinite(opts.cap) || opts.cap < 0) {
+            throw new Error(`[RudderJS AI] BudgetStorage: cap must be a non-negative finite number, got ${opts.cap}`);
+        }
+        if (!Number.isFinite(opts.costUsd) || opts.costUsd < 0) {
+            throw new Error(`[RudderJS AI] BudgetStorage: costUsd must be a non-negative finite number, got ${opts.costUsd}`);
+        }
+        const now = opts.now ?? new Date();
+        const key = buildPeriodKey(opts.period, now, opts.timezone);
+        const existing = await BudgetUsageRecord
+            .where('userId', opts.userId)
+            .where('period', opts.period)
+            .where('periodKey', key)
+            .first();
+        // ─── No row yet — first write for this period ─────────
+        if (!existing) {
+            // Pure-read on an empty bucket — still empty after.
+            if (opts.costUsd === 0) {
+                return { allowed: true, spent: 0, cap: opts.cap };
+            }
+            // Single debit larger than cap — refuse before creating the row,
+            // so we don't pollute storage with denied requests.
+            if (opts.costUsd > opts.cap) {
+                return { allowed: false, spent: 0, cap: opts.cap };
+            }
+            try {
+                await BudgetUsageRecord.create({
+                    userId: opts.userId,
+                    period: opts.period,
+                    periodKey: key,
+                    spent: opts.costUsd,
+                });
+                return { allowed: true, spent: opts.costUsd, cap: opts.cap };
+            }
+            catch (e) {
+                // Race: another caller created the row between our `first()` and
+                // `create()`. Re-read and fall through to the increment path.
+                // We deliberately don't sniff the error type — any create failure
+                // means the row may now exist; let the re-read decide.
+                const refetched = await BudgetUsageRecord
+                    .where('userId', opts.userId)
+                    .where('period', opts.period)
+                    .where('periodKey', key)
+                    .first();
+                if (!refetched)
+                    throw e; // not a unique-constraint race; surface the original error
+                return this._applyIncrementPath(refetched, opts);
+            }
+        }
+        return this._applyIncrementPath(existing, opts);
+    }
+    /** Apply the read-then-conditional-increment path on an existing row. */
+    async _applyIncrementPath(row, opts) {
+        const current = Number(row.spent ?? 0);
+        // Pure read.
+        if (opts.costUsd === 0) {
+            return { allowed: true, spent: current, cap: opts.cap };
+        }
+        // Cap check — read-then-decide. Atomic under single-writer; under
+        // concurrent writers, see the module-level atomicity caveat.
+        if (current + opts.costUsd > opts.cap) {
+            return { allowed: false, spent: current, cap: opts.cap };
+        }
+        const updated = await BudgetUsageRecord.increment(row.id, 'spent', opts.costUsd);
+        const newSpent = Number(updated?.spent ?? current + opts.costUsd);
+        return { allowed: true, spent: newSpent, cap: opts.cap };
+    }
+    async reset(userId, period, now, timezone) {
+        const key = buildPeriodKey(period, now ?? new Date(), timezone);
+        await BudgetUsageRecord
+            .where('userId', userId)
+            .where('period', period)
+            .where('periodKey', key)
+            .deleteAll();
+    }
+}
+/**
+ * Convenience factory — returns a fresh {@link OrmBudgetStorage}
+ * instance. Prefer this over `new OrmBudgetStorage()` for symmetry with
+ * `memoryBudgetStorage()`.
+ */
+export function ormBudgetStorage() {
+    return new OrmBudgetStorage();
+}
+// ─── Schema reference ─────────────────────────────────────
+/**
+ * Reference Prisma schema for `OrmBudgetStorage`. Copy into your
+ * `prisma/schema/<file>.prisma` (or paste alongside an existing model).
+ *
+ * The `@@unique([userId, period, periodKey])` constraint is required —
+ * without it the find-or-create path can race and produce duplicate
+ * rows, breaking cap accounting.
+ *
+ * SQLite stores `Float` as `REAL`; Postgres / MySQL as `DOUBLE
+ * PRECISION` / `DOUBLE`. All three give 15+ significant digits — more
+ * than enough for sub-cent budget tracking.
+ */
+export const budgetUsagePrismaSchema = `model BudgetUsage {
+  id        String   @id @default(cuid())
+  userId    String
+  /// 'daily' | 'monthly'
+  period    String
+  /// YYYY-MM-DD (daily) or YYYY-MM (monthly), in the configured timezone
+  periodKey String
+  /// Cumulative USD spend in this period
+  spent     Float    @default(0)
+  createdAt DateTime @default(now())
+  updatedAt DateTime @updatedAt
+
+  @@unique([userId, period, periodKey])
+  @@index([userId])
+}
+`;
+//# sourceMappingURL=index.js.map

package/dist/budget-orm/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../src/budget-orm/index.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAuCG;AAEH,OAAO,EAAE,KAAK,EAAE,MAAM,eAAe,CAAA;AACrC,OAAO,EAKL,SAAS,IAAI,cAAc,GAC5B,MAAM,sBAAsB,CAAA;AAE7B,6DAA6D;AAE7D;;;;;;;;;GASG;AACH,MAAM,OAAO,iBAAkB,SAAQ,KAAK;IAC1C,MAAM,CAAU,KAAK,GAAM,aAAa,CAAA;IACxC,MAAM,CAAU,QAAQ,GAAG,CAAC,QAAQ,EAAE,QAAQ,EAAE,WAAW,EAAE,OAAO,CAAC,CAAA;;AAcvE,6DAA6D;AAE7D;;;GAGG;AACH,MAAM,OAAO,gBAAgB;IAC3B,KAAK,CAAC,aAAa,CAAC,IAAwB;QAC1C,IAAI,CAAC,MAAM,CAAC,QAAQ,CAAC,IAAI,CAAC,GAAG,CAAC,IAAI,IAAI,CAAC,GAAG,GAAG,CAAC,EAAE,CAAC;YAC/C,MAAM,IAAI,KAAK,CAAC,8EAA8E,IAAI,CAAC,GAAG,EAAE,CAAC,CAAA;QAC3G,CAAC;QACD,IAAI,CAAC,MAAM,CAAC,QAAQ,CAAC,IAAI,CAAC,OAAO,CAAC,IAAI,IAAI,CAAC,OAAO,GAAG,CAAC,EAAE,CAAC;YACvD,MAAM,IAAI,KAAK,CAAC,kFAAkF,IAAI,CAAC,OAAO,EAAE,CAAC,CAAA;QACnH,CAAC;QAED,MAAM,GAAG,GAAG,IAAI,CAAC,GAAG,IAAI,IAAI,IAAI,EAAE,CAAA;QAClC,MAAM,GAAG,GAAG,cAAc,CAAC,IAAI,CAAC,MAAM,EAAE,GAAG,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAA;QAE3D,MAAM,QAAQ,GAAG,MAAM,iBAAiB;aACrC,KAAK,CAAC,QAAQ,EAAK,IAAI,CAAC,MAAM,CAAC;aAC/B,KAAK,CAAC,QAAQ,EAAK,IAAI,CAAC,MAAM,CAAC;aAC/B,KAAK,CAAC,WAAW,EAAE,GAAG,CAAC;aACvB,KAAK,EAAyC,CAAA;QAEjD,yDAAyD;QACzD,IAAI,CAAC,QAAQ,EAAE,CAAC;YACd,oDAAoD;YACpD,IAAI,IAAI,CAAC,OAAO,KAAK,CAAC,EAAE,CAAC;gBACvB,OAAO,EAAE,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,CAAC,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,CAAA;YACnD,CAAC;YACD,iEAAiE;YACjE,oDAAoD;YACpD,IAAI,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,GAAG,EAAE,CAAC;gBAC5B,OAAO,EAAE,OAAO,EAAE,KAAK,EAAE,KAAK,EAAE,CAAC,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,CAAA;YACpD,CAAC;YAED,IAAI,CAAC;gBACH,MAAM,iBAAiB,CAAC,MAAM,CAAC;oBAC7B,MAAM,EAAK,IAAI,CAAC,MAAM;oBACtB,MAAM,EAAK,IAAI,CAAC,MAAM;oBACtB,SAAS,EAAE,GAAG;oBACd,KAAK,EAAM,IAAI,CAAC,OAAO;iBACxB,CAAC,CAAA;gBACF,OAAO,EAAE,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,IAAI,CAAC,OAAO,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,CAAA;YAC9D,CAAC;YAAC,OAAO,CAAC,EAAE,CAAC;gBACX,iEAAiE;gBACjE,8DAA8D;gBAC9D,kEAAkE;gBAClE,uDAAuD;gBACvD,MAAM,SAAS,GAAG,MAAM,iBAAiB;qBACtC,KAAK,CAAC,QAAQ,EAAK,IAAI,CAAC,MAAM,CAAC;qBAC/B,KAAK,CAAC,QAAQ,EAAK,IAAI,CAAC,MAAM,CAAC;qBAC/B,KAAK,CAAC,WAAW,EAAE,GAAG,CAAC;qBACvB,KAAK,EAAyC,CAAA;gBACjD,IAAI,CAAC,SAAS;oBAAE,MAAM,CAAC,CAAA,CAAE,2DAA2D;gBACpF,OAAO,IAAI,CAAC,mBAAmB,CAAC,SAAS,EAAE,IAAI,CAAC,CAAA;YAClD,CAAC;QACH,CAAC;QAED,OAAO,IAAI,CAAC,mBAAmB,CAAC,QAAQ,EAAE,IAAI,CAAC,CAAA;IACjD,CAAC;IAED,yEAAyE;IACjE,KAAK,CAAC,mBAAmB,CAC/B,GAAuB,EACvB,IAAwB;QAExB,MAAM,OAAO,GAAG,MAAM,CAAC,GAAG,CAAC,KAAK,IAAI,CAAC,CAAC,CAAA;QAEtC,aAAa;QACb,IAAI,IAAI,CAAC,OAAO,KAAK,CAAC,EAAE,CAAC;YACvB,OAAO,EAAE,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,OAAO,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,CAAA;QACzD,CAAC;QAED,kEAAkE;QAClE,6DAA6D;QAC7D,IAAI,OAAO,GAAG,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,GAAG,EAAE,CAAC;YACtC,OAAO,EAAE,OAAO,EAAE,KAAK,EAAE,KAAK,EAAE,OAAO,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,CAAA;QAC1D,CAAC;QAED,MAAM,OAAO,GAAG,MAAM,iBAAiB,CAAC,SAAS,CAAC,GAAG,CAAC,EAAE,EAAE,OAAO,EAAE,IAAI,CAAC,OAAO,CAAiC,CAAA;QAChH,MAAM,QAAQ,GAAG,MAAM,CAAC,OAAO,EAAE,KAAK,IAAI,OAAO,GAAG,IAAI,CAAC,OAAO,CAAC,CAAA;QACjE,OAAO,EAAE,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,QAAQ,EAAE,GAAG,EAAE,IAAI,CAAC,GAAG,EAAE,CAAA;IAC1D,CAAC;IAED,KAAK,CAAC,KAAK,CAAC,MAAc,EAAE,MAAoB,EAAE,GAAU,EAAE,QAAiB;QAC7E,MAAM,GAAG,GAAG,cAAc,CAAC,MAAM,EAAE,GAAG,IAAI,IAAI,IAAI,EAAE,EAAE,QAAQ,CAAC,CAAA;QAC/D,MAAM,iBAAiB;aACpB,KAAK,CAAC,QAAQ,EAAK,MAAM,CAAC;aAC1B,KAAK,CAAC,QAAQ,EAAK,MAAM,CAAC;aAC1B,KAAK,CAAC,WAAW,EAAE,GAAG,CAAC;aACvB,SAAS,EAAE,CAAA;IAChB,CAAC;CACF;AAED;;;;GAIG;AACH,MAAM,UAAU,gBAAgB;IAC9B,OAAO,IAAI,gBAAgB,EAAE,CAAA;AAC/B,CAAC;AAED,6DAA6D;AAE7D;;;;;;;;;;;GAWG;AACH,MAAM,CAAC,MAAM,uBAAuB,GAAG;;;;;;;;;;;;;;;CAetC,CAAA"}

package/dist/commands/ai-eval.d.ts

@@ -0,0 +1,93 @@
+/**
+ * `pnpm rudder ai:eval` — discover `evals/**\/*.eval.ts` suites,
+ * run each, and report. Console reporter by default; `--json` emits
+ * a machine-readable envelope to stdout for CI.
+ *
+ * Registered from the CLI loader (`packages/cli/src/index.ts`)
+ * — the AiProvider doesn't own this so it surfaces even when the
+ * user app fails to boot, matching the `command:list --json`
+ * graceful-degradation pattern from #349.
+ */
+import type { EvalSuite } from '../eval/index.js';
+type Rudder = {
+    command(name: string, handler: (args: string[]) => void | Promise<void>): {
+        description(text: string): unknown;
+    };
+};
+/** CLI flags + positional name filter. */
+export interface AiEvalOptions {
+    /** Substring filter (case-insensitive) applied to suite names. */
+    filter?: string;
+    /** Stop on the first failing suite. */
+    bail: boolean;
+    /** Emit `{ suites: [...] }` JSON to stdout. */
+    json: boolean;
+    /**
+     * Run against the real provider, capture each case's assistant
+     * turns to `evals/__fixtures__/<suite>/<case>.json`. Existing
+     * fixtures are overwritten — diff in your VCS to see what changed.
+     * Default `false`.
+     */
+    record?: boolean;
+    /**
+     * Swap the runtime with `AiFake.fake()` and feed each case its
+     * recorded fixture via `respondWithSequence`. Zero API calls,
+     * deterministic regression tests. Cases without a fixture fall
+     * through to a normal run with a stderr warning. Default `false`.
+     */
+    replay?: boolean;
+    /**
+     * Path for a self-contained HTML report (#A5 Phase 5). Pasteable
+     * into PR comments / Slack threads. Coexists with `--json` (JSON
+     * still goes to stdout, HTML goes to disk).
+     */
+    html?: string;
+}
+/**
+ * Test seam — every external dependency gets an injectable
+ * override. The CLI handler defaults each to its real impl.
+ */
+export interface AiEvalDeps {
+    cwd?: string;
+    stdout?: {
+        write(s: string): boolean | void;
+    };
+    stderr?: {
+        write(s: string): boolean | void;
+    };
+    /** Override the file walk (test harness returns a virtual list). */
+    discover?: (cwd: string, pattern: string) => Promise<string[]>;
+    /** Override file → suite loader (test harness uses an in-memory map). */
+    loadSuite?: (absPath: string) => Promise<EvalSuite | null>;
+    /** Override config lookup (test harness skips `@rudderjs/core`). */
+    configPattern?: () => string | null | Promise<string | null>;
+    /**
+     * Override fixtures directory (defaults to `<cwd>/evals/__fixtures__`).
+     * Tests point to a tmpdir to keep round-trips off the source tree.
+     */
+    fixturesDir?: string;
+}
+/** Register the `ai:eval` command on the rudder runner. */
+export declare function registerAiEvalCommand(rudder: Rudder): void;
+/**
+ * Parse the rest-of-line. Recognizes:
+ *  - boolean flags: `--bail`, `--json`, `--record`, `--replay`
+ *  - value flags  : `--html <path>` or `--html=<path>`
+ *  - one positional name filter (anything not consumed above)
+ */
+export declare function parseArgs(args: string[]): AiEvalOptions;
+/**
+ * Execute the CLI flow. Returns the process exit code (0 = all pass,
+ * 1 = at least one suite had a failure or no suites discovered).
+ *
+ * The handler is `process.exit`-free so tests can drive it directly.
+ */
+export declare function runEvalCli(opts: AiEvalOptions, deps?: AiEvalDeps): Promise<number>;
+/**
+ * Recursive walk constrained to a `<dir>/**\/*<suffix>` shape.
+ * Returns absolute paths sorted lexicographically for stable test
+ * output and predictable `--bail` ordering.
+ */
+export declare function discoverSuiteFiles(cwd: string, pattern: string): Promise<string[]>;
+export {};
+//# sourceMappingURL=ai-eval.d.ts.map

package/dist/commands/ai-eval.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"ai-eval.d.ts","sourceRoot":"","sources":["../../src/commands/ai-eval.ts"],"names":[],"mappings":"AAAA;;;;;;;;;GASG;AAMH,OAAO,KAAK,EAAE,SAAS,EAAiC,MAAM,kBAAkB,CAAA;AAUhF,KAAK,MAAM,GAAG;IACZ,OAAO,CACL,IAAI,EAAE,MAAM,EACZ,OAAO,EAAE,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,GAChD;QAAE,WAAW,CAAC,IAAI,EAAE,MAAM,GAAG,OAAO,CAAA;KAAE,CAAA;CAC1C,CAAA;AAED,0CAA0C;AAC1C,MAAM,WAAW,aAAa;IAC5B,kEAAkE;IAClE,MAAM,CAAC,EAAE,MAAM,CAAA;IACf,uCAAuC;IACvC,IAAI,EAAK,OAAO,CAAA;IAChB,+CAA+C;IAC/C,IAAI,EAAK,OAAO,CAAA;IAChB;;;;;OAKG;IACH,MAAM,CAAC,EAAE,OAAO,CAAA;IAChB;;;;;OAKG;IACH,MAAM,CAAC,EAAE,OAAO,CAAA;IAChB;;;;OAIG;IACH,IAAI,CAAC,EAAE,MAAM,CAAA;CACd;AAED;;;GAGG;AACH,MAAM,WAAW,UAAU;IACzB,GAAG,CAAC,EAAS,MAAM,CAAA;IACnB,MAAM,CAAC,EAAM;QAAE,KAAK,CAAC,CAAC,EAAE,MAAM,GAAG,OAAO,GAAG,IAAI,CAAA;KAAE,CAAA;IACjD,MAAM,CAAC,EAAM;QAAE,KAAK,CAAC,CAAC,EAAE,MAAM,GAAG,OAAO,GAAG,IAAI,CAAA;KAAE,CAAA;IACjD,oEAAoE;IACpE,QAAQ,CAAC,EAAI,CAAC,GAAG,EAAE,MAAM,EAAE,OAAO,EAAE,MAAM,KAAK,OAAO,CAAC,MAAM,EAAE,CAAC,CAAA;IAChE,yEAAyE;IACzE,SAAS,CAAC,EAAG,CAAC,OAAO,EAAE,MAAM,KAAK,OAAO,CAAC,SAAS,GAAG,IAAI,CAAC,CAAA;IAC3D,oEAAoE;IACpE,aAAa,CAAC,EAAE,MAAM,MAAM,GAAG,IAAI,GAAG,OAAO,CAAC,MAAM,GAAG,IAAI,CAAC,CAAA;IAC5D;;;OAGG;IACH,WAAW,CAAC,EAAE,MAAM,CAAA;CACrB;AAED,2DAA2D;AAC3D,wBAAgB,qBAAqB,CAAC,MAAM,EAAE,MAAM,GAAG,IAAI,CAO1D;AAMD;;;;;GAKG;AACH,wBAAgB,SAAS,CAAC,IAAI,EAAE,MAAM,EAAE,GAAG,aAAa,CA8BvD;AAID;;;;;GAKG;AACH,wBAAsB,UAAU,CAAC,IAAI,EAAE,aAAa,EAAE,IAAI,GAAE,UAAe,GAAG,OAAO,CAAC,MAAM,CAAC,CAsE5F;AA4JD;;;;GAIG;AACH,wBAAsB,kBAAkB,CAAC,GAAG,EAAE,MAAM,EAAE,OAAO,EAAE,MAAM,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC,CAMxF"}