@rudderjs/ai 1.4.0 → 1.6.0

package/dist/budget/storage.d.ts

@@ -0,0 +1,104 @@
+/**
+ * Persistence contract for `withBudget(...)` middleware (#A6 Phase 3).
+ *
+ * Apps configure a `BudgetStorage` instance — the middleware calls
+ * `checkAndDebit()` once before each request (estimated input cost) and
+ * once per `usage` chunk (actual output cost). The storage is responsible
+ * for atomic increment + cap enforcement.
+ *
+ * # Atomicity
+ *
+ * `checkAndDebit` MUST be atomic: a concurrent caller cannot observe a
+ * stale `spent` value between the read and the write. Otherwise two
+ * requests both pass the check before either debits, and a user can
+ * exceed their cap by `costUsd × concurrency`.
+ *
+ * Implementations:
+ * - `memoryBudgetStorage()` — single-process Map; atomic because we never
+ *   `await` between the read and the write.
+ * - `ormBudgetStorage()` (Phase 4) — uses `UPDATE … RETURNING` (Postgres)
+ *   or row-level `SELECT … FOR UPDATE` (MySQL/SQLite) to keep the read +
+ *   write in a single transaction.
+ * - Redis / Cache impls — would use `INCRBY` + a `EXPIRE`, treating the
+ *   counter as authoritative.
+ *
+ * # Period semantics
+ *
+ * `daily` rolls at the user's local midnight when `timezone` is provided
+ * (IANA name, e.g. `'America/New_York'`); UTC otherwise. `monthly` rolls
+ * on the calendar-month boundary in the same TZ. Phase 2 does not
+ * implement rolling-30-days windows — adds storage complexity for marginal
+ * value.
+ */
+export type BudgetPeriod = 'daily' | 'monthly';
+export interface BudgetCheckOptions {
+    /** Identifier for the budget owner (typically a user id). */
+    userId: string;
+    /** Which cap to check against. */
+    period: BudgetPeriod;
+    /** Cap in USD. Caller resolves `budget()` once per request. */
+    cap: number;
+    /** Cost to debit in USD. Pass `0` to read the current spent value without mutating. */
+    costUsd: number;
+    /** Override the clock — useful for tests + period-rollover assertions. */
+    now?: Date;
+    /** IANA timezone for period boundaries. Defaults to UTC. */
+    timezone?: string;
+}
+export interface BudgetCheckResult {
+    /** True when the debit was applied; false when it would have exceeded `cap`. */
+    allowed: boolean;
+    /**
+     * Current spend for the period.
+     *
+     * - When `allowed: true`, this is the value AFTER the debit (i.e.
+     *   includes `costUsd`).
+     * - When `allowed: false`, this is the value BEFORE the rejected debit
+     *   (i.e. the prior `spent`). This matches the value `BudgetExceededError`
+     *   should report — the user spent `spent`, the cap was `cap`, the
+     *   request would have pushed past.
+     */
+    spent: number;
+    /** The cap passed in `BudgetCheckOptions.cap`, echoed for caller convenience. */
+    cap: number;
+}
+export interface BudgetStorage {
+    /**
+     * Atomically read the current `spent` value, add `costUsd` if the result
+     * stays within `cap`, and return the outcome.
+     *
+     * Pass `costUsd: 0` for a pure read (e.g. for a "you've spent $X today"
+     * status display) — the result reflects current spend without mutating.
+     */
+    checkAndDebit(opts: BudgetCheckOptions): Promise<BudgetCheckResult>;
+    /**
+     * Reset the budget counter for `(userId, period)` at the period
+     * containing `now` (or "now" by clock). Useful for tests and admin
+     * overrides; not required by the middleware.
+     */
+    reset?(userId: string, period: BudgetPeriod, now?: Date, timezone?: string): Promise<void>;
+}
+/**
+ * Format a `Date` as the bucket key for a billing period.
+ *
+ * - `daily`   → `YYYY-MM-DD` in the given timezone (or UTC)
+ * - `monthly` → `YYYY-MM` in the given timezone (or UTC)
+ *
+ * Used by `memoryBudgetStorage` and `ormBudgetStorage` (Phase 4) to
+ * partition the counter by billing window. Exposed publicly so apps
+ * computing per-period totals outside the middleware (admin dashboards,
+ * tests) get the same bucketing.
+ */
+export declare function periodKey(period: BudgetPeriod, now: Date, timezone?: string): string;
+/**
+ * In-process `BudgetStorage` backed by a `Map`. Suitable for tests and
+ * single-process dev servers.
+ *
+ * **Cross-process caveat:** counters live in process-local memory, so
+ * queue workers (or any second process) see their own copy. For any app
+ * with workers or horizontal scaling, use `ormBudgetStorage()` (Phase 4)
+ * or a Redis-backed storage. Use `memoryBudgetStorage()` only when you're
+ * sure all budget-relevant code paths run in the same Node process.
+ */
+export declare function memoryBudgetStorage(): BudgetStorage;
+//# sourceMappingURL=storage.d.ts.map

package/dist/budget/storage.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"storage.d.ts","sourceRoot":"","sources":["../../src/budget/storage.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA+BG;AAEH,MAAM,MAAM,YAAY,GAAG,OAAO,GAAG,SAAS,CAAA;AAE9C,MAAM,WAAW,kBAAkB;IACjC,6DAA6D;IAC7D,MAAM,EAAK,MAAM,CAAA;IACjB,kCAAkC;IAClC,MAAM,EAAK,YAAY,CAAA;IACvB,+DAA+D;IAC/D,GAAG,EAAQ,MAAM,CAAA;IACjB,uFAAuF;IACvF,OAAO,EAAI,MAAM,CAAA;IACjB,0EAA0E;IAC1E,GAAG,CAAC,EAAO,IAAI,CAAA;IACf,4DAA4D;IAC5D,QAAQ,CAAC,EAAE,MAAM,CAAA;CAClB;AAED,MAAM,WAAW,iBAAiB;IAChC,gFAAgF;IAChF,OAAO,EAAE,OAAO,CAAA;IAChB;;;;;;;;;OASG;IACH,KAAK,EAAE,MAAM,CAAA;IACb,iFAAiF;IACjF,GAAG,EAAE,MAAM,CAAA;CACZ;AAED,MAAM,WAAW,aAAa;IAC5B;;;;;;OAMG;IACH,aAAa,CAAC,IAAI,EAAE,kBAAkB,GAAG,OAAO,CAAC,iBAAiB,CAAC,CAAA;IAEnE;;;;OAIG;IACH,KAAK,CAAC,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,YAAY,EAAE,GAAG,CAAC,EAAE,IAAI,EAAE,QAAQ,CAAC,EAAE,MAAM,GAAG,OAAO,CAAC,IAAI,CAAC,CAAA;CAC3F;AAED;;;;;;;;;;GAUG;AACH,wBAAgB,SAAS,CAAC,MAAM,EAAE,YAAY,EAAE,GAAG,EAAE,IAAI,EAAE,QAAQ,CAAC,EAAE,MAAM,GAAG,MAAM,CAUpF;AAED;;;;;;;;;GASG;AACH,wBAAgB,mBAAmB,IAAI,aAAa,CAgCnD"}

package/dist/budget/storage.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"storage.js","sourceRoot":"","sources":["../../src/budget/storage.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA+BG;AAuDH;;;;;;;;;;GAUG;AACH,MAAM,UAAU,SAAS,CAAC,MAAoB,EAAE,GAAS,EAAE,QAAiB;IAC1E,kEAAkE;IAClE,MAAM,GAAG,GAAG,IAAI,IAAI,CAAC,cAAc,CAAC,OAAO,EAAE;QAC3C,QAAQ,EAAE,QAAQ,IAAI,KAAK;QAC3B,IAAI,EAAG,SAAS;QAChB,KAAK,EAAE,SAAS;QAChB,GAAG,EAAI,SAAS;KACjB,CAAC,CAAA;IACF,MAAM,GAAG,GAAG,GAAG,CAAC,MAAM,CAAC,GAAG,CAAC,CAAA,CAAC,eAAe;IAC3C,OAAO,MAAM,KAAK,OAAO,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,GAAG,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAA,CAAC,YAAY;AAChE,CAAC;AAED;;;;;;;;;GASG;AACH,MAAM,UAAU,mBAAmB;IACjC,kEAAkE;IAClE,8DAA8D;IAC9D,MAAM,IAAI,GAAG,IAAI,GAAG,EAAkB,CAAA;IACtC,MAAM,GAAG,GAAG,CAAC,MAAc,EAAE,MAAoB,EAAE,GAAS,EAAE,EAAW,EAAU,EAAE,CACnF,GAAG,MAAM,IAAI,MAAM,IAAI,SAAS,CAAC,MAAM,EAAE,GAAG,EAAE,EAAE,CAAC,EAAE,CAAA;IAErD,OAAO;QACL,sEAAsE;QACtE,sEAAsE;QACtE,yEAAyE;QACzE,sEAAsE;QACtE,gEAAgE;QAChE,aAAa;QACb,KAAK,CAAC,aAAa,CAAC,EAAE,MAAM,EAAE,MAAM,EAAE,GAAG,EAAE,OAAO,EAAE,GAAG,GAAG,IAAI,IAAI,EAAE,EAAE,QAAQ,EAAE;YAC9E,IAAI,CAAC,MAAM,CAAC,QAAQ,CAAC,GAAG,CAAC,IAAI,GAAG,GAAG,CAAC;gBAAU,MAAM,IAAI,KAAK,CAAC,8EAA8E,GAAG,EAAE,CAAC,CAAA;YAClJ,IAAI,CAAC,MAAM,CAAC,QAAQ,CAAC,OAAO,CAAC,IAAI,OAAO,GAAG,CAAC;gBAAE,MAAM,IAAI,KAAK,CAAC,kFAAkF,OAAO,EAAE,CAAC,CAAA;YAE1J,MAAM,CAAC,GAAG,GAAG,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG,EAAE,QAAQ,CAAC,CAAA;YAC5C,MAAM,OAAO,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,CAAC,IAAI,CAAC,CAAA;YAChC,MAAM,SAAS,GAAG,OAAO,GAAG,OAAO,CAAA;YACnC,IAAI,SAAS,GAAG,GAAG,EAAE,CAAC;gBACpB,OAAO,EAAE,OAAO,EAAE,KAAK,EAAE,KAAK,EAAE,OAAO,EAAE,GAAG,EAAE,CAAA;YAChD,CAAC;YACD,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,SAAS,CAAC,CAAA;YACtB,OAAO,EAAE,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,SAAS,EAAE,GAAG,EAAE,CAAA;QACjD,CAAC;QAED,KAAK,CAAC,KAAK,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG,GAAG,IAAI,IAAI,EAAE,EAAE,QAAQ;YACpD,IAAI,CAAC,MAAM,CAAC,GAAG,CAAC,MAAM,EAAE,MAAM,EAAE,GAAG,EAAE,QAAQ,CAAC,CAAC,CAAA;QACjD,CAAC;KACF,CAAA;AACH,CAAC"}

package/dist/budget/with-budget.d.ts

@@ -0,0 +1,119 @@
+import type { AiMiddleware, MiddlewareContext } from '../types.js';
+import { type ModelPriceEntry } from './pricing.js';
+import type { BudgetPeriod, BudgetStorage } from './storage.js';
+/**
+ * Per-period caps in USD. Either or both can be set; periods with `null`
+ * / `undefined` are not enforced.
+ */
+export interface BudgetCaps {
+    daily?: number | null;
+    monthly?: number | null;
+}
+export interface BudgetExceededArgs {
+    userId: string;
+    period: BudgetPeriod;
+    /** Spend recorded BEFORE the rejected debit. */
+    spent: number;
+    cap: number;
+    ctx: MiddlewareContext;
+}
+export interface WithBudgetOptions {
+    /**
+     * Resolves the user identifier for a given request. Return `null` (or
+     * `undefined`) to bypass budget enforcement entirely — useful for
+     * unauthenticated paths or admin tooling.
+     *
+     * Async return supported.
+     */
+    user(ctx: MiddlewareContext): string | null | undefined | Promise<string | null | undefined>;
+    /**
+     * USD caps for the resolved user. Called once per request (per agent
+     * step, more precisely). Caps may be set per-tier / per-plan by reading
+     * the user from your DB inside the callback.
+     */
+    budget(args: {
+        userId: string;
+        ctx: MiddlewareContext;
+    }): BudgetCaps | Promise<BudgetCaps>;
+    /**
+     * Where counters persist. Use {@link memoryBudgetStorage} for tests +
+     * single-process dev; `ormBudgetStorage` (#A6 Phase 4) for production.
+     */
+    storage: BudgetStorage;
+    /**
+     * Pricing catalog. Defaults to the shipped {@link ModelPricing}; spread
+     * to override entries for negotiated rates:
+     *
+     * ```ts
+     * pricing: { ...ModelPricing, 'anthropic/claude-opus-4-7': { ... } }
+     * ```
+     */
+    pricing?: Record<string, ModelPriceEntry>;
+    /**
+     * Called when a debit would exceed a cap. Default throws
+     * {@link BudgetExceededError}; supply your own to log/alert before
+     * throwing, or to throw a different error class. Must throw — return
+     * value is ignored. The throw aborts the agent run before the model
+     * call.
+     */
+    onExceeded?(args: BudgetExceededArgs): never | Promise<never>;
+    /**
+     * IANA timezone for daily / monthly period rollover. Defaults to UTC.
+     * Use the user's tz to roll caps at user-local midnight (matches
+     * billing dashboards).
+     */
+    timezone?: string;
+    /**
+     * Approximate-tokens estimator used for the pre-debit. Defaults to
+     * `Math.ceil(text.length / 4)` — fine for English-heavy prompts. Pass
+     * a tiktoken-backed estimator for accuracy.
+     */
+    estimateTokens?: (text: string) => number;
+}
+/**
+ * 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 declare function withBudget(opts: WithBudgetOptions): AiMiddleware;
+//# sourceMappingURL=with-budget.d.ts.map

package/dist/budget/with-budget.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"with-budget.d.ts","sourceRoot":"","sources":["../../src/budget/with-budget.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAEV,YAAY,EAEZ,iBAAiB,EAClB,MAAM,aAAa,CAAA;AACpB,OAAO,EAIL,KAAK,eAAe,EACrB,MAAM,cAAc,CAAA;AACrB,OAAO,KAAK,EACV,YAAY,EACZ,aAAa,EACd,MAAM,cAAc,CAAA;AAErB;;;GAGG;AACH,MAAM,WAAW,UAAU;IACzB,KAAK,CAAC,EAAI,MAAM,GAAG,IAAI,CAAA;IACvB,OAAO,CAAC,EAAE,MAAM,GAAG,IAAI,CAAA;CACxB;AAED,MAAM,WAAW,kBAAkB;IACjC,MAAM,EAAE,MAAM,CAAA;IACd,MAAM,EAAE,YAAY,CAAA;IACpB,gDAAgD;IAChD,KAAK,EAAG,MAAM,CAAA;IACd,GAAG,EAAK,MAAM,CAAA;IACd,GAAG,EAAK,iBAAiB,CAAA;CAC1B;AAED,MAAM,WAAW,iBAAiB;IAChC;;;;;;OAMG;IACH,IAAI,CAAC,GAAG,EAAE,iBAAiB,GAAG,MAAM,GAAG,IAAI,GAAG,SAAS,GAAG,OAAO,CAAC,MAAM,GAAG,IAAI,GAAG,SAAS,CAAC,CAAA;IAE5F;;;;OAIG;IACH,MAAM,CAAC,IAAI,EAAE;QAAE,MAAM,EAAE,MAAM,CAAC;QAAC,GAAG,EAAE,iBAAiB,CAAA;KAAE,GAAG,UAAU,GAAG,OAAO,CAAC,UAAU,CAAC,CAAA;IAE1F;;;OAGG;IACH,OAAO,EAAE,aAAa,CAAA;IAEtB;;;;;;;OAOG;IACH,OAAO,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,eAAe,CAAC,CAAA;IAEzC;;;;;;OAMG;IACH,UAAU,CAAC,CAAC,IAAI,EAAE,kBAAkB,GAAG,KAAK,GAAG,OAAO,CAAC,KAAK,CAAC,CAAA;IAE7D;;;;OAIG;IACH,QAAQ,CAAC,EAAE,MAAM,CAAA;IAEjB;;;;OAIG;IACH,cAAc,CAAC,EAAE,CAAC,IAAI,EAAE,MAAM,KAAK,MAAM,CAAA;CAC1C;AAgBD;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA4CG;AACH,wBAAgB,UAAU,CAAC,IAAI,EAAE,iBAAiB,GAAG,YAAY,CAmGhE"}

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"}