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

package/src/client/index.ts

@@ -41,21 +41,40 @@ export type AIBudgetApi = UseApi<typeof api>;
 /** @deprecated use AIBudgetApi */
 export type AIGatewayApi = AIBudgetApi;
 
-/** Fired when a request is admitted over a *soft* limit. */
-export type SoftLimitInfo = {
+/**
+ * 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 };
+
+/** Common shape for budget-event callbacks. */
+export type BudgetEventInfo = {
   userId: string;
   action?: string;
-  requestId: string;
-  warnings: string[];
+  tags?: Tag[];
+  requestId?: string;
+  /** soft-cap warnings (onSoftLimit) or approaching-cap notices (onThreshold). */
+  messages: string[];
+  /** rejection code/reason (onLimitReached only). */
+  code?: string;
+  reason?: string;
 };
+/** @deprecated use BudgetEventInfo */
+export type SoftLimitInfo = BudgetEventInfo & { warnings: string[] };
 export type AIBudgetOptions = {
   defaultModel?: string;
   /**
-   * Called when a soft limit is exceeded (the request is still allowed). Lets
-   * you surface budget warnings even on the languageModel/Agent path, where
-   * they can't be returned. Errors thrown here are swallowed.
+   * A *soft* limit was exceeded (request still allowed). Lets you surface budget
+   * warnings even on the languageModel/Agent path where they can't be returned.
+   * Errors thrown in any of these callbacks are swallowed.
    */
   onSoftLimit?: (info: SoftLimitInfo) => void | Promise<void>;
+  /** Usage crossed a bucket's warnAtPct threshold (approaching a cap). */
+  onThreshold?: (info: BudgetEventInfo) => void | Promise<void>;
+  /** A *hard* limit blocked the request (fires just before chat/model throws). */
+  onLimitReached?: (info: BudgetEventInfo) => void | Promise<void>;
 };
 
 type RunQueryCtx = {
@@ -114,6 +133,8 @@ export type ChatResult = {
   cachedTokens: number;
   /** Soft-limit warnings raised at admission (empty unless a soft cap was hit). */
   warnings: string[];
+  /** Approaching-cap notices (empty unless a warnAtPct threshold was crossed). */
+  notices: string[];
 };
 
 // ---------- helpers ----------
@@ -134,10 +155,12 @@ function extractUsage(usage: any): {
   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`.
+    // 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?.cachedInputTokens ??
+      usage?.inputTokenDetails?.cacheReadTokens ??
+        usage?.cachedInputTokens ??
         usage?.promptTokensDetails?.cachedTokens ??
         usage?.prompt_tokens_details?.cached_tokens ??
         usage?.cached_tokens
@@ -145,6 +168,20 @@ function extractUsage(usage: any): {
   };
 }
 
+// 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: any): number | undefined {
+  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: any): Message[] {
   if (!Array.isArray(prompt)) return [];
@@ -178,23 +215,71 @@ function extractText(result: any): string {
 
 // ---------- client ----------
 
+/** Limits/controls settable on any budget bucket (user, action, or tag). */
+export type BucketLimits = {
+  requestsPerMinute?: number;
+  maxConcurrent?: number;
+  dailySpendLimitNanos?: number;
+  monthlySpendLimitNanos?: number;
+  lifetimeSpendLimitNanos?: number;
+  dailyTokenLimit?: number;
+  monthlyTokenLimit?: number;
+  lifetimeTokenLimit?: number;
+  /** Fire an approaching-limit alert at this fraction of a cap (e.g. 0.8). */
+  warnAtPct?: number;
+  enforcement?: "hard" | "soft";
+  blocked?: boolean;
+};
+/** One-time bump amounts, added on top of a standing cap. */
+export type BumpArgs = {
+  dailyNanos?: number;
+  monthlyNanos?: number;
+  lifetimeNanos?: number;
+};
+
 export class AIBudget {
   public defaultModel: string;
   private onSoftLimit?: AIBudgetOptions["onSoftLimit"];
+  private onThreshold?: AIBudgetOptions["onThreshold"];
+  private onLimitReached?: AIBudgetOptions["onLimitReached"];
   constructor(
     public component: AIBudgetApi,
     options?: AIBudgetOptions
   ) {
     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.
+  private async fireBudgetEvents(
+    base: Omit<BudgetEventInfo, "messages">,
+    warnings: string[],
+    notices: string[]
+  ) {
+    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 */
+      }
+    }
   }
 
-  private async fireSoftLimit(info: SoftLimitInfo) {
-    if (info.warnings.length === 0 || !this.onSoftLimit) return;
+  private async fireLimitReached(info: BudgetEventInfo) {
+    if (!this.onLimitReached) return;
     try {
-      await this.onSoftLimit(info);
+      await this.onLimitReached(info);
     } catch {
-      // never let a callback error break a request
+      /* swallow */
     }
   }
 
@@ -213,6 +298,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,11 +310,20 @@ export class AIBudget {
     const started = await ctx.runMutation(this.component.lib.startRequest, {
       userId,
       actionName,
+      tags: args.tags,
       model,
       messages,
       rerunOf: args.rerunOf as any,
     });
     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,
@@ -236,7 +332,12 @@ 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,
@@ -259,10 +360,11 @@ 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, {
         requestId,
@@ -281,34 +383,48 @@ 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;
-    const fireSoftLimit = this.fireSoftLimit.bind(this);
+    const fireBudgetEvents = this.fireBudgetEvents.bind(this);
+    const fireLimitReached = this.fireLimitReached.bind(this);
 
     const begin = async (params: any) => {
       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 (
@@ -318,6 +434,8 @@ export class AIBudget {
         error?: string;
         promptTokens?: number;
         completionTokens?: number;
+        cachedTokens?: number;
+        costNanos?: number;
         latencyMs?: number;
       }
     ) => ctx.runMutation(component.lib.finishRequest, { requestId, ...fields });
@@ -333,6 +451,7 @@ export class AIBudget {
             await finish(requestId, {
               responseText: extractText(result),
               ...extractUsage(result.usage),
+              costNanos: extractGatewayCostNanos(result),
               latencyMs: Date.now() - start,
             });
             return result;
@@ -409,6 +528,7 @@ export class AIBudget {
       messages: args.messages ?? original.messages,
       rerunOf: args.requestId,
       action: original.actionName,
+      tags: original.tags,
     });
   }
 
@@ -418,8 +538,16 @@ export class AIBudget {
   get requests() {
     const c = this.component;
     return {
-      list: (ctx: RunQueryCtx, args: { userId?: string; limit?: number } = {}) =>
-        ctx.runQuery(c.lib.listRequests, args),
+      /** Filter by userId, or by any {dimension, value} (incl. custom tags). */
+      list: (
+        ctx: RunQueryCtx,
+        args: {
+          userId?: string;
+          dimension?: string;
+          value?: string;
+          limit?: number;
+        } = {}
+      ) => ctx.runQuery(c.lib.listRequests, args),
       /** Ancestors up to the original, plus direct re-runs. */
       lineage: (ctx: RunQueryCtx, args: { requestId: string }) =>
         ctx.runQuery(c.lib.lineage, { requestId: args.requestId as any }),
@@ -431,60 +559,97 @@ export class AIBudget {
     };
   }
 
-  /** 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: string) {
+    return this.dimensionApi(dimension, (a: { value: string }) => a.value);
+  }
+
+  // Shared implementation behind tag()/users/actions. `key` maps the namespace's
+  // id field (value/userId/name) to the bucket value.
+  private dimensionApi<A extends Record<string, any>>(
+    dimension: string,
+    key: (a: A) => string
+  ) {
     const c = this.component;
     return {
-      list: (ctx: RunQueryCtx) => ctx.runQuery(c.lib.listUsers, {}),
-      setLimits: (
-        ctx: RunMutationCtx,
-        args: {
-          userId: string;
-          requestsPerMinute?: number;
-          dailySpendLimitNanos?: number;
-          lifetimeSpendLimitNanos?: number;
-          dailyTokenLimit?: number;
-          lifetimeTokenLimit?: number;
-          enforcement?: "hard" | "soft";
-          blocked?: boolean;
-        }
-      ) => ctx.runMutation(c.lib.setLimits, args),
-      /** One-time "approve another $X" bump (daily is today-only). */
-      bump: (
+      /** 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: A) =>
+        ctx.runQuery(c.lib.getBucket, { dimension, value: key(args) }),
+      setLimits: (ctx: RunMutationCtx, args: A & BucketLimits) => {
+        const { value, userId, name, ...limits } = args as any;
+        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: RunMutationCtx, args: A & BumpArgs) =>
+        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: RunMutationCtx,
-        args: { userId: string; dailyNanos?: number; lifetimeNanos?: number }
-      ) => ctx.runMutation(c.lib.bumpUser, args),
-      /** Delete a user and all their request rows. */
-      delete: (ctx: RunMutationCtx, args: { userId: string }) =>
-        ctx.runMutation(c.lib.deleteUser, args),
+        args: A & { deltaNanos: number; tokens?: number; reason?: string }
+      ) =>
+        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: RunQueryCtx,
+        args: A & { period?: "day" | "month"; limit?: number }
+      ) =>
+        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: RunQueryCtx, args: A & { limit?: number }) =>
+        ctx.runQuery(c.lib.listAdjustments, {
+          dimension,
+          value: key(args),
+          limit: args.limit,
+        }),
+      /** Delete the bucket (for "user", also its request rows). */
+      delete: (ctx: RunMutationCtx, args: A) =>
+        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<{ userId: string }>("user", (a) => a.userId);
+  }
+
+  /** 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, {}),
-      setLimits: (
-        ctx: RunMutationCtx,
-        args: {
-          name: string;
-          dailySpendLimitNanos?: number;
-          lifetimeSpendLimitNanos?: number;
-          dailyTokenLimit?: number;
-          lifetimeTokenLimit?: number;
-          enforcement?: "hard" | "soft";
-          disabled?: boolean;
-        }
-      ) => ctx.runMutation(c.lib.setActionLimits, args),
-      bump: (
-        ctx: RunMutationCtx,
-        args: { name: string; dailyNanos?: number; lifetimeNanos?: number }
-      ) => ctx.runMutation(c.lib.bumpAction, args),
-    };
+    return this.dimensionApi<{ name: string }>("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 {
@@ -503,6 +668,9 @@ export class AIBudget {
         ctx: RunMutationCtx,
         args: { dailyNanos?: number; lifetimeNanos?: number }
       ) => ctx.runMutation(c.lib.bumpGlobal, args),
+      /** Default approaching-limit alert threshold (fraction of a cap, e.g. 0.8). */
+      setAlertDefaults: (ctx: RunMutationCtx, args: { warnAtPct?: number }) =>
+        ctx.runMutation(c.lib.setAlertDefaults, args),
       /** Request-row retention window in ms (default 1h; 0 disables). */
       setRetention: (ctx: RunMutationCtx, args: { retentionMs: number }) =>
         ctx.runMutation(c.lib.setRetention, args),
@@ -533,6 +701,8 @@ export class AIBudget {
           model: string;
           inputNanosPerMTok: number;
           outputNanosPerMTok: number;
+          /** Cache-read rate; defaults to a discount off input if omitted. */
+          cachedNanosPerMTok?: number;
         }
       ) => ctx.runMutation(c.lib.setPrice, args),
     };

package/src/component/_generated/component.ts

@@ -24,31 +24,43 @@ import type { FunctionReference } from "convex/server";
 export type ComponentApi<Name extends string | undefined = string | undefined> =
   {
     lib: {
-      bumpAction: FunctionReference<
+      adjustBucket: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number; name: string },
+        {
+          deltaNanos: number;
+          dimension: string;
+          reason?: string;
+          tokens?: number;
+          value: string;
+        },
         null,
         Name
       >;
-      bumpGlobal: FunctionReference<
+      bumpBucket: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number },
+        {
+          dailyNanos?: number;
+          dimension: string;
+          lifetimeNanos?: number;
+          monthlyNanos?: number;
+          value: string;
+        },
         null,
         Name
       >;
-      bumpUser: FunctionReference<
+      bumpGlobal: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number; userId: string },
+        { dailyNanos?: number; lifetimeNanos?: number; monthlyNanos?: number },
         null,
         Name
       >;
-      deleteUser: FunctionReference<
+      deleteBucket: FunctionReference<
         "mutation",
         "internal",
-        { userId: string },
+        { dimension: string; value: string },
         { deletedThisBatch: number; done: boolean },
         Name
       >;
@@ -58,6 +70,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         {
           cachedTokens?: number;
           completionTokens?: number;
+          costNanos?: number;
           error?: string;
           latencyMs?: number;
           promptTokens?: number;
@@ -67,6 +80,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,54 +121,63 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         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",
-        { limit?: number; userId?: string },
+        { dimension?: string; limit?: number; userId?: string; value?: string },
         any,
         Name
       >;
-      listUsers: FunctionReference<"query", "internal", {}, any, Name>;
-      setActionLimits: FunctionReference<
+      setAlertDefaults: FunctionReference<
         "mutation",
         "internal",
-        {
-          dailySpendLimitNanos?: number;
-          dailyTokenLimit?: number;
-          disabled?: boolean;
-          enforcement?: "hard" | "soft";
-          lifetimeSpendLimitNanos?: number;
-          lifetimeTokenLimit?: number;
-          name: string;
-        },
+        { warnAtPct?: number },
         null,
         Name
       >;
-      setGlobalLimits: FunctionReference<
+      setBucketLimits: FunctionReference<
         "mutation",
         "internal",
         {
+          blocked?: boolean;
           dailySpendLimitNanos?: number;
+          dailyTokenLimit?: number;
+          dimension: string;
           enforcement?: "hard" | "soft";
           lifetimeSpendLimitNanos?: number;
+          lifetimeTokenLimit?: number;
+          maxConcurrent?: number;
+          monthlySpendLimitNanos?: number;
+          monthlyTokenLimit?: number;
+          requestsPerMinute?: number;
+          value: string;
+          warnAtPct?: number;
         },
         null,
         Name
       >;
-      setLimits: FunctionReference<
+      setGlobalLimits: FunctionReference<
         "mutation",
         "internal",
         {
-          blocked?: boolean;
           dailySpendLimitNanos?: number;
-          dailyTokenLimit?: number;
           enforcement?: "hard" | "soft";
           lifetimeSpendLimitNanos?: number;
-          lifetimeTokenLimit?: number;
-          requestsPerMinute?: number;
-          userId: string;
         },
         null,
         Name
@@ -164,6 +193,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         "mutation",
         "internal",
         {
+          cachedNanosPerMTok?: number;
           inputNanosPerMTok: number;
           model: string;
           outputNanosPerMTok: number;
@@ -186,11 +216,29 @@ 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> }
+        | {
+            allowed: true;
+            notices: Array<string>;
+            requestId: string;
+            warnings: Array<string>;
+          }
         | { allowed: false; code: string; reason: string },
         Name
       >;
+      usageHistory: FunctionReference<
+        "query",
+        "internal",
+        {
+          dimension: string;
+          limit?: number;
+          period: "day" | "month";
+          value: string;
+        },
+        any,
+        Name
+      >;
     };
   };