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

package/README.md

@@ -25,10 +25,11 @@ full audit log you can replay later.
 |---|---|
 | **Usage & cost tracking** | Every request stored with messages, response, tokens, latency, and per-request cost. |
 | **Attribution** | Each call is attributed to a `userId` **and** to the Convex action that made it — auto-detected via `ctx.meta`, no manual tagging. Running totals per user and per action. |
+| **Tagged budgets** | `user` and `action` are just built-in *dimensions* — add your own (team, project, customer, env, feature…) by passing `tags`, and cap any of them with `ai.tag("customer").setLimits(...)`. One request can be billed to several buckets at once. |
 | **Spend & token limits** | Per-user daily / lifetime **spend** and **token** budgets, plus a requests-per-minute rate limit and a block switch. |
 | **Concurrency-safe caps** | A reserve-then-settle design makes admission a true atomic check — concurrent in-flight requests can't blow past the cap (a naive implementation overshoots ~40×). |
 | **Hard or soft** | Each limit either **blocks** (`hard`) or **allows-with-a-warning** (`soft`). |
-| **Per-feature budgets** | Cap or disable a whole action (e.g. `summarize`) independently of any user. |
+| **Per-feature budgets** | Cap or block a whole action (e.g. `summarize`) independently of any user. |
 | **Global killswitch** | A deployment-wide spend cap across all users and actions (sharded for throughput; enforced approximately). |
 | **One-time bumps** | "Approve another $X" at any level (user / action / global) without changing the standing cap — daily bumps are today-only, lifetime bumps permanent. |
 | **Model policy** | Allow/deny lists for models; an unknown/unpriced model **fails closed** (charged a conservative max, never $0). |
@@ -136,6 +137,7 @@ ai.chat(ctx, {
   messages?,         // [{ role, content }]
   model?,            // defaults to defaultModel
   action?,           // attribution name; defaults to the calling Convex action
+  tags?,             // extra dimensions: [{ dimension: "customer", value: "acme" }, …]
 }): Promise<{ text, requestId, costNanos, promptTokens, completionTokens, warnings }>
 ```
 
@@ -208,10 +210,42 @@ ai.actions.setLimits(ctx, {
   dailyTokenLimit?,
   lifetimeTokenLimit?,
   enforcement?,                  // "hard" | "soft"
-  disabled?,                     // kill switch for the whole feature
+  blocked?,                      // kill switch for the whole feature
 })
 ```
 
+### Tagged budgets (custom dimensions)
+
+`user` and `action` are the two built-in *dimensions*. To classify or budget
+along any other axis — team, project, tenant, customer, environment, feature —
+attach `tags` to a call and cap a value with `ai.tag(dimension)`:
+
+```ts
+// Bill this call to a user, an action (implicit), AND a customer + env.
+await ai.chat(ctx, {
+  prompt,
+  tags: [
+    { dimension: "customer", value: "acme" },
+    { dimension: "env", value: "prod" },
+  ],
+});
+
+// Cap the customer "acme" to $50/day — independent of any per-user cap.
+await ai.tag("customer").setLimits(ctx, {
+  value: "acme",
+  dailySpendLimitNanos: 50 * 1_000_000_000,
+});
+await ai.tag("customer").bump(ctx, { value: "acme", dailyNanos: 10 * 1_000_000_000 });
+await ai.tag("customer").list(ctx);                 // every customer's spend & caps
+await ai.tag("customer").get(ctx, { value: "acme" }); // one bucket
+```
+
+A single request is admitted only if it fits **every** bucket it touches (user,
+action, and each tag) — the same exact reserve-then-settle check runs per bucket.
+Uncapped buckets never serialize, so adding tags you don't cap is free at
+admission; their running totals still accrue for reporting. `ai.users.*` and
+`ai.actions.*` are simply sugar over `ai.tag("user")` / `ai.tag("action")`.
+
 ### Global (deployment-wide) budget
 
 ```ts
@@ -279,6 +313,7 @@ flagged `unpricedModel: true` so you know to add a real price.
 ```ts
 ai.users.list(ctx)                      // per-user spend today / total / tokens / limits
 ai.actions.list(ctx)                    // per-action spend & totals
+ai.tag("customer").list(ctx)            // spend & caps for any custom dimension
 ai.requests.list(ctx, { userId?, limit? })  // the audit log (blocked attempts included)
 ```
 
@@ -305,15 +340,18 @@ same pre-spend total and all pass, so spend blows past the cap (measured at
    **exactly-once** (a terminal request is never re-folded), so a slow request
    that the reconciler already swept can't double-count when it finally returns.
 
-Reservations are only taken on entities that actually have a cap, so uncapped
-traffic never serializes. The `error.md` file documents the adversarial audits
-this design survived, with live repros.
-
-**One guarantee, all scopes.** Per-user, per-action, and global caps run through
-the *same* admission check — a request is admitted only when
-`committed + reserved + estimate ≤ cap` (bumps included). The only thing that
-differs is the holder: per-user and per-action reserve on a single document, an
-exact atomic check-and-reserve; the **global** killswitch is backed by a sharded
+Reservations are only taken on buckets that actually have a cap, so uncapped
+traffic never serializes — this is what makes arbitrary `tags` cheap: a request
+reserves on one row per *capped* dimension it carries, and nothing else. The
+`error.md` file documents the adversarial audits this design survived, with live
+repros.
+
+**One guarantee, all scopes.** Every per-bucket cap — user, action, or any
+custom tag dimension — and the global cap run through the *same* admission check:
+a request is admitted only when `committed + reserved + estimate ≤ cap` (bumps
+included) for **every** bucket it touches. The only thing that differs is the
+holder: each per-bucket cap reserves on a single document, an exact atomic
+check-and-reserve; the **global** killswitch is backed by a sharded
 counter for throughput, so its committed total is read as an eventually-consistent
 sum with no cross-request reservation. That makes the global cap **approximate** —
 it can overshoot by a bounded amount under a burst — the deliberate

package/dist/client/index.d.ts

@@ -11,10 +11,21 @@ type UseApi<API> = Expand<{
 export type AIBudgetApi = UseApi<typeof api>;
 /** @deprecated use AIBudgetApi */
 export type AIGatewayApi = AIBudgetApi;
+/**
+ * One attribution tag: a (dimension, value) pair, e.g. {dimension:"customer",
+ * value:"acme"}. `user` and `action` are built-in dimensions (set via
+ * userId/action); use tags for anything else — team, project, tenant, env, ….
+ * Any tagged bucket can carry its own budget (see `ai.tag(dimension)`).
+ */
+export type Tag = {
+    dimension: string;
+    value: string;
+};
 /** Fired when a request is admitted over a *soft* limit. */
 export type SoftLimitInfo = {
     userId: string;
     action?: string;
+    tags?: Tag[];
     requestId: string;
     warnings: string[];
 };
@@ -76,6 +87,8 @@ export declare class AIBudget {
         rerunOf?: string;
         /** Attribute spend to this action name. Defaults to the calling Convex action. */
         action?: string;
+        /** Extra attribution dimensions to bill/limit (team, customer, env, …). */
+        tags?: Tag[];
     }): Promise<ChatResult>;
     /**
      * An AI SDK LanguageModel that enforces limits and records usage/cost for
@@ -87,6 +100,8 @@ export declare class AIBudget {
         userId?: string;
         model?: string;
         action?: string;
+        /** Extra attribution dimensions to bill/limit (team, customer, env, …). */
+        tags?: Tag[];
     }): LanguageModel;
     private rerunImpl;
     /** The request audit log, replay, and re-run lineage. */
@@ -98,6 +113,10 @@ export declare class AIBudget {
             _id: string;
             _creationTime: number;
             actionName?: string | undefined;
+            tags?: {
+                dimension: string;
+                value: string;
+            }[] | undefined;
             estimatedNanos?: number | undefined;
             estimatedTokens?: number | undefined;
             unpricedModel?: boolean | undefined;
@@ -127,6 +146,10 @@ export declare class AIBudget {
                 _id: string;
                 _creationTime: number;
                 actionName?: string | undefined;
+                tags?: {
+                    dimension: string;
+                    value: string;
+                }[] | undefined;
                 estimatedNanos?: number | undefined;
                 estimatedTokens?: number | undefined;
                 unpricedModel?: boolean | undefined;
@@ -152,6 +175,10 @@ export declare class AIBudget {
                 _id: string;
                 _creationTime: number;
                 actionName?: string | undefined;
+                tags?: {
+                    dimension: string;
+                    value: string;
+                }[] | undefined;
                 estimatedNanos?: number | undefined;
                 estimatedTokens?: number | undefined;
                 unpricedModel?: boolean | undefined;
@@ -181,7 +208,100 @@ export declare class AIBudget {
             model?: string;
         }) => Promise<ChatResult>;
     };
-    /** Per-user budgets and controls. */
+    /**
+     * 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", dailySpendLimitNanos });
+     *   ai.tag("customer").list(ctx);
+     *
+     * Attribute a call to it by passing `tags` to `chat`/`languageModel`.
+     */
+    tag(dimension: string): {
+        /** All buckets in this dimension. */
+        list: (ctx: RunQueryCtx) => Promise<{
+            spendTodayNanos: number;
+            _id: string;
+            _creationTime: number;
+            requestsPerMinute?: number | undefined;
+            dailySpendLimitNanos?: number | undefined;
+            lifetimeSpendLimitNanos?: number | undefined;
+            dailyTokenLimit?: number | undefined;
+            lifetimeTokenLimit?: number | undefined;
+            blocked?: boolean | undefined;
+            enforcement?: "hard" | "soft" | undefined;
+            dailyBumpNanos?: number | undefined;
+            lifetimeBumpNanos?: number | undefined;
+            bumpDayStamp?: string | undefined;
+            tokensToday?: number | undefined;
+            reservedTodayNanos?: number | undefined;
+            reservedTotalNanos?: number | undefined;
+            reservedTodayTokens?: number | undefined;
+            reservedTotalTokens?: number | undefined;
+            pendingCount?: number | undefined;
+            dimension: string;
+            value: string;
+            totalSpendNanos: number;
+            totalRequests: number;
+            totalTokens: number;
+            dayStamp: string;
+        }[]>;
+        /** One bucket's limits + spend (null if it has none yet). */
+        get: (ctx: RunQueryCtx, args: {
+            value: string;
+        }) => Promise<{
+            spendTodayNanos: number;
+            _id: string;
+            _creationTime: number;
+            requestsPerMinute?: number | undefined;
+            dailySpendLimitNanos?: number | undefined;
+            lifetimeSpendLimitNanos?: number | undefined;
+            dailyTokenLimit?: number | undefined;
+            lifetimeTokenLimit?: number | undefined;
+            blocked?: boolean | undefined;
+            enforcement?: "hard" | "soft" | undefined;
+            dailyBumpNanos?: number | undefined;
+            lifetimeBumpNanos?: number | undefined;
+            bumpDayStamp?: string | undefined;
+            tokensToday?: number | undefined;
+            reservedTodayNanos?: number | undefined;
+            reservedTotalNanos?: number | undefined;
+            reservedTodayTokens?: number | undefined;
+            reservedTotalTokens?: number | undefined;
+            pendingCount?: number | undefined;
+            dimension: string;
+            value: string;
+            totalSpendNanos: number;
+            totalRequests: number;
+            totalTokens: number;
+            dayStamp: string;
+        } | null>;
+        setLimits: (ctx: RunMutationCtx, args: {
+            value: string;
+            requestsPerMinute?: number;
+            dailySpendLimitNanos?: number;
+            lifetimeSpendLimitNanos?: number;
+            dailyTokenLimit?: number;
+            lifetimeTokenLimit?: number;
+            enforcement?: "hard" | "soft";
+            blocked?: boolean;
+        }) => Promise<null>;
+        /** One-time "approve another $X" bump (daily is today-only). */
+        bump: (ctx: RunMutationCtx, args: {
+            value: string;
+            dailyNanos?: number;
+            lifetimeNanos?: number;
+        }) => Promise<null>;
+        /** Delete the bucket (for "user", also its request rows). */
+        delete: (ctx: RunMutationCtx, args: {
+            value: string;
+        }) => Promise<{
+            deletedThisBatch: number;
+            done: boolean;
+        }>;
+    };
+    /** Per-user budgets and controls — sugar over the "user" dimension. */
     get users(): {
         list: (ctx: RunQueryCtx) => Promise<{
             spendTodayNanos: number;
@@ -203,7 +323,8 @@ export declare class AIBudget {
             reservedTodayTokens?: number | undefined;
             reservedTotalTokens?: number | undefined;
             pendingCount?: number | undefined;
-            userId: string;
+            dimension: string;
+            value: string;
             totalSpendNanos: number;
             totalRequests: number;
             totalTokens: number;
@@ -233,16 +354,18 @@ export declare class AIBudget {
             done: boolean;
         }>;
     };
-    /** Per-action (per-feature) budgets. */
+    /** Per-action (per-feature) budgets — sugar over the "action" dimension. */
     get actions(): {
         list: (ctx: RunQueryCtx) => Promise<{
             spendTodayNanos: number;
             _id: string;
             _creationTime: number;
+            requestsPerMinute?: number | undefined;
             dailySpendLimitNanos?: number | undefined;
             lifetimeSpendLimitNanos?: number | undefined;
             dailyTokenLimit?: number | undefined;
             lifetimeTokenLimit?: number | undefined;
+            blocked?: boolean | undefined;
             enforcement?: "hard" | "soft" | undefined;
             dailyBumpNanos?: number | undefined;
             lifetimeBumpNanos?: number | undefined;
@@ -253,12 +376,12 @@ export declare class AIBudget {
             reservedTodayTokens?: number | undefined;
             reservedTotalTokens?: number | undefined;
             pendingCount?: number | undefined;
-            disabled?: boolean | undefined;
+            dimension: string;
+            value: string;
             totalSpendNanos: number;
             totalRequests: number;
             totalTokens: number;
             dayStamp: string;
-            name: string;
         }[]>;
         setLimits: (ctx: RunMutationCtx, args: {
             name: string;
@@ -267,7 +390,8 @@ export declare class AIBudget {
             dailyTokenLimit?: number;
             lifetimeTokenLimit?: number;
             enforcement?: "hard" | "soft";
-            disabled?: boolean;
+            /** Hard-block this action (was `disabled`). */
+            blocked?: boolean;
         }) => Promise<null>;
         bump: (ctx: RunMutationCtx, args: {
             name: string;

package/dist/client/index.js

@@ -110,6 +110,7 @@ export class AIBudget {
         const started = await ctx.runMutation(this.component.lib.startRequest, {
             userId,
             actionName,
+            tags: args.tags,
             model,
             messages,
             rerunOf: args.rerunOf,
@@ -123,7 +124,13 @@ export class AIBudget {
         }
         const requestId = started.requestId;
         const warnings = started.warnings;
-        await this.fireSoftLimit({ userId, action: actionName, requestId, warnings });
+        await this.fireSoftLimit({
+            userId,
+            action: actionName,
+            tags: args.tags,
+            requestId,
+            warnings,
+        });
         const start = Date.now();
         try {
             // The full chain (incl. system) is stored on the request for audit/replay,
@@ -172,6 +179,7 @@ export class AIBudget {
             const started = await ctx.runMutation(component.lib.startRequest, {
                 userId,
                 actionName,
+                tags: opts.tags,
                 model: modelId,
                 messages: simplifyPrompt(params.prompt),
             });
@@ -185,6 +193,7 @@ export class AIBudget {
             await fireSoftLimit({
                 userId,
                 action: actionName,
+                tags: opts.tags,
                 requestId: started.requestId,
                 warnings: started.warnings,
             });
@@ -278,6 +287,7 @@ export class AIBudget {
             messages: args.messages ?? original.messages,
             rerunOf: args.requestId,
             action: original.actionName,
+            tags: original.tags,
         });
     }
     // ---------- namespaced admin API ----------
@@ -292,25 +302,76 @@ export class AIBudget {
             rerun: (ctx, args) => this.rerunImpl(ctx, args),
         };
     }
-    /** Per-user budgets and controls. */
+    /**
+     * 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", dailySpendLimitNanos });
+     *   ai.tag("customer").list(ctx);
+     *
+     * Attribute a call to it by passing `tags` to `chat`/`languageModel`.
+     */
+    tag(dimension) {
+        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.listUsers, {}),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setLimits, args),
+            list: (ctx) => ctx.runQuery(c.lib.listBuckets, { dimension: "user" }),
+            setLimits: (ctx, args) => {
+                const { userId, ...limits } = args;
+                return ctx.runMutation(c.lib.setBucketLimits, {
+                    dimension: "user",
+                    value: userId,
+                    ...limits,
+                });
+            },
             /** One-time "approve another $X" bump (daily is today-only). */
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpUser, args),
+            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
+                dimension: "user",
+                value: args.userId,
+                dailyNanos: args.dailyNanos,
+                lifetimeNanos: args.lifetimeNanos,
+            }),
             /** Delete a user and all their request rows. */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteUser, args),
+            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, {
+                dimension: "user",
+                value: args.userId,
+            }),
         };
     }
-    /** Per-action (per-feature) budgets. */
+    /** 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),
+            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,
+            }),
         };
     }
     /** The deployment-wide budget and retention config. */

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

@@ -20,22 +20,19 @@ import type { FunctionReference } from "convex/server";
  */
 export type ComponentApi<Name extends string | undefined = string | undefined> = {
     lib: {
-        bumpAction: FunctionReference<"mutation", "internal", {
+        bumpBucket: FunctionReference<"mutation", "internal", {
             dailyNanos?: number;
+            dimension: string;
             lifetimeNanos?: number;
-            name: string;
+            value: string;
         }, null, Name>;
         bumpGlobal: FunctionReference<"mutation", "internal", {
             dailyNanos?: number;
             lifetimeNanos?: 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;
@@ -51,6 +48,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,37 +69,30 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         lineage: FunctionReference<"query", "internal", {
             requestId: string;
         }, any, Name>;
-        listActions: FunctionReference<"query", "internal", {}, any, Name>;
+        listBuckets: FunctionReference<"query", "internal", {
+            dimension?: string;
+        }, any, Name>;
         listPrices: FunctionReference<"query", "internal", {}, any, Name>;
         listRequests: FunctionReference<"query", "internal", {
             limit?: number;
             userId?: string;
         }, any, Name>;
-        listUsers: FunctionReference<"query", "internal", {}, any, Name>;
-        setActionLimits: FunctionReference<"mutation", "internal", {
+        setBucketLimits: FunctionReference<"mutation", "internal", {
+            blocked?: boolean;
             dailySpendLimitNanos?: number;
             dailyTokenLimit?: number;
-            disabled?: boolean;
+            dimension: string;
             enforcement?: "hard" | "soft";
             lifetimeSpendLimitNanos?: number;
             lifetimeTokenLimit?: number;
-            name: string;
+            requestsPerMinute?: number;
+            value: string;
         }, 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>;
@@ -119,6 +113,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             }>;
             model: string;
             rerunOf?: string;
+            tags?: Array<{
+                dimension: string;
+                value: string;
+            }>;
             userId: string;
         }, {
             allowed: true;