@convex-dev/ai-budget 0.0.2-alpha.4 → 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;
     }
-    async fireSoftLimit(info) {
-        if (info.warnings.length === 0 || !this.onSoftLimit)
+    // 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 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 */
         }
     }
     /**
@@ -116,6 +154,14 @@ export class AIBudget {
             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,
@@ -124,13 +170,8 @@ export class AIBudget {
         }
         const requestId = started.requestId;
         const warnings = started.warnings;
-        await this.fireSoftLimit({
-            userId,
-            action: actionName,
-            tags: args.tags,
-            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,
@@ -150,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, {
@@ -172,10 +214,12 @@ 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,
@@ -184,19 +228,19 @@ export class AIBudget {
                 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,
-                tags: opts.tags,
-                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 });
@@ -211,6 +255,7 @@ export class AIBudget {
                         await finish(requestId, {
                             responseText: extractText(result),
                             ...extractUsage(result.usage),
+                            costNanos: extractGatewayCostNanos(result),
                             latencyMs: Date.now() - start,
                         });
                         return result;
@@ -295,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 }),
@@ -307,74 +353,73 @@ export class AIBudget {
      * 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", dailySpendLimitNanos });
-     *   ai.tag("customer").list(ctx);
+     *   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 {
             /** 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: args.value }),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setBucketLimits, { dimension, ...args }),
-            /** One-time "approve another $X" bump (daily is today-only). */
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, { dimension, ...args }),
-            /** Delete the bucket (for "user", also its request rows). */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, { dimension, value: args.value }),
-        };
-    }
-    /** Per-user budgets and controls — sugar over the "user" dimension. */
-    get users() {
-        const c = this.component;
-        return {
-            list: (ctx) => ctx.runQuery(c.lib.listBuckets, { dimension: "user" }),
+            get: (ctx, args) => ctx.runQuery(c.lib.getBucket, { dimension, value: key(args) }),
             setLimits: (ctx, args) => {
-                const { userId, ...limits } = args;
+                const { value, userId, name, ...limits } = args;
                 return ctx.runMutation(c.lib.setBucketLimits, {
-                    dimension: "user",
-                    value: userId,
+                    dimension,
+                    value: key(args),
                     ...limits,
                 });
             },
-            /** One-time "approve another $X" bump (daily is today-only). */
+            /** One-time "approve another $X" bump (daily/monthly reset with the window). */
             bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
-                dimension: "user",
-                value: args.userId,
+                dimension,
+                value: key(args),
                 dailyNanos: args.dailyNanos,
+                monthlyNanos: args.monthlyNanos,
                 lifetimeNanos: args.lifetimeNanos,
             }),
-            /** Delete a user and all their request rows. */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, {
-                dimension: "user",
-                value: args.userId,
+            /** 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-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.listBuckets, { dimension: "action" }),
-            setLimits: (ctx, args) => {
-                const { name, ...limits } = args;
-                return ctx.runMutation(c.lib.setBucketLimits, {
-                    dimension: "action",
-                    value: name,
-                    ...limits,
-                });
-            },
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
-                dimension: "action",
-                value: args.name,
-                dailyNanos: args.dailyNanos,
-                lifetimeNanos: args.lifetimeNanos,
-            }),
-        };
+        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 {
@@ -383,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,15 +20,24 @@ import type { FunctionReference } from "convex/server";
  */
 export type ComponentApi<Name extends string | undefined = string | undefined> = {
     lib: {
+        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;
+            monthlyNanos?: number;
             value: string;
         }, null, Name>;
         bumpGlobal: FunctionReference<"mutation", "internal", {
             dailyNanos?: number;
             lifetimeNanos?: number;
+            monthlyNanos?: number;
         }, null, Name>;
         deleteBucket: FunctionReference<"mutation", "internal", {
             dimension: string;
@@ -40,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;
@@ -69,14 +79,24 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         lineage: FunctionReference<"query", "internal", {
             requestId: string;
         }, 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>;
+        setAlertDefaults: FunctionReference<"mutation", "internal", {
+            warnAtPct?: number;
+        }, null, Name>;
         setBucketLimits: FunctionReference<"mutation", "internal", {
             blocked?: boolean;
             dailySpendLimitNanos?: number;
@@ -85,8 +105,12 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             enforcement?: "hard" | "soft";
             lifetimeSpendLimitNanos?: number;
             lifetimeTokenLimit?: number;
+            maxConcurrent?: number;
+            monthlySpendLimitNanos?: number;
+            monthlyTokenLimit?: number;
             requestsPerMinute?: number;
             value: string;
+            warnAtPct?: number;
         }, null, Name>;
         setGlobalLimits: FunctionReference<"mutation", "internal", {
             dailySpendLimitNanos?: number;
@@ -98,6 +122,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             models: Array<string>;
         }, null, Name>;
         setPrice: FunctionReference<"mutation", "internal", {
+            cachedNanosPerMTok?: number;
             inputNanosPerMTok: number;
             model: string;
             outputNanosPerMTok: number;
@@ -120,6 +145,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             userId: string;
         }, {
             allowed: true;
+            notices: Array<string>;
             requestId: string;
             warnings: Array<string>;
         } | {
@@ -127,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>;
     };
 };

package/dist/component/lib.d.ts

@@ -19,6 +19,7 @@ export declare const startRequest: import("convex/server").RegisteredMutation<"p
     allowed: true;
     requestId: import("convex/values").GenericId<"requests">;
     warnings: string[];
+    notices: string[];
 }>>;
 export declare const finishRequest: import("convex/server").RegisteredMutation<"public", {
     error?: string | undefined;
@@ -26,6 +27,7 @@ export declare const finishRequest: import("convex/server").RegisteredMutation<"
     promptTokens?: number | undefined;
     completionTokens?: number | undefined;
     cachedTokens?: number | undefined;
+    costNanos?: number | undefined;
     latencyMs?: number | undefined;
     requestId: import("convex/values").GenericId<"requests">;
 }, Promise<{
@@ -136,6 +138,8 @@ export declare const getRequest: import("convex/server").RegisteredQuery<"public
     status: "blocked" | "pending" | "success" | "error";
 } | null>>;
 export declare const listRequests: import("convex/server").RegisteredQuery<"public", {
+    dimension?: string | undefined;
+    value?: string | undefined;
     userId?: string | undefined;
     limit?: number | undefined;
 }, Promise<{
@@ -171,22 +175,33 @@ export declare const listBuckets: import("convex/server").RegisteredQuery<"publi
     dimension?: string | undefined;
 }, Promise<{
     spendTodayNanos: number;
+    spendThisMonthNanos: number;
     _id: import("convex/values").GenericId<"buckets">;
     _creationTime: number;
     requestsPerMinute?: number | undefined;
+    maxConcurrent?: number | undefined;
     dailySpendLimitNanos?: number | undefined;
+    monthlySpendLimitNanos?: number | undefined;
     lifetimeSpendLimitNanos?: number | undefined;
     dailyTokenLimit?: number | undefined;
+    monthlyTokenLimit?: number | undefined;
     lifetimeTokenLimit?: number | undefined;
     blocked?: boolean | undefined;
+    warnAtPct?: number | undefined;
     enforcement?: "hard" | "soft" | undefined;
     dailyBumpNanos?: number | undefined;
+    monthlyBumpNanos?: number | undefined;
     lifetimeBumpNanos?: number | undefined;
     bumpDayStamp?: string | undefined;
+    bumpMonthStamp?: string | undefined;
     tokensToday?: number | undefined;
+    monthStamp?: string | undefined;
+    tokensThisMonth?: number | undefined;
     reservedTodayNanos?: number | undefined;
+    reservedMonthNanos?: number | undefined;
     reservedTotalNanos?: number | undefined;
     reservedTodayTokens?: number | undefined;
+    reservedMonthTokens?: number | undefined;
     reservedTotalTokens?: number | undefined;
     pendingCount?: number | undefined;
     dimension: string;
@@ -201,22 +216,33 @@ export declare const getBucket: import("convex/server").RegisteredQuery<"public"
     value: string;
 }, Promise<{
     spendTodayNanos: number;
+    spendThisMonthNanos: number;
     _id: import("convex/values").GenericId<"buckets">;
     _creationTime: number;
     requestsPerMinute?: number | undefined;
+    maxConcurrent?: number | undefined;
     dailySpendLimitNanos?: number | undefined;
+    monthlySpendLimitNanos?: number | undefined;
     lifetimeSpendLimitNanos?: number | undefined;
     dailyTokenLimit?: number | undefined;
+    monthlyTokenLimit?: number | undefined;
     lifetimeTokenLimit?: number | undefined;
     blocked?: boolean | undefined;
+    warnAtPct?: number | undefined;
     enforcement?: "hard" | "soft" | undefined;
     dailyBumpNanos?: number | undefined;
+    monthlyBumpNanos?: number | undefined;
     lifetimeBumpNanos?: number | undefined;
     bumpDayStamp?: string | undefined;
+    bumpMonthStamp?: string | undefined;
     tokensToday?: number | undefined;
+    monthStamp?: string | undefined;
+    tokensThisMonth?: number | undefined;
     reservedTodayNanos?: number | undefined;
+    reservedMonthNanos?: number | undefined;
     reservedTotalNanos?: number | undefined;
     reservedTodayTokens?: number | undefined;
+    reservedMonthTokens?: number | undefined;
     reservedTotalTokens?: number | undefined;
     pendingCount?: number | undefined;
     dimension: string;
@@ -228,21 +254,65 @@ export declare const getBucket: import("convex/server").RegisteredQuery<"public"
 } | null>>;
 export declare const setBucketLimits: import("convex/server").RegisteredMutation<"public", {
     requestsPerMinute?: number | undefined;
+    maxConcurrent?: number | undefined;
     dailySpendLimitNanos?: number | undefined;
+    monthlySpendLimitNanos?: number | undefined;
     lifetimeSpendLimitNanos?: number | undefined;
     dailyTokenLimit?: number | undefined;
+    monthlyTokenLimit?: number | undefined;
     lifetimeTokenLimit?: number | undefined;
     blocked?: boolean | undefined;
+    warnAtPct?: number | undefined;
     enforcement?: "hard" | "soft" | undefined;
     dimension: string;
     value: string;
 }, Promise<null>>;
 export declare const bumpBucket: import("convex/server").RegisteredMutation<"public", {
     dailyNanos?: number | undefined;
+    monthlyNanos?: number | undefined;
     lifetimeNanos?: number | undefined;
     dimension: string;
     value: string;
 }, Promise<null>>;
+export declare const adjustBucket: import("convex/server").RegisteredMutation<"public", {
+    tokens?: number | undefined;
+    reason?: string | undefined;
+    dimension: string;
+    value: string;
+    deltaNanos: number;
+}, Promise<null>>;
+export declare const listAdjustments: import("convex/server").RegisteredQuery<"public", {
+    limit?: number | undefined;
+    dimension: string;
+    value: string;
+}, Promise<{
+    _id: import("convex/values").GenericId<"adjustments">;
+    _creationTime: number;
+    tokens?: number | undefined;
+    reason?: string | undefined;
+    dimension: string;
+    value: string;
+    deltaNanos: number;
+}[]>>;
+export declare const usageHistory: import("convex/server").RegisteredQuery<"public", {
+    limit?: number | undefined;
+    dimension: string;
+    value: string;
+    period: "day" | "month";
+}, Promise<{
+    _id: import("convex/values").GenericId<"usage">;
+    _creationTime: number;
+    dimension: string;
+    value: string;
+    period: "day" | "month";
+    stamp: string;
+    spendNanos: number;
+    tokens: number;
+    requests: number;
+}[]>>;
+export declare const setAlertDefaults: import("convex/server").RegisteredMutation<"public", {
+    warnAtPct?: number | undefined;
+}, Promise<null>>;
 export declare const deleteBucket: import("convex/server").RegisteredMutation<"public", {
     dimension: string;
     value: string;
@@ -268,6 +338,7 @@ export declare const setGlobalLimits: import("convex/server").RegisteredMutation
 }, Promise<null>>;
 export declare const bumpGlobal: import("convex/server").RegisteredMutation<"public", {
     dailyNanos?: number | undefined;
+    monthlyNanos?: number | undefined;
     lifetimeNanos?: number | undefined;
 }, Promise<null>>;
 export declare const setModelPolicy: import("convex/server").RegisteredMutation<"public", {
@@ -275,6 +346,7 @@ export declare const setModelPolicy: import("convex/server").RegisteredMutation<
     mode: "open" | "allowlist" | "denylist";
 }, Promise<null>>;
 export declare const setPrice: import("convex/server").RegisteredMutation<"public", {
+    cachedNanosPerMTok?: number | undefined;
     model: string;
     inputNanosPerMTok: number;
     outputNanosPerMTok: number;
@@ -282,5 +354,6 @@ export declare const setPrice: import("convex/server").RegisteredMutation<"publi
 export declare const listPrices: import("convex/server").RegisteredQuery<"public", {}, Promise<Record<string, {
     input: number;
     output: number;
+    cached?: number;
     overridden: boolean;
 }>>>;