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

package/dist/component/schema.d.ts

@@ -5,8 +5,15 @@ export declare const vMessage: import("convex/values").VObject<{
     role: import("convex/values").VString<string, "required">;
     content: import("convex/values").VString<string, "required">;
 }, "required", "role" | "content">;
+export declare const vTag: import("convex/values").VObject<{
+    dimension: string;
+    value: string;
+}, {
+    dimension: import("convex/values").VString<string, "required">;
+    value: import("convex/values").VString<string, "required">;
+}, "required", "dimension" | "value">;
 declare const _default: import("convex/server").SchemaDefinition<{
-    users: import("convex/server").TableDefinition<import("convex/values").VObject<{
+    buckets: import("convex/server").TableDefinition<import("convex/values").VObject<{
         requestsPerMinute?: number | undefined;
         dailySpendLimitNanos?: number | undefined;
         lifetimeSpendLimitNanos?: number | undefined;
@@ -23,14 +30,16 @@ declare const _default: import("convex/server").SchemaDefinition<{
         reservedTodayTokens?: number | undefined;
         reservedTotalTokens?: number | undefined;
         pendingCount?: number | undefined;
-        userId: string;
+        dimension: string;
+        value: string;
         totalSpendNanos: number;
         totalRequests: number;
         totalTokens: number;
         dayStamp: string;
         spendTodayNanos: number;
     }, {
-        userId: import("convex/values").VString<string, "required">;
+        dimension: import("convex/values").VString<string, "required">;
+        value: import("convex/values").VString<string, "required">;
         requestsPerMinute: import("convex/values").VFloat64<number | undefined, "optional">;
         dailySpendLimitNanos: import("convex/values").VFloat64<number | undefined, "optional">;
         lifetimeSpendLimitNanos: import("convex/values").VFloat64<number | undefined, "optional">;
@@ -52,58 +61,16 @@ declare const _default: import("convex/server").SchemaDefinition<{
         reservedTodayTokens: import("convex/values").VFloat64<number | undefined, "optional">;
         reservedTotalTokens: import("convex/values").VFloat64<number | undefined, "optional">;
         pendingCount: import("convex/values").VFloat64<number | undefined, "optional">;
-    }, "required", "userId" | "requestsPerMinute" | "dailySpendLimitNanos" | "lifetimeSpendLimitNanos" | "dailyTokenLimit" | "lifetimeTokenLimit" | "blocked" | "enforcement" | "dailyBumpNanos" | "lifetimeBumpNanos" | "bumpDayStamp" | "totalSpendNanos" | "totalRequests" | "totalTokens" | "dayStamp" | "spendTodayNanos" | "tokensToday" | "reservedTodayNanos" | "reservedTotalNanos" | "reservedTodayTokens" | "reservedTotalTokens" | "pendingCount">, {
-        userId: ["userId", "_creationTime"];
-    }, {}, {}>;
-    actions: import("convex/server").TableDefinition<import("convex/values").VObject<{
-        dailySpendLimitNanos?: number | undefined;
-        lifetimeSpendLimitNanos?: number | undefined;
-        dailyTokenLimit?: number | undefined;
-        lifetimeTokenLimit?: number | 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;
-        disabled?: boolean | undefined;
-        totalSpendNanos: number;
-        totalRequests: number;
-        totalTokens: number;
-        dayStamp: string;
-        spendTodayNanos: number;
-        name: string;
-    }, {
-        name: import("convex/values").VString<string, "required">;
-        dailySpendLimitNanos: import("convex/values").VFloat64<number | undefined, "optional">;
-        lifetimeSpendLimitNanos: import("convex/values").VFloat64<number | undefined, "optional">;
-        dailyTokenLimit: import("convex/values").VFloat64<number | undefined, "optional">;
-        lifetimeTokenLimit: import("convex/values").VFloat64<number | undefined, "optional">;
-        disabled: import("convex/values").VBoolean<boolean | undefined, "optional">;
-        enforcement: import("convex/values").VUnion<"hard" | "soft" | undefined, [import("convex/values").VLiteral<"hard", "required">, import("convex/values").VLiteral<"soft", "required">], "optional", never>;
-        dailyBumpNanos: import("convex/values").VFloat64<number | undefined, "optional">;
-        lifetimeBumpNanos: import("convex/values").VFloat64<number | undefined, "optional">;
-        bumpDayStamp: import("convex/values").VString<string | undefined, "optional">;
-        totalSpendNanos: import("convex/values").VFloat64<number, "required">;
-        totalRequests: import("convex/values").VFloat64<number, "required">;
-        totalTokens: import("convex/values").VFloat64<number, "required">;
-        dayStamp: import("convex/values").VString<string, "required">;
-        spendTodayNanos: import("convex/values").VFloat64<number, "required">;
-        tokensToday: import("convex/values").VFloat64<number | undefined, "optional">;
-        reservedTodayNanos: import("convex/values").VFloat64<number | undefined, "optional">;
-        reservedTotalNanos: import("convex/values").VFloat64<number | undefined, "optional">;
-        reservedTodayTokens: import("convex/values").VFloat64<number | undefined, "optional">;
-        reservedTotalTokens: import("convex/values").VFloat64<number | undefined, "optional">;
-        pendingCount: import("convex/values").VFloat64<number | undefined, "optional">;
-    }, "required", "dailySpendLimitNanos" | "lifetimeSpendLimitNanos" | "dailyTokenLimit" | "lifetimeTokenLimit" | "enforcement" | "dailyBumpNanos" | "lifetimeBumpNanos" | "bumpDayStamp" | "totalSpendNanos" | "totalRequests" | "totalTokens" | "dayStamp" | "spendTodayNanos" | "tokensToday" | "reservedTodayNanos" | "reservedTotalNanos" | "reservedTodayTokens" | "reservedTotalTokens" | "pendingCount" | "name" | "disabled">, {
-        name: ["name", "_creationTime"];
+    }, "required", "dimension" | "value" | "requestsPerMinute" | "dailySpendLimitNanos" | "lifetimeSpendLimitNanos" | "dailyTokenLimit" | "lifetimeTokenLimit" | "blocked" | "enforcement" | "dailyBumpNanos" | "lifetimeBumpNanos" | "bumpDayStamp" | "totalSpendNanos" | "totalRequests" | "totalTokens" | "dayStamp" | "spendTodayNanos" | "tokensToday" | "reservedTodayNanos" | "reservedTotalNanos" | "reservedTodayTokens" | "reservedTotalTokens" | "pendingCount">, {
+        dim_value: ["dimension", "value", "_creationTime"];
+        dimension: ["dimension", "_creationTime"];
     }, {}, {}>;
     requests: import("convex/server").TableDefinition<import("convex/values").VObject<{
         actionName?: string | undefined;
+        tags?: {
+            dimension: string;
+            value: string;
+        }[] | undefined;
         estimatedNanos?: number | undefined;
         estimatedTokens?: number | undefined;
         unpricedModel?: boolean | undefined;
@@ -127,6 +94,16 @@ declare const _default: import("convex/server").SchemaDefinition<{
     }, {
         userId: import("convex/values").VString<string, "required">;
         actionName: import("convex/values").VString<string | undefined, "optional">;
+        tags: import("convex/values").VArray<{
+            dimension: string;
+            value: string;
+        }[] | undefined, import("convex/values").VObject<{
+            dimension: string;
+            value: string;
+        }, {
+            dimension: import("convex/values").VString<string, "required">;
+            value: import("convex/values").VString<string, "required">;
+        }, "required", "dimension" | "value">, "optional">;
         model: import("convex/values").VString<string, "required">;
         estimatedNanos: import("convex/values").VFloat64<number | undefined, "optional">;
         estimatedTokens: import("convex/values").VFloat64<number | undefined, "optional">;
@@ -152,7 +129,7 @@ declare const _default: import("convex/server").SchemaDefinition<{
         costNanos: import("convex/values").VFloat64<number | undefined, "optional">;
         latencyMs: import("convex/values").VFloat64<number | undefined, "optional">;
         rerunOf: import("convex/values").VId<import("convex/values").GenericId<"requests"> | undefined, "optional">;
-    }, "required", "userId" | "actionName" | "model" | "estimatedNanos" | "estimatedTokens" | "unpricedModel" | "overBudget" | "settled" | "messages" | "status" | "error" | "responseText" | "promptTokens" | "completionTokens" | "cachedTokens" | "costNanos" | "latencyMs" | "rerunOf">, {
+    }, "required", "userId" | "actionName" | "tags" | "model" | "estimatedNanos" | "estimatedTokens" | "unpricedModel" | "overBudget" | "settled" | "messages" | "status" | "error" | "responseText" | "promptTokens" | "completionTokens" | "cachedTokens" | "costNanos" | "latencyMs" | "rerunOf">, {
         userId: ["userId", "_creationTime"];
         status: ["status", "_creationTime"];
         rerunOf: ["rerunOf", "_creationTime"];

package/dist/component/schema.js

@@ -4,20 +4,28 @@ export const vMessage = v.object({
     role: v.string(), // "system" | "user" | "assistant" | "tool"
     content: v.string(),
 });
+// One attribution tag: a (dimension, value) pair, e.g. {dimension:"user",
+// value:"alice"} or {dimension:"customer", value:"acme"}. `user` and `action`
+// are built-in dimensions; apps can add any others (team, project, env, …).
+export const vTag = v.object({ dimension: v.string(), value: v.string() });
 export default defineSchema({
-    users: defineTable({
-        userId: v.string(), // app-provided stable key (your user id)
+    // A budget holder, keyed by (dimension, value). Unifies what used to be the
+    // `users` and `actions` tables — those are just the "user" and "action"
+    // dimensions now. Any tag a request carries can have its own budget here.
+    buckets: defineTable({
+        dimension: v.string(),
+        value: v.string(),
         // limits (all optional — unlimited by default)
-        requestsPerMinute: v.optional(v.number()),
+        requestsPerMinute: v.optional(v.number()), // enforced on the "user" dimension
         dailySpendLimitNanos: v.optional(v.number()),
         lifetimeSpendLimitNanos: v.optional(v.number()),
         dailyTokenLimit: v.optional(v.number()),
         lifetimeTokenLimit: v.optional(v.number()),
-        blocked: v.optional(v.boolean()),
+        blocked: v.optional(v.boolean()), // hard block (was `blocked`/`disabled`)
         // "hard" (default): exceeding a budget blocks. "soft": warn but allow.
         enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
-        // one-time bumps ("approve another $X") added on top of the cap. Daily bump
-        // is scoped to bumpDayStamp (resets with the day); lifetime bump is permanent.
+        // one-time bumps ("approve another $X"). Daily bump is scoped to bumpDayStamp
+        // (resets with the day); lifetime bump is permanent.
         dailyBumpNanos: v.optional(v.number()),
         lifetimeBumpNanos: v.optional(v.number()),
         bumpDayStamp: v.optional(v.string()),
@@ -26,43 +34,24 @@ export default defineSchema({
         totalRequests: v.number(),
         totalTokens: v.number(),
         // daily window
-        dayStamp: v.string(), // e.g. "2026-08-27" (UTC)
-        spendTodayNanos: v.number(),
-        tokensToday: v.optional(v.number()),
-        // in-flight reservations (pessimistic holds; released on settle/expiry)
-        reservedTodayNanos: v.optional(v.number()),
-        reservedTotalNanos: v.optional(v.number()),
-        reservedTodayTokens: v.optional(v.number()),
-        reservedTotalTokens: v.optional(v.number()),
-        pendingCount: v.optional(v.number()),
-    }).index("userId", ["userId"]),
-    // per-action-name budgets and running totals (e.g. "chat", "summarize")
-    actions: defineTable({
-        name: v.string(),
-        dailySpendLimitNanos: v.optional(v.number()),
-        lifetimeSpendLimitNanos: v.optional(v.number()),
-        dailyTokenLimit: v.optional(v.number()),
-        lifetimeTokenLimit: v.optional(v.number()),
-        disabled: v.optional(v.boolean()),
-        enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
-        dailyBumpNanos: v.optional(v.number()),
-        lifetimeBumpNanos: v.optional(v.number()),
-        bumpDayStamp: v.optional(v.string()),
-        totalSpendNanos: v.number(),
-        totalRequests: v.number(),
-        totalTokens: v.number(),
         dayStamp: v.string(),
         spendTodayNanos: v.number(),
         tokensToday: v.optional(v.number()),
+        // in-flight reservations (pessimistic holds; released on settle/expiry)
         reservedTodayNanos: v.optional(v.number()),
         reservedTotalNanos: v.optional(v.number()),
         reservedTodayTokens: v.optional(v.number()),
         reservedTotalTokens: v.optional(v.number()),
         pendingCount: v.optional(v.number()),
-    }).index("name", ["name"]),
+    })
+        .index("dim_value", ["dimension", "value"])
+        .index("dimension", ["dimension"]),
     requests: defineTable({
+        // `user` and `action` stay first-class + indexed (the hot-path filters and
+        // rate limiting); the full attribution incl. extra tags lives in `tags`.
         userId: v.string(),
         actionName: v.optional(v.string()),
+        tags: v.optional(v.array(vTag)),
         model: v.string(),
         // pessimistic holds placed at start; reconciled to actual on settle
         estimatedNanos: v.optional(v.number()),
@@ -105,12 +94,12 @@ export default defineSchema({
         // "denylist": any model except the listed ones.
         modelMode: v.optional(v.union(v.literal("open"), v.literal("allowlist"), v.literal("denylist"))),
         models: v.optional(v.array(v.string())),
-        // Deployment-wide ("global") spend cap across ALL users and actions. The
-        // running totals live in a sharded counter (high write throughput); only
-        // the limit config lives here. The cap is enforced approximately — the
-        // sharded total is read without a reservation, so under heavy concurrency
-        // it can overshoot by a bounded amount. Right for a global killswitch;
-        // per-user/per-action caps stay exact via reserve/settle.
+        // Deployment-wide ("global") spend cap across ALL requests. Running totals
+        // live in a sharded counter (high write throughput) since every request
+        // touches it; only the limit config lives here. Enforced approximately —
+        // the sharded total is read without a reservation, so under heavy
+        // concurrency it can overshoot by a bounded amount. Right for a global
+        // killswitch; per-bucket caps stay exact via reserve/settle.
         globalDailySpendLimitNanos: v.optional(v.number()),
         globalLifetimeSpendLimitNanos: v.optional(v.number()),
         globalEnforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),

package/package.json

@@ -7,7 +7,7 @@
     "email": "support@convex.dev",
     "url": "https://github.com/get-convex/ai-budget/issues"
   },
-  "version": "0.0.2-alpha.0",
+  "version": "0.0.2-alpha.4",
   "license": "Apache-2.0",
   "type": "module",
   "keywords": [
@@ -64,10 +64,10 @@
     "convex": "^1.45.0"
   },
   "devDependencies": {
-    "@edge-runtime/vm": "^5.0.0",
-    "convex": "^1.45.0",
-    "convex-test": "^0.0.38",
+    "@edge-runtime/vm": "5.0.0",
+    "convex": "1.45.0",
+    "convex-test": "0.0.56",
     "typescript": "^5.7.0",
-    "vitest": "^4.0.0"
+    "vitest": "4.1.11"
   }
 }

package/src/client/index.ts

@@ -41,10 +41,19 @@ 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[];
 };
@@ -213,6 +222,8 @@ export 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> {
     const model = args.model ?? this.defaultModel;
@@ -223,6 +234,7 @@ export class AIBudget {
     const started = await ctx.runMutation(this.component.lib.startRequest, {
       userId,
       actionName,
+      tags: args.tags,
       model,
       messages,
       rerunOf: args.rerunOf as any,
@@ -236,7 +248,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,
@@ -281,7 +299,13 @@ export class AIBudget {
    */
   languageModel(
     ctx: RunMutationCtx,
-    opts: { userId?: string; model?: string; action?: string } = {}
+    opts: {
+      userId?: string;
+      model?: string;
+      action?: string;
+      /** Extra attribution dimensions to bill/limit (team, customer, env, …). */
+      tags?: Tag[];
+    } = {}
   ): LanguageModel {
     const modelId = opts.model ?? this.defaultModel;
     const component = this.component;
@@ -293,6 +317,7 @@ export class AIBudget {
       const started = await ctx.runMutation(component.lib.startRequest, {
         userId,
         actionName,
+        tags: opts.tags,
         model: modelId,
         messages: simplifyPrompt(params.prompt),
       });
@@ -306,6 +331,7 @@ export class AIBudget {
       await fireSoftLimit({
         userId,
         action: actionName,
+        tags: opts.tags,
         requestId: started.requestId,
         warnings: started.warnings,
       });
@@ -409,6 +435,7 @@ export class AIBudget {
       messages: args.messages ?? original.messages,
       rerunOf: args.requestId,
       action: original.actionName,
+      tags: original.tags,
     });
   }
 
@@ -431,11 +458,55 @@ export class AIBudget {
     };
   }
 
-  /** 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) {
+    const c = this.component;
+    return {
+      /** All buckets in this dimension. */
+      list: (ctx: RunQueryCtx) =>
+        ctx.runQuery(c.lib.listBuckets, { dimension }),
+      /** One bucket's limits + spend (null if it has none yet). */
+      get: (ctx: RunQueryCtx, args: { value: string }) =>
+        ctx.runQuery(c.lib.getBucket, { dimension, value: args.value }),
+      setLimits: (
+        ctx: RunMutationCtx,
+        args: {
+          value: string;
+          requestsPerMinute?: number;
+          dailySpendLimitNanos?: number;
+          lifetimeSpendLimitNanos?: number;
+          dailyTokenLimit?: number;
+          lifetimeTokenLimit?: number;
+          enforcement?: "hard" | "soft";
+          blocked?: boolean;
+        }
+      ) => ctx.runMutation(c.lib.setBucketLimits, { dimension, ...args }),
+      /** One-time "approve another $X" bump (daily is today-only). */
+      bump: (
+        ctx: RunMutationCtx,
+        args: { value: string; dailyNanos?: number; lifetimeNanos?: number }
+      ) => ctx.runMutation(c.lib.bumpBucket, { dimension, ...args }),
+      /** Delete the bucket (for "user", also its request rows). */
+      delete: (ctx: RunMutationCtx, args: { value: string }) =>
+        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: RunQueryCtx) => ctx.runQuery(c.lib.listUsers, {}),
+      list: (ctx: RunQueryCtx) =>
+        ctx.runQuery(c.lib.listBuckets, { dimension: "user" }),
       setLimits: (
         ctx: RunMutationCtx,
         args: {
@@ -448,23 +519,40 @@ export class AIBudget {
           enforcement?: "hard" | "soft";
           blocked?: boolean;
         }
-      ) => ctx.runMutation(c.lib.setLimits, 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: RunMutationCtx,
         args: { userId: string; dailyNanos?: number; lifetimeNanos?: number }
-      ) => ctx.runMutation(c.lib.bumpUser, 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: RunMutationCtx, args: { userId: string }) =>
-        ctx.runMutation(c.lib.deleteUser, 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: RunQueryCtx) => ctx.runQuery(c.lib.listActions, {}),
+      list: (ctx: RunQueryCtx) =>
+        ctx.runQuery(c.lib.listBuckets, { dimension: "action" }),
       setLimits: (
         ctx: RunMutationCtx,
         args: {
@@ -474,13 +562,27 @@ export class AIBudget {
           dailyTokenLimit?: number;
           lifetimeTokenLimit?: number;
           enforcement?: "hard" | "soft";
-          disabled?: boolean;
+          /** Hard-block this action (was `disabled`). */
+          blocked?: boolean;
         }
-      ) => ctx.runMutation(c.lib.setActionLimits, args),
+      ) => {
+        const { name, ...limits } = args;
+        return ctx.runMutation(c.lib.setBucketLimits, {
+          dimension: "action",
+          value: name,
+          ...limits,
+        });
+      },
       bump: (
         ctx: RunMutationCtx,
         args: { name: string; dailyNanos?: number; lifetimeNanos?: number }
-      ) => ctx.runMutation(c.lib.bumpAction, args),
+      ) =>
+        ctx.runMutation(c.lib.bumpBucket, {
+          dimension: "action",
+          value: args.name,
+          dailyNanos: args.dailyNanos,
+          lifetimeNanos: args.lifetimeNanos,
+        }),
     };
   }
 

package/src/component/_generated/component.ts

@@ -24,10 +24,15 @@ import type { FunctionReference } from "convex/server";
 export type ComponentApi<Name extends string | undefined = string | undefined> =
   {
     lib: {
-      bumpAction: FunctionReference<
+      bumpBucket: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number; name: string },
+        {
+          dailyNanos?: number;
+          dimension: string;
+          lifetimeNanos?: number;
+          value: string;
+        },
         null,
         Name
       >;
@@ -38,17 +43,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         null,
         Name
       >;
-      bumpUser: FunctionReference<
+      deleteBucket: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number; userId: string },
-        null,
-        Name
-      >;
-      deleteUser: FunctionReference<
-        "mutation",
-        "internal",
-        { userId: string },
+        { dimension: string; value: string },
         { deletedThisBatch: number; done: boolean },
         Name
       >;
@@ -67,6 +65,13 @@ 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",
@@ -101,7 +106,13 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         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",
@@ -110,18 +121,19 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         any,
         Name
       >;
-      listUsers: FunctionReference<"query", "internal", {}, any, Name>;
-      setActionLimits: FunctionReference<
+      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
@@ -137,22 +149,6 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         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",
@@ -186,6 +182,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           messages: Array<{ content: string; role: string }>;
           model: string;
           rerunOf?: string;
+          tags?: Array<{ dimension: string; value: string }>;
           userId: string;
         },
         | { allowed: true; requestId: string; warnings: Array<string> }