@convex-dev/ai-budget 0.0.2-alpha.10 → 0.0.2-alpha.13

package/src/client/index.ts

@@ -306,97 +306,233 @@ export class AIBudget {
   }
 
   /**
-   * One-shot chat through the AI Gateway with tracking + limits.
-   * Call from an action. `userId` defaults to the authenticated caller.
+   * Reserve budget for a call WITHOUT running it — the async half of the
+   * lifecycle. Use for long-running jobs (video generation) where the result
+   * arrives minutes later via a poll or webhook: `begin` here, then `settle`
+   * from that later context with the request's `requestId`. Returns the
+   * admission result (does NOT throw over a cap — check `allowed`). Set
+   * `reserveTtlMs` to the job's max duration so the hold isn't reaped mid-flight.
    */
-  async chat(
+  async begin(
     ctx: RunMutationCtx,
-    args: {
-      /** Whom to bill. Defaults to the authenticated user (ctx.auth). */
-      userId?: string;
-      prompt?: string;
+    opts: {
+      model: string;
       messages?: Message[];
-      model?: string;
-      rerunOf?: string;
-      /** Attribute spend to this action name. Defaults to the calling Convex action. */
+      userId?: string;
       action?: string;
-      /** Extra attribution dimensions to bill/limit (team, customer, env, …). */
       tags?: Tag[];
-    } = {}
-  ): Promise<ChatResult> {
-    const model = args.model ?? this.defaultModel;
-    const userId = await resolveUserId(ctx, args.userId);
-    const actionName = await resolveActionName(ctx, args.action);
-    const messages: Message[] =
-      args.messages ?? [{ role: "user", content: args.prompt ?? "" }];
+      /** Reserve this exact amount (nanodollars) when the cost is known up front. */
+      estimatedCostNanos?: number;
+      /** Hold the reservation up to this long (ms) for long async jobs. */
+      reserveTtlMs?: number;
+      rerunOf?: string;
+    }
+  ): Promise<
+    | { allowed: true; requestId: string; warnings: string[]; notices: string[] }
+    | { allowed: false; code: string; reason: string }
+  > {
+    const userId = await resolveUserId(ctx, opts.userId);
+    const actionName = await resolveActionName(ctx, opts.action);
     const started = await ctx.runMutation(this.component.lib.startRequest, {
       userId,
       actionName,
-      tags: args.tags,
-      model,
-      messages,
-      rerunOf: args.rerunOf as any,
+      tags: opts.tags,
+      model: opts.model,
+      messages: opts.messages ?? [],
+      estimatedCostNanos: opts.estimatedCostNanos,
+      reserveTtlMs: opts.reserveTtlMs,
+      rerunOf: opts.rerunOf as any,
     });
-    if (!started.allowed) {
+    if (started.allowed) {
+      await this.fireBudgetEvents(
+        { userId, action: actionName, tags: opts.tags, requestId: started.requestId },
+        started.warnings,
+        started.notices
+      );
+    } else {
       await this.fireLimitReached({
         userId,
         action: actionName,
-        tags: args.tags,
+        tags: opts.tags,
         messages: [started.reason],
         code: started.code,
         reason: started.reason,
       });
+    }
+    return started;
+  }
+
+  /**
+   * Record the actual usage/cost of a `begin`-reserved request and release its
+   * reservation. Idempotent (exactly-once server-side). Pass a raw provider
+   * `usage` (auto-normalized) or explicit token counts, plus optional
+   * `serverToolUses` and an authoritative `costNanos`.
+   */
+  async settle(
+    ctx: RunMutationCtx,
+    args: {
+      requestId: string;
+      responseText?: string;
+      error?: string;
+      usage?: any;
+      promptTokens?: number;
+      completionTokens?: number;
+      cachedTokens?: number;
+      serverToolUses?: Record<string, number>;
+      costNanos?: number;
+      latencyMs?: number;
+    }
+  ): Promise<{ costNanos: number }> {
+    const { requestId, usage, promptTokens, completionTokens, cachedTokens, ...rest } = args;
+    const tokens =
+      promptTokens !== undefined ||
+      completionTokens !== undefined ||
+      cachedTokens !== undefined
+        ? {
+            promptTokens: promptTokens ?? 0,
+            completionTokens: completionTokens ?? 0,
+            cachedTokens: cachedTokens ?? 0,
+          }
+        : usage !== undefined
+          ? extractUsage(usage)
+          : {};
+    return ctx.runMutation(this.component.lib.finishRequest, {
+      requestId: requestId as any,
+      ...tokens,
+      ...rest,
+    });
+  }
+
+  /**
+   * Meter ANY synchronous LLM call — gateway, a provider SDK, a raw fetch — with
+   * the same budgets, audit log, and cost tracking. Reserves before your `run`
+   * (throwing a ConvexError over a hard cap), runs it, then records the actual
+   * usage/cost. The provider-agnostic core; `chat` is sugar over it. (For async
+   * jobs that settle later, use `begin`/`settle` instead.)
+   *
+   * `run` returns what happened. Pass a raw provider `usage` object (auto-
+   * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
+   * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`) and an
+   * authoritative `costNanos` (used verbatim if present).
+   */
+  async meter(
+    ctx: RunMutationCtx,
+    opts: {
+      model: string;
+      messages: Message[];
+      userId?: string;
+      action?: string;
+      tags?: Tag[];
+      rerunOf?: string;
+      /** Reserve this exact amount (nanodollars) instead of the token estimate. */
+      estimatedCostNanos?: number;
+    },
+    run: () => Promise<{
+      text?: string;
+      usage?: any;
+      promptTokens?: number;
+      completionTokens?: number;
+      cachedTokens?: number;
+      serverToolUses?: Record<string, number>;
+      costNanos?: number;
+    }>
+  ): Promise<ChatResult> {
+    const started = await this.begin(ctx, opts);
+    if (!started.allowed) {
       throw new ConvexError({
         kind: "AIBudgetLimit",
         code: started.code,
         reason: started.reason,
       });
     }
-    const requestId = started.requestId;
-    const warnings = started.warnings;
-    const notices = started.notices;
-    await this.fireBudgetEvents(
-      { userId, action: actionName, tags: args.tags, requestId },
-      warnings,
-      notices
-    );
+    const { requestId, warnings, notices } = started;
     const start = Date.now();
     try {
-      // The full chain (incl. system) is stored on the request for audit/replay,
-      // but the AI SDK wants system prompts in the `system` option, not messages.
-      const system =
-        messages
-          .filter((m) => m.role === "system")
-          .map((m) => m.content)
-          .join("\n\n") || undefined;
-      const convo = messages.filter((m) => m.role !== "system");
-      const result = await generateText({
-        model: convexGateway(model),
-        ...(system ? { system } : {}),
-        messages: convo as any,
-      });
-      const usage = extractUsage(result.usage);
-      const { costNanos } = await ctx.runMutation(
-        this.component.lib.finishRequest,
-        {
-          requestId,
-          responseText: result.text,
-          ...usage,
-          costNanos: extractGatewayCostNanos(result),
-          latencyMs: Date.now() - start,
-        }
-      );
-      return { text: result.text, requestId, costNanos, warnings, notices, ...usage };
-    } catch (e) {
-      await ctx.runMutation(this.component.lib.finishRequest, {
+      const out = await run();
+      const { costNanos } = await this.settle(ctx, {
         requestId,
-        error: String(e),
+        responseText: out.text,
+        usage: out.usage,
+        promptTokens: out.promptTokens,
+        completionTokens: out.completionTokens,
+        cachedTokens: out.cachedTokens,
+        serverToolUses: out.serverToolUses,
+        costNanos: out.costNanos,
         latencyMs: Date.now() - start,
       });
+      // Re-derive the recorded usage for the return value.
+      const usage =
+        out.promptTokens !== undefined ||
+        out.completionTokens !== undefined ||
+        out.cachedTokens !== undefined
+          ? {
+              promptTokens: out.promptTokens ?? 0,
+              completionTokens: out.completionTokens ?? 0,
+              cachedTokens: out.cachedTokens ?? 0,
+            }
+          : extractUsage(out.usage);
+      return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
+    } catch (e) {
+      await this.settle(ctx, { requestId, error: String(e), latencyMs: Date.now() - start });
       throw e;
     }
   }
 
+  /**
+   * One-shot chat through the AI Gateway with tracking + limits — sugar over
+   * `meter`. Call from an action. `userId` defaults to the authenticated caller.
+   */
+  async chat(
+    ctx: RunMutationCtx,
+    args: {
+      /** Whom to bill. Defaults to the authenticated user (ctx.auth). */
+      userId?: string;
+      prompt?: string;
+      messages?: Message[];
+      model?: string;
+      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;
+    const messages: Message[] =
+      args.messages ?? [{ role: "user", content: args.prompt ?? "" }];
+    return this.meter(
+      ctx,
+      {
+        model,
+        messages,
+        userId: args.userId,
+        action: args.action,
+        tags: args.tags,
+        rerunOf: args.rerunOf,
+      },
+      async () => {
+        // The full chain (incl. system) is stored for audit/replay, but the AI
+        // SDK wants system prompts in the `system` option, not messages.
+        const system =
+          messages
+            .filter((m) => m.role === "system")
+            .map((m) => m.content)
+            .join("\n\n") || undefined;
+        const convo = messages.filter((m) => m.role !== "system");
+        const result = await generateText({
+          model: convexGateway(model),
+          ...(system ? { system } : {}),
+          messages: convo as any,
+        });
+        return {
+          text: result.text,
+          usage: result.usage,
+          costNanos: extractGatewayCostNanos(result),
+        };
+      }
+    );
+  }
+
   /**
    * An AI SDK LanguageModel that enforces limits and records usage/cost for
    * `userId` on every call. Drop it into `generateText`, `streamText`, or the
@@ -720,7 +856,7 @@ export class AIBudget {
     };
   }
 
-  /** Per-model prices (cents per million tokens). */
+  /** Per-model prices (nanodollars per million tokens) + server-tool fees. */
   get prices() {
     const c = this.component;
     return {
@@ -735,6 +871,14 @@ export class AIBudget {
           cachedNanosPerMTok?: number;
         }
       ) => ctx.runMutation(c.lib.setPrice, args),
+      /** Per-call fees for provider server tools (web search, etc.). */
+      listServerTools: (ctx: RunQueryCtx) =>
+        ctx.runQuery(c.lib.listServerToolPrices, {}),
+      /** Set a server-tool's per-call price, e.g. { tool: "web_search", nanosPerCall }. */
+      setServerTool: (
+        ctx: RunMutationCtx,
+        args: { tool: string; nanosPerCall: number }
+      ) => ctx.runMutation(c.lib.setServerToolPrice, args),
     };
   }
 
@@ -871,6 +1015,60 @@ export class AIBudget {
     http.route({ pathPrefix: `${prefix}/`, method: "GET", handler });
     http.route({ pathPrefix: `${prefix}/`, method: "POST", handler });
   }
+
+  /**
+   * Mount a POST webhook that settles a `begin`-reserved request from a
+   * provider's completion callback (async image/video jobs). You supply
+   * `resolve` — verify the payload's signature and map the provider's job id to
+   * your stored `requestId` + final usage/cost; the helper calls `settle`.
+   * Return `null` to ignore an unrecognized/duplicate callback (HTTP 202).
+   *
+   *   ai.registerWebhook(http, {
+   *     path: "/aibudget/video-done",
+   *     resolve: async (ctx, request, body) => {
+   *       if (!verifySignature(request, body)) return null;
+   *       const job = await lookupJob(ctx, body.id);   // your table: { requestId }
+   *       return { requestId: job.requestId, serverToolUses: { video_seconds: body.seconds } };
+   *     },
+   *   });
+   */
+  registerWebhook(
+    http: HttpRouter,
+    opts: {
+      path?: string;
+      resolve: (
+        ctx: any,
+        request: Request,
+        body: any
+      ) => Promise<
+        | ({ requestId: string } & {
+            responseText?: string;
+            error?: string;
+            usage?: any;
+            promptTokens?: number;
+            completionTokens?: number;
+            cachedTokens?: number;
+            serverToolUses?: Record<string, number>;
+            costNanos?: number;
+          })
+        | null
+      >;
+    }
+  ) {
+    const path = opts.path ?? "/aibudget/webhook";
+    const self = this;
+    http.route({
+      path,
+      method: "POST",
+      handler: httpActionGeneric(async (ctx: any, request: Request) => {
+        const body = await request.json().catch(() => ({}));
+        const settle = await opts.resolve(ctx, request, body);
+        if (!settle) return new Response("ignored", { status: 202 });
+        await self.settle(ctx, settle);
+        return new Response("ok");
+      }),
+    });
+  }
 }
 
 /** @deprecated Renamed to `AIBudget`. */

package/src/component/_generated/component.ts

@@ -76,6 +76,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           promptTokens?: number;
           requestId: string;
           responseText?: string;
+          serverToolUses?: Record<string, number>;
         },
         { costNanos: number },
         Name
@@ -145,6 +146,13 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         any,
         Name
       >;
+      listServerToolPrices: FunctionReference<
+        "query",
+        "internal",
+        {},
+        any,
+        Name
+      >;
       setAlertDefaults: FunctionReference<
         "mutation",
         "internal",
@@ -210,14 +218,23 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         null,
         Name
       >;
+      setServerToolPrice: FunctionReference<
+        "mutation",
+        "internal",
+        { nanosPerCall: number; tool: string },
+        null,
+        Name
+      >;
       startRequest: FunctionReference<
         "mutation",
         "internal",
         {
           actionName?: string;
+          estimatedCostNanos?: number;
           messages: Array<{ content: string; role: string }>;
           model: string;
           rerunOf?: string;
+          reserveTtlMs?: number;
           tags?: Array<{ dimension: string; value: string }>;
           userId: string;
         },

package/src/component/lib.test.ts

@@ -116,6 +116,109 @@ describe("cache-aware pricing", () => {
   });
 });
 
+describe("server-tool pricing", () => {
+  test("server-tool uses add a per-call fee on top of tokens", async () => {
+    const t = convexTest(schema, modules);
+    const r = await start(t, { userId: "u" });
+    // 0 tokens; 3 web searches at the $0.01 default = 30_000_000 nano.
+    await settleWith(t, r.requestId, {
+      promptTokens: 0,
+      completionTokens: 0,
+      serverToolUses: { web_search: 3 },
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(30_000_000);
+    expect(req.serverToolUses).toEqual({ web_search: 3 });
+  });
+
+  test("an override price is applied", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setServerToolPrice, {
+      tool: "web_search",
+      nanosPerCall: 12_000_000,
+    });
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, {
+      promptTokens: 0,
+      completionTokens: 0,
+      serverToolUses: { web_search: 2 },
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(24_000_000);
+  });
+
+  test("an authoritative cost already includes tool fees (not double-charged)", async () => {
+    const t = convexTest(schema, modules);
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, {
+      promptTokens: 1_000_000,
+      completionTokens: 0,
+      serverToolUses: { web_search: 5 },
+      costNanos: 999,
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(999);
+  });
+});
+
+describe("cost known up front (image gen, per-call APIs)", () => {
+  test("estimatedCostNanos drives the reservation for a hard cap", async () => {
+    const t = convexTest(schema, modules);
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 100_000_000 }); // $0.10
+    // A $0.13 image is known before the call; reserving it exceeds the cap,
+    // even though the token estimate for the prompt alone would pass.
+    const r = await start(t, {
+      userId: "u",
+      model: "openai/gpt-image-1",
+      estimatedCostNanos: 130_000_000,
+    });
+    expect(r.allowed).toBe(false);
+    expect(r.code).toBe("user_daily_spend_limit");
+  });
+
+  test("admits when it fits, then settles to the real per-image cost", async () => {
+    const t = convexTest(schema, modules);
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 500_000_000 });
+    const r = await start(t, {
+      userId: "u",
+      model: "openai/gpt-image-1",
+      estimatedCostNanos: 130_000_000,
+    });
+    expect(r.allowed).toBe(true);
+    await settleWith(t, r.requestId, { costNanos: 130_000_000 });
+    const u = await userOf(t, "u");
+    expect(u.totalSpendNanos).toBe(130_000_000);
+    expect(u.reservedTotalNanos ?? 0).toBe(0);
+  });
+});
+
+describe("async lifecycle (video jobs): begin now, settle later", () => {
+  test("reserveTtlMs is stored, and settle records the real cost", async () => {
+    const t = convexTest(schema, modules);
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 5_000_000_000 });
+    // Reserve $2 for a long job that will settle minutes later.
+    const r = await start(t, {
+      userId: "u",
+      model: "openai/sora",
+      estimatedCostNanos: 2_000_000_000,
+      reserveTtlMs: 30 * 60 * 1000,
+    });
+    expect(r.allowed).toBe(true);
+    const pending = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(pending.status).toBe("pending");
+    expect(pending.reserveTtlMs).toBe(30 * 60 * 1000);
+    // held reservation
+    let u = await userOf(t, "u");
+    expect(u.reservedTotalNanos).toBe(2_000_000_000);
+
+    // …later: the webhook fires and settles the actual cost.
+    await settleWith(t, r.requestId, { costNanos: 1_800_000_000 });
+    u = await userOf(t, "u");
+    expect(u.totalSpendNanos).toBe(1_800_000_000);
+    expect(u.reservedTotalNanos ?? 0).toBe(0);
+  });
+});
+
 describe("durable usage history", () => {
   test("settled spend lands in a per-day usage row", async () => {
     const t = convexTest(schema, modules);