@convex-dev/ai-budget 0.0.2-alpha.3 → 0.0.2-alpha.5

package/dist/client/index.js

@@ -39,14 +39,30 @@ function extractUsage(usage) {
     return {
         promptTokens: toTokenCount(usage?.inputTokens ?? usage?.promptTokens),
         completionTokens: toTokenCount(usage?.outputTokens ?? usage?.completionTokens),
-        // cached prompt tokens: AI SDK v5 `cachedInputTokens`, OpenAI-compat
-        // `prompt_tokens_details.cached_tokens` / `cached_tokens`.
-        cachedTokens: toTokenCount(usage?.cachedInputTokens ??
+        // cached prompt tokens. The Convex gateway reports these at
+        // `usage.inputTokenDetails.cacheReadTokens`; the other paths cover AI SDK v5
+        // (`cachedInputTokens`) and raw OpenAI-compatible shapes.
+        cachedTokens: toTokenCount(usage?.inputTokenDetails?.cacheReadTokens ??
+            usage?.cachedInputTokens ??
             usage?.promptTokensDetails?.cachedTokens ??
             usage?.prompt_tokens_details?.cached_tokens ??
             usage?.cached_tokens),
     };
 }
+// The gateway doesn't report a dollar cost today, but if a future response ever
+// carries an unambiguous nanodollar cost we pass it straight through as
+// authoritative (finishRequest prefers it over the token-based estimate). Only
+// an explicitly nano-denominated field is trusted — a bare `cost` could be in
+// dollars and silently mis-bill by 1e9×.
+function extractGatewayCostNanos(result) {
+    const meta = result?.providerMetadata?.convexGateway;
+    const candidates = [meta?.costNanos, result?.usage?.costNanos];
+    for (const c of candidates) {
+        if (typeof c === "number" && Number.isFinite(c) && c >= 0)
+            return c;
+    }
+    return undefined;
+}
 // Flatten an AI SDK prompt (roles + content parts) into simple storable messages.
 function simplifyPrompt(prompt) {
     if (!Array.isArray(prompt))
@@ -78,24 +94,46 @@ function extractText(result) {
     }
     return "";
 }
-// ---------- client ----------
 export class AIBudget {
     component;
     defaultModel;
     onSoftLimit;
+    onThreshold;
+    onLimitReached;
     constructor(component, options) {
         this.component = component;
         this.defaultModel = options?.defaultModel ?? "openai/gpt-4o-mini";
         this.onSoftLimit = options?.onSoftLimit;
+        this.onThreshold = options?.onThreshold;
+        this.onLimitReached = options?.onLimitReached;
+    }
+    // Fire the soft-limit + threshold callbacks from a startRequest result.
+    async fireBudgetEvents(base, warnings, notices) {
+        if (warnings.length > 0 && this.onSoftLimit) {
+            try {
+                await this.onSoftLimit({ ...base, messages: warnings, warnings });
+            }
+            catch {
+                /* never let a callback error break a request */
+            }
+        }
+        if (notices.length > 0 && this.onThreshold) {
+            try {
+                await this.onThreshold({ ...base, messages: notices });
+            }
+            catch {
+                /* swallow */
+            }
+        }
     }
-    async fireSoftLimit(info) {
-        if (info.warnings.length === 0 || !this.onSoftLimit)
+    async fireLimitReached(info) {
+        if (!this.onLimitReached)
             return;
         try {
-            await this.onSoftLimit(info);
+            await this.onLimitReached(info);
         }
         catch {
-            // never let a callback error break a request
+            /* swallow */
         }
     }
     /**
@@ -110,11 +148,20 @@ export class AIBudget {
         const started = await ctx.runMutation(this.component.lib.startRequest, {
             userId,
             actionName,
+            tags: args.tags,
             model,
             messages,
             rerunOf: args.rerunOf,
         });
         if (!started.allowed) {
+            await this.fireLimitReached({
+                userId,
+                action: actionName,
+                tags: args.tags,
+                messages: [started.reason],
+                code: started.code,
+                reason: started.reason,
+            });
             throw new ConvexError({
                 kind: "AIBudgetLimit",
                 code: started.code,
@@ -123,7 +170,8 @@ export class AIBudget {
         }
         const requestId = started.requestId;
         const warnings = started.warnings;
-        await this.fireSoftLimit({ userId, action: actionName, requestId, warnings });
+        const notices = started.notices;
+        await this.fireBudgetEvents({ userId, action: actionName, tags: args.tags, requestId }, warnings, notices);
         const start = Date.now();
         try {
             // The full chain (incl. system) is stored on the request for audit/replay,
@@ -143,9 +191,10 @@ export class AIBudget {
                 requestId,
                 responseText: result.text,
                 ...usage,
+                costNanos: extractGatewayCostNanos(result),
                 latencyMs: Date.now() - start,
             });
-            return { text: result.text, requestId, costNanos, warnings, ...usage };
+            return { text: result.text, requestId, costNanos, warnings, notices, ...usage };
         }
         catch (e) {
             await ctx.runMutation(this.component.lib.finishRequest, {
@@ -165,29 +214,33 @@ export class AIBudget {
     languageModel(ctx, opts = {}) {
         const modelId = opts.model ?? this.defaultModel;
         const component = this.component;
-        const fireSoftLimit = this.fireSoftLimit.bind(this);
+        const fireBudgetEvents = this.fireBudgetEvents.bind(this);
+        const fireLimitReached = this.fireLimitReached.bind(this);
         const begin = async (params) => {
             const userId = await resolveUserId(ctx, opts.userId);
             const actionName = await resolveActionName(ctx, opts.action);
+            const base = { userId, action: actionName, tags: opts.tags };
             const started = await ctx.runMutation(component.lib.startRequest, {
                 userId,
                 actionName,
+                tags: opts.tags,
                 model: modelId,
                 messages: simplifyPrompt(params.prompt),
             });
             if (!started.allowed) {
+                await fireLimitReached({
+                    ...base,
+                    messages: [started.reason],
+                    code: started.code,
+                    reason: started.reason,
+                });
                 throw new ConvexError({
                     kind: "AIBudgetLimit",
                     code: started.code,
                     reason: started.reason,
                 });
             }
-            await fireSoftLimit({
-                userId,
-                action: actionName,
-                requestId: started.requestId,
-                warnings: started.warnings,
-            });
+            await fireBudgetEvents({ ...base, requestId: started.requestId }, started.warnings, started.notices);
             return started.requestId;
         };
         const finish = async (requestId, fields) => ctx.runMutation(component.lib.finishRequest, { requestId, ...fields });
@@ -202,6 +255,7 @@ export class AIBudget {
                         await finish(requestId, {
                             responseText: extractText(result),
                             ...extractUsage(result.usage),
+                            costNanos: extractGatewayCostNanos(result),
                             latencyMs: Date.now() - start,
                         });
                         return result;
@@ -278,6 +332,7 @@ export class AIBudget {
             messages: args.messages ?? original.messages,
             rerunOf: args.requestId,
             action: original.actionName,
+            tags: original.tags,
         });
     }
     // ---------- namespaced admin API ----------
@@ -285,6 +340,7 @@ export class AIBudget {
     get requests() {
         const c = this.component;
         return {
+            /** Filter by userId, or by any {dimension, value} (incl. custom tags). */
             list: (ctx, args = {}) => ctx.runQuery(c.lib.listRequests, args),
             /** Ancestors up to the original, plus direct re-runs. */
             lineage: (ctx, args) => ctx.runQuery(c.lib.lineage, { requestId: args.requestId }),
@@ -292,28 +348,78 @@ export class AIBudget {
             rerun: (ctx, args) => this.rerunImpl(ctx, args),
         };
     }
-    /** Per-user budgets and controls. */
-    get users() {
+    /**
+     * Budgets and controls for an arbitrary attribution dimension — the
+     * generalization of `users`/`actions`. Give it any dimension name (team,
+     * project, tenant, customer, env, feature, …) and set caps per value:
+     *
+     *   ai.tag("customer").setLimits(ctx, { value: "acme", monthlySpendLimitNanos });
+     *   ai.tag("customer").history(ctx, { value: "acme", period: "day" });
+     *
+     * Attribute a call to it by passing `tags` to `chat`/`languageModel`.
+     */
+    tag(dimension) {
+        return this.dimensionApi(dimension, (a) => a.value);
+    }
+    // Shared implementation behind tag()/users/actions. `key` maps the namespace's
+    // id field (value/userId/name) to the bucket value.
+    dimensionApi(dimension, key) {
         const c = this.component;
         return {
-            list: (ctx) => ctx.runQuery(c.lib.listUsers, {}),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setLimits, args),
-            /** One-time "approve another $X" bump (daily is today-only). */
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpUser, args),
-            /** Delete a user and all their request rows. */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteUser, args),
+            /** All buckets in this dimension. */
+            list: (ctx) => ctx.runQuery(c.lib.listBuckets, { dimension }),
+            /** One bucket's limits + spend (null if it has none yet). */
+            get: (ctx, args) => ctx.runQuery(c.lib.getBucket, { dimension, value: key(args) }),
+            setLimits: (ctx, args) => {
+                const { value, userId, name, ...limits } = args;
+                return ctx.runMutation(c.lib.setBucketLimits, {
+                    dimension,
+                    value: key(args),
+                    ...limits,
+                });
+            },
+            /** One-time "approve another $X" bump (daily/monthly reset with the window). */
+            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
+                dimension,
+                value: key(args),
+                dailyNanos: args.dailyNanos,
+                monthlyNanos: args.monthlyNanos,
+                lifetimeNanos: args.lifetimeNanos,
+            }),
+            /** Manually credit (negative) or debit (positive) this bucket. */
+            adjust: (ctx, args) => ctx.runMutation(c.lib.adjustBucket, {
+                dimension,
+                value: key(args),
+                deltaNanos: args.deltaNanos,
+                tokens: args.tokens,
+                reason: args.reason,
+            }),
+            /** Durable spend history for this bucket (per day or per month). */
+            history: (ctx, args) => ctx.runQuery(c.lib.usageHistory, {
+                dimension,
+                value: key(args),
+                period: args.period ?? "day",
+                limit: args.limit,
+            }),
+            /** Manual-adjustment audit log for this bucket. */
+            adjustments: (ctx, args) => ctx.runQuery(c.lib.listAdjustments, {
+                dimension,
+                value: key(args),
+                limit: args.limit,
+            }),
+            /** Delete the bucket (for "user", also its request rows). */
+            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, { dimension, value: key(args) }),
         };
     }
-    /** Per-action (per-feature) budgets. */
+    /** Per-user budgets and controls — sugar over the "user" dimension. */
+    get users() {
+        return this.dimensionApi("user", (a) => a.userId);
+    }
+    /** Per-action (per-feature) budgets — sugar over the "action" dimension. */
     get actions() {
-        const c = this.component;
-        return {
-            list: (ctx) => ctx.runQuery(c.lib.listActions, {}),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setActionLimits, args),
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpAction, args),
-        };
+        return this.dimensionApi("action", (a) => a.name);
     }
-    /** The deployment-wide budget and retention config. */
+    /** The deployment-wide budget, alerts, and retention config. */
     get global() {
         const c = this.component;
         return {
@@ -322,6 +428,8 @@ export class AIBudget {
             /** A killswitch spend cap across all users/actions (enforced approximately). */
             setLimits: (ctx, args) => ctx.runMutation(c.lib.setGlobalLimits, args),
             bump: (ctx, args) => ctx.runMutation(c.lib.bumpGlobal, args),
+            /** Default approaching-limit alert threshold (fraction of a cap, e.g. 0.8). */
+            setAlertDefaults: (ctx, args) => ctx.runMutation(c.lib.setAlertDefaults, args),
             /** Request-row retention window in ms (default 1h; 0 disables). */
             setRetention: (ctx, args) => ctx.runMutation(c.lib.setRetention, args),
         };

package/dist/component/_generated/component.d.ts

@@ -20,22 +20,28 @@ import type { FunctionReference } from "convex/server";
  */
 export type ComponentApi<Name extends string | undefined = string | undefined> = {
     lib: {
-        bumpAction: FunctionReference<"mutation", "internal", {
+        adjustBucket: FunctionReference<"mutation", "internal", {
+            deltaNanos: number;
+            dimension: string;
+            reason?: string;
+            tokens?: number;
+            value: string;
+        }, null, Name>;
+        bumpBucket: FunctionReference<"mutation", "internal", {
             dailyNanos?: number;
+            dimension: string;
             lifetimeNanos?: number;
-            name: string;
+            monthlyNanos?: number;
+            value: string;
         }, null, Name>;
         bumpGlobal: FunctionReference<"mutation", "internal", {
             dailyNanos?: number;
             lifetimeNanos?: number;
+            monthlyNanos?: number;
         }, null, Name>;
-        bumpUser: FunctionReference<"mutation", "internal", {
-            dailyNanos?: number;
-            lifetimeNanos?: number;
-            userId: string;
-        }, null, Name>;
-        deleteUser: FunctionReference<"mutation", "internal", {
-            userId: string;
+        deleteBucket: FunctionReference<"mutation", "internal", {
+            dimension: string;
+            value: string;
         }, {
             deletedThisBatch: number;
             done: boolean;
@@ -43,6 +49,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         finishRequest: FunctionReference<"mutation", "internal", {
             cachedTokens?: number;
             completionTokens?: number;
+            costNanos?: number;
             error?: string;
             latencyMs?: number;
             promptTokens?: number;
@@ -51,6 +58,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         }, {
             costNanos: number;
         }, Name>;
+        getBucket: FunctionReference<"query", "internal", {
+            dimension: string;
+            value: string;
+        }, any, Name>;
         getGlobalStatus: FunctionReference<"query", "internal", {}, {
             dailySpendLimitNanos: number | null;
             enforcement: "hard" | "soft";
@@ -68,42 +79,50 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         lineage: FunctionReference<"query", "internal", {
             requestId: string;
         }, any, Name>;
-        listActions: FunctionReference<"query", "internal", {}, any, Name>;
+        listAdjustments: FunctionReference<"query", "internal", {
+            dimension: string;
+            limit?: number;
+            value: string;
+        }, any, Name>;
+        listBuckets: FunctionReference<"query", "internal", {
+            dimension?: string;
+        }, any, Name>;
         listPrices: FunctionReference<"query", "internal", {}, any, Name>;
         listRequests: FunctionReference<"query", "internal", {
+            dimension?: string;
             limit?: number;
             userId?: string;
+            value?: string;
         }, any, Name>;
-        listUsers: FunctionReference<"query", "internal", {}, any, Name>;
-        setActionLimits: FunctionReference<"mutation", "internal", {
+        setAlertDefaults: FunctionReference<"mutation", "internal", {
+            warnAtPct?: number;
+        }, null, Name>;
+        setBucketLimits: FunctionReference<"mutation", "internal", {
+            blocked?: boolean;
             dailySpendLimitNanos?: number;
             dailyTokenLimit?: number;
-            disabled?: boolean;
+            dimension: string;
             enforcement?: "hard" | "soft";
             lifetimeSpendLimitNanos?: number;
             lifetimeTokenLimit?: number;
-            name: string;
+            maxConcurrent?: number;
+            monthlySpendLimitNanos?: number;
+            monthlyTokenLimit?: number;
+            requestsPerMinute?: number;
+            value: string;
+            warnAtPct?: number;
         }, null, Name>;
         setGlobalLimits: FunctionReference<"mutation", "internal", {
             dailySpendLimitNanos?: number;
             enforcement?: "hard" | "soft";
             lifetimeSpendLimitNanos?: number;
         }, null, Name>;
-        setLimits: FunctionReference<"mutation", "internal", {
-            blocked?: boolean;
-            dailySpendLimitNanos?: number;
-            dailyTokenLimit?: number;
-            enforcement?: "hard" | "soft";
-            lifetimeSpendLimitNanos?: number;
-            lifetimeTokenLimit?: number;
-            requestsPerMinute?: number;
-            userId: string;
-        }, null, Name>;
         setModelPolicy: FunctionReference<"mutation", "internal", {
             mode: "open" | "allowlist" | "denylist";
             models: Array<string>;
         }, null, Name>;
         setPrice: FunctionReference<"mutation", "internal", {
+            cachedNanosPerMTok?: number;
             inputNanosPerMTok: number;
             model: string;
             outputNanosPerMTok: number;
@@ -119,9 +138,14 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             }>;
             model: string;
             rerunOf?: string;
+            tags?: Array<{
+                dimension: string;
+                value: string;
+            }>;
             userId: string;
         }, {
             allowed: true;
+            notices: Array<string>;
             requestId: string;
             warnings: Array<string>;
         } | {
@@ -129,5 +153,11 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             code: string;
             reason: string;
         }, Name>;
+        usageHistory: FunctionReference<"query", "internal", {
+            dimension: string;
+            limit?: number;
+            period: "day" | "month";
+            value: string;
+        }, any, Name>;
     };
 };