@convex-dev/ai-budget 0.0.2-alpha.12 → 0.0.2-alpha.14

package/README.md

@@ -234,6 +234,55 @@ await ai.meter(ctx, { userId, model: "anthropic/claude-…", messages }, async (
 normalized) or explicit `promptTokens`/`completionTokens`/`cachedTokens`, plus
 optional `serverToolUses` (see [pricing](#pricing--cost)) and `costNanos`.
 
+**Cost known up front (image gen, per-call APIs).** When you know the price
+before the call — image generation (`n` × per-image), audio, anything per-call —
+pass `estimatedCostNanos` so the *reservation* holds the real amount and a hard
+cap is exact (the token estimate is meaningless for these). Price the units with
+`serverToolUses` + [`setServerTool`](#pricing--cost):
+
+```ts
+const IMG = 130_000_000; // $0.13/image
+await ai.meter(ctx,
+  { userId, model: "openai/gpt-image-1", messages: [{ role: "user", content: prompt }],
+    estimatedCostNanos: IMG * n },
+  async () => {
+    const res = await openrouter.images.generate({ model: "openai/gpt-image-1", prompt, n });
+    return { serverToolUses: { image: n } };   // priced via setServerTool({ tool: "image", … })
+  });
+```
+
+### `ai.begin` / `ai.settle` — long async jobs (video)
+
+A video job is submit → wait minutes → poll/webhook → done, spanning multiple
+Convex functions, so the synchronous `meter` bracket doesn't fit. Reserve with
+`begin` at submit, `settle` from the later context. Set `reserveTtlMs` to the
+job's max duration so the reconciler doesn't reap the hold mid-flight:
+
+```ts
+// submit (action): reserve, then kick off the job
+const started = await ai.begin(ctx, {
+  userId, model: "openai/sora",
+  estimatedCostNanos: perSecond * seconds,
+  reserveTtlMs: 20 * 60 * 1000,          // hold up to 20 min
+});
+if (!started.allowed) throw new ConvexError(started.reason);
+const job = await sora.videos.create({ … });
+await ctx.db.insert("videoJobs", { requestId: started.requestId, jobId: job.id });
+
+// later — settle from a poll, or a provider webhook:
+ai.registerWebhook(http, {
+  path: "/aibudget/video-done",
+  resolve: async (ctx, request, body) => {
+    if (!verifySignature(request, body)) return null;             // 202, ignored
+    const job = await lookupJob(ctx, body.id);                    // → { requestId }
+    return { requestId: job.requestId, serverToolUses: { video_seconds: body.seconds } };
+  },
+});
+```
+
+`begin` returns the admission result (it doesn't throw — check `allowed`).
+`settle` is idempotent (exactly-once server-side), so a retried webhook is safe.
+
 ### Replay
 
 ```ts
@@ -555,14 +604,19 @@ endpoints verbatim. In production:
 
 `example/` is a full working demo: chat as different personas on the left; a live
 admin panel on the right — the request audit log (inspect → edit → re-run, with
-lineage), a users table (limits, soft toggle, block, bump), and per-action budgets.
+lineage), a users table (limits, soft toggle, block, bump), per-action budgets,
+and a **⚡ Burst** tab that fires N real concurrent AI requests against one
+tightly-capped budget: watch reservations appear live, some requests get
+admitted (with real settled costs), and the rest get atomically rejected by the
+cap — the reserve-then-settle admission design, visible.
 
 ![Users & Limits admin table](docs/users.png)
 ![Actions & Budgets admin table](docs/actions.png)
 
 ```sh
 cd example
-npm install
+npm install         # also links the repo-root node_modules the demo's
+                    # ../../src component imports resolve through (postinstall)
 npx convex dev      # terminal 1 — provisions a dev deployment
 npm run dev         # terminal 2 — Vite app
 ```

package/dist/client/index.d.ts

@@ -113,15 +113,65 @@ export declare class AIBudget {
     private fireBudgetEvents;
     private fireLimitReached;
     /**
-     * Meter ANY 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.
-     * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
+     * 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.
+     */
+    begin(ctx: RunMutationCtx, opts: {
+        model: string;
+        messages?: Message[];
+        userId?: string;
+        action?: string;
+        tags?: Tag[];
+        /** 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;
+    }>;
+    /**
+     * 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`.
+     */
+    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;
+    }>;
+    /**
+     * 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 }`, priced on top of
-     * tokens) and an authoritative `costNanos` (used verbatim if present).
+     * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`) and an
+     * authoritative `costNanos` (used verbatim if present).
      */
     meter(ctx: RunMutationCtx, opts: {
         model: string;
@@ -130,6 +180,8 @@ export declare class AIBudget {
         action?: string;
         tags?: Tag[];
         rerunOf?: string;
+        /** Reserve this exact amount (nanodollars) instead of the token estimate. */
+        estimatedCostNanos?: number;
     }, run: () => Promise<{
         text?: string;
         usage?: any;
@@ -198,6 +250,7 @@ export declare class AIBudget {
             serverToolUses?: {
                 [x: string]: number;
             } | undefined;
+            reserveTtlMs?: number | undefined;
             costNanos?: number | undefined;
             latencyMs?: number | undefined;
             rerunOf?: string | undefined;
@@ -233,6 +286,7 @@ export declare class AIBudget {
             serverToolUses?: {
                 [x: string]: number;
             } | undefined;
+            reserveTtlMs?: number | undefined;
             costNanos?: number | undefined;
             latencyMs?: number | undefined;
             rerunOf?: string | undefined;
@@ -269,6 +323,7 @@ export declare class AIBudget {
                 serverToolUses?: {
                     [x: string]: number;
                 } | undefined;
+                reserveTtlMs?: number | undefined;
                 costNanos?: number | undefined;
                 latencyMs?: number | undefined;
                 rerunOf?: string | undefined;
@@ -301,6 +356,7 @@ export declare class AIBudget {
                 serverToolUses?: {
                     [x: string]: number;
                 } | undefined;
+                reserveTtlMs?: number | undefined;
                 costNanos?: number | undefined;
                 latencyMs?: number | undefined;
                 rerunOf?: string | undefined;
@@ -834,6 +890,37 @@ export declare class AIBudget {
         /** Return true to allow the request. Runs on the HTML page and every API call. */
         authorize?: (ctx: any, request: Request) => boolean | Promise<boolean>;
     }): void;
+    /**
+     * 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>;
+    }): void;
 }
 /** @deprecated Renamed to `AIBudget`. */
 export declare const WorryFreeAI: typeof AIBudget;

package/dist/client/index.js

@@ -155,17 +155,14 @@ export class AIBudget {
         }
     }
     /**
-     * Meter ANY 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.
-     * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
-     *
-     * `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 }`, priced on top of
-     * tokens) and an authoritative `costNanos` (used verbatim if present).
+     * 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 meter(ctx, opts, run) {
+    async begin(ctx, opts) {
         const userId = await resolveUserId(ctx, opts.userId);
         const actionName = await resolveActionName(ctx, opts.action);
         const started = await ctx.runMutation(this.component.lib.startRequest, {
@@ -173,10 +170,15 @@ export class AIBudget {
             actionName,
             tags: opts.tags,
             model: opts.model,
-            messages: opts.messages,
+            messages: opts.messages ?? [],
+            estimatedCostNanos: opts.estimatedCostNanos,
+            reserveTtlMs: opts.reserveTtlMs,
             rerunOf: opts.rerunOf,
         });
-        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,
@@ -185,19 +187,71 @@ export class AIBudget {
                 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, args) {
+        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,
+            ...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, opts, run) {
+        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, notices } = started;
-        await this.fireBudgetEvents({ userId, action: actionName, tags: opts.tags, requestId }, warnings, notices);
+        const { requestId, warnings, notices } = started;
         const start = Date.now();
         try {
             const out = await run();
-            // Explicit token fields win; otherwise normalize a raw provider usage.
+            const { costNanos } = await this.settle(ctx, {
+                requestId,
+                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
@@ -207,22 +261,10 @@ export class AIBudget {
                     cachedTokens: out.cachedTokens ?? 0,
                 }
                 : extractUsage(out.usage);
-            const { costNanos } = await ctx.runMutation(this.component.lib.finishRequest, {
-                requestId,
-                responseText: out.text,
-                ...usage,
-                serverToolUses: out.serverToolUses,
-                costNanos: out.costNanos,
-                latencyMs: Date.now() - start,
-            });
             return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
         }
         catch (e) {
-            await ctx.runMutation(this.component.lib.finishRequest, {
-                requestId,
-                error: String(e),
-                latencyMs: Date.now() - start,
-            });
+            await this.settle(ctx, { requestId, error: String(e), latencyMs: Date.now() - start });
             throw e;
         }
     }
@@ -627,6 +669,38 @@ 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, opts) {
+        const path = opts.path ?? "/aibudget/webhook";
+        const self = this;
+        http.route({
+            path,
+            method: "POST",
+            handler: httpActionGeneric(async (ctx, 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`. */
 export const WorryFreeAI = AIBudget;

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

@@ -140,12 +140,14 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         }, 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;

package/dist/component/lib.d.ts

@@ -4,7 +4,9 @@ export declare const startRequest: import("convex/server").RegisteredMutation<"p
         dimension: string;
         value: string;
     }[] | undefined;
+    reserveTtlMs?: number | undefined;
     rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
+    estimatedCostNanos?: number | undefined;
     userId: string;
     model: string;
     messages: {
@@ -67,6 +69,7 @@ export declare const lineage: import("convex/server").RegisteredQuery<"public",
         completionTokens?: number | undefined;
         cachedTokens?: number | undefined;
         serverToolUses?: Record<string, number> | undefined;
+        reserveTtlMs?: number | undefined;
         costNanos?: number | undefined;
         latencyMs?: number | undefined;
         rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -97,6 +100,7 @@ export declare const lineage: import("convex/server").RegisteredQuery<"public",
         completionTokens?: number | undefined;
         cachedTokens?: number | undefined;
         serverToolUses?: Record<string, number> | undefined;
+        reserveTtlMs?: number | undefined;
         costNanos?: number | undefined;
         latencyMs?: number | undefined;
         rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -130,6 +134,7 @@ export declare const getRequest: import("convex/server").RegisteredQuery<"public
     completionTokens?: number | undefined;
     cachedTokens?: number | undefined;
     serverToolUses?: Record<string, number> | undefined;
+    reserveTtlMs?: number | undefined;
     costNanos?: number | undefined;
     latencyMs?: number | undefined;
     rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -165,6 +170,7 @@ export declare const listRequests: import("convex/server").RegisteredQuery<"publ
     completionTokens?: number | undefined;
     cachedTokens?: number | undefined;
     serverToolUses?: Record<string, number> | undefined;
+    reserveTtlMs?: number | undefined;
     costNanos?: number | undefined;
     latencyMs?: number | undefined;
     rerunOf?: import("convex/values").GenericId<"requests"> | undefined;

package/dist/component/lib.js

@@ -303,6 +303,15 @@ export const startRequest = mutation({
         tags: v.optional(v.array(vTag)),
         model: v.string(),
         messages: v.array(vMessage),
+        // Reserve this exact amount (nanodollars) instead of the token-based
+        // estimate. Use it whenever the cost is known up front — image generation
+        // (n × per-image), audio, per-call APIs — so a hard cap reserves the real
+        // amount rather than a meaningless token guess.
+        estimatedCostNanos: v.optional(v.number()),
+        // Hold the reservation this long (ms) before the reconciler may reap it —
+        // for long async jobs (video) that settle minutes later. Extends the 30-min
+        // default floor.
+        reserveTtlMs: v.optional(v.number()),
         rerunOf: v.optional(v.id("requests")),
     },
     returns: vStartResult,
@@ -315,7 +324,7 @@ export const startRequest = mutation({
         // bloat the 60s rate-limit window read below.
         const reject = async (code, reason, persist = true) => {
             if (persist) {
-                await ctx.db.insert("requests", {
+                const requestId = await ctx.db.insert("requests", {
                     userId: args.userId,
                     actionName: args.actionName,
                     ...(extraTags.length ? { tags: extraTags } : {}),
@@ -325,6 +334,15 @@ export const startRequest = mutation({
                     status: "blocked",
                     error: reason,
                 });
+                // Reverse-index the blocked attempt too, so tag-filtered request logs
+                // show rejections alongside admitted traffic.
+                for (const t of extraTags) {
+                    await ctx.db.insert("requestTags", {
+                        dimension: t.dimension,
+                        value: t.value,
+                        requestId,
+                    });
+                }
             }
             return { allowed: false, code, reason };
         };
@@ -332,6 +350,11 @@ export const startRequest = mutation({
         const month = monthStamp();
         const priceInfo = await getPrice(ctx, args.model);
         const est = estimateUsage(args.messages, priceInfo);
+        // A caller-supplied known cost (image gen, audio, per-call APIs) reserves
+        // the real amount up front; the token estimate stays as the token reserve.
+        if (args.estimatedCostNanos !== undefined && args.estimatedCostNanos >= 0) {
+            est.cost = Math.round(args.estimatedCostNanos);
+        }
         const warnings = [];
         const notices = [];
         // Model allow/deny policy (component-wide).
@@ -396,13 +419,17 @@ export const startRequest = mutation({
                 recentCount = recent.filter((r) => r.status !== "blocked").length;
             }
             else {
-                recentCount = (await ctx.db
+                // Tag rows also cover persisted blocked attempts; fetch each request
+                // to exclude them, matching the user/action paths above.
+                const tagRows = await ctx.db
                     .query("requestTags")
                     .withIndex("dim_value", (q) => q
                     .eq("dimension", b.dimension)
                     .eq("value", b.value)
                     .gt("_creationTime", rateCutoff))
-                    .take(limit)).length;
+                    .take(limit + 50);
+                const recent = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
+                recentCount = recent.filter((r) => r !== null && r.status !== "blocked").length;
             }
             if (recentCount >= limit) {
                 const code = b.dimension === USER_DIM ? "rate_limit" : `${b.dimension}_rate_limit`;
@@ -521,6 +548,7 @@ export const startRequest = mutation({
             model: args.model,
             messages: args.messages,
             rerunOf: args.rerunOf,
+            ...(args.reserveTtlMs !== undefined ? { reserveTtlMs: args.reserveTtlMs } : {}),
             status: "pending",
             estimatedNanos: est.cost,
             estimatedTokens: est.tokens,
@@ -693,12 +721,19 @@ export const reconcile = internalMutation({
             .take(200);
         for (const req of toFold)
             await foldOne(ctx, req);
+        // Reap dead reservations: pending rows older than the default floor, but a
+        // per-request reserveTtlMs (set for long async jobs like video) holds the
+        // reservation until *its* deadline so a still-running job isn't reaped.
         const cutoff = Date.now() - STALE_PENDING_MS;
-        const stale = await ctx.db
+        const candidates = await ctx.db
             .query("requests")
             .withIndex("status", (q) => q.eq("status", "pending").lt("_creationTime", cutoff))
             .take(200);
-        for (const req of stale) {
+        let expired = 0;
+        for (const req of candidates) {
+            const ttl = req.reserveTtlMs ?? STALE_PENDING_MS;
+            if (Date.now() - req._creationTime <= ttl)
+                continue; // still within its window
             await ctx.db.patch(req._id, {
                 status: "error",
                 error: "Timed out before settling; reservation released",
@@ -706,6 +741,7 @@ export const reconcile = internalMutation({
                 settled: false,
             });
             await foldOne(ctx, await ctx.db.get(req._id));
+            expired++;
         }
         // Retention: delete terminal, fully-accounted request rows past the window.
         const settings = await getSettings(ctx);
@@ -727,7 +763,7 @@ export const reconcile = internalMutation({
                 }
             }
         }
-        return { folded: toFold.length, expired: stale.length, purged };
+        return { folded: toFold.length, expired, purged };
     },
 });
 export const setRetention = mutation({

package/dist/component/schema.d.ts

@@ -151,6 +151,7 @@ declare const _default: import("convex/server").SchemaDefinition<{
         completionTokens?: number | undefined;
         cachedTokens?: number | undefined;
         serverToolUses?: Record<string, number> | undefined;
+        reserveTtlMs?: number | undefined;
         costNanos?: number | undefined;
         latencyMs?: number | undefined;
         rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -197,10 +198,11 @@ declare const _default: import("convex/server").SchemaDefinition<{
         completionTokens: import("convex/values").VFloat64<number | undefined, "optional">;
         cachedTokens: import("convex/values").VFloat64<number | undefined, "optional">;
         serverToolUses: import("convex/values").VRecord<Record<string, number> | undefined, import("convex/values").VString<string, "required">, import("convex/values").VFloat64<number, "required">, "optional", string>;
+        reserveTtlMs: import("convex/values").VFloat64<number | undefined, "optional">;
         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" | "tags" | "model" | "estimatedNanos" | "estimatedTokens" | "unpricedModel" | "overBudget" | "settled" | "messages" | "status" | "error" | "responseText" | "promptTokens" | "completionTokens" | "cachedTokens" | "serverToolUses" | "costNanos" | "latencyMs" | "rerunOf" | `serverToolUses.${string}`>, {
+    }, "required", "userId" | "actionName" | "tags" | "model" | "estimatedNanos" | "estimatedTokens" | "unpricedModel" | "overBudget" | "settled" | "messages" | "status" | "error" | "responseText" | "promptTokens" | "completionTokens" | "cachedTokens" | "serverToolUses" | "reserveTtlMs" | "costNanos" | "latencyMs" | "rerunOf" | `serverToolUses.${string}`>, {
         userId: ["userId", "_creationTime"];
         status: ["status", "_creationTime"];
         rerunOf: ["rerunOf", "_creationTime"];

package/dist/component/schema.js

@@ -123,6 +123,11 @@ export default defineSchema({
         // server-side tool invocations that bill a per-call fee on top of tokens
         // (e.g. { web_search: 3 }). Priced via serverToolPrices at settle.
         serverToolUses: v.optional(v.record(v.string(), v.number())),
+        // How long the reservation may stay held before the reconciler reaps it as
+        // dead (ms). For long async jobs (video generation) set this to the job's
+        // max duration so the hold isn't released mid-flight. Extends the default
+        // 30-min floor; only stored while pending.
+        reserveTtlMs: v.optional(v.number()),
         costNanos: v.optional(v.number()),
         latencyMs: v.optional(v.number()),
         rerunOf: v.optional(v.id("requests")),

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.12",
+  "version": "0.0.2-alpha.14",
   "license": "Apache-2.0",
   "type": "module",
   "keywords": [

package/src/client/index.ts

@@ -306,36 +306,31 @@ export class AIBudget {
   }
 
   /**
-   * Meter ANY 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.
-   * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
-   *
-   * `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 }`, priced on top of
-   * tokens) and an authoritative `costNanos` (used verbatim if present).
+   * 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 meter(
+  async begin(
     ctx: RunMutationCtx,
     opts: {
       model: string;
-      messages: Message[];
+      messages?: Message[];
       userId?: string;
       action?: string;
       tags?: Tag[];
+      /** 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;
-    },
-    run: () => Promise<{
-      text?: string;
-      usage?: any;
-      promptTokens?: number;
-      completionTokens?: number;
-      cachedTokens?: number;
-      serverToolUses?: Record<string, number>;
-      costNanos?: number;
-    }>
-  ): Promise<ChatResult> {
+    }
+  ): 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, {
@@ -343,10 +338,18 @@ export class AIBudget {
       actionName,
       tags: opts.tags,
       model: opts.model,
-      messages: opts.messages,
+      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,
@@ -355,23 +358,109 @@ export class AIBudget {
         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, notices } = started;
-    await this.fireBudgetEvents(
-      { userId, action: actionName, tags: opts.tags, requestId },
-      warnings,
-      notices
-    );
+    const { requestId, warnings, notices } = started;
     const start = Date.now();
     try {
       const out = await run();
-      // Explicit token fields win; otherwise normalize a raw provider usage.
+      const { costNanos } = await this.settle(ctx, {
+        requestId,
+        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 ||
@@ -382,24 +471,9 @@ export class AIBudget {
               cachedTokens: out.cachedTokens ?? 0,
             }
           : extractUsage(out.usage);
-      const { costNanos } = await ctx.runMutation(
-        this.component.lib.finishRequest,
-        {
-          requestId,
-          responseText: out.text,
-          ...usage,
-          serverToolUses: out.serverToolUses,
-          costNanos: out.costNanos,
-          latencyMs: Date.now() - start,
-        }
-      );
       return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
     } catch (e) {
-      await ctx.runMutation(this.component.lib.finishRequest, {
-        requestId,
-        error: String(e),
-        latencyMs: Date.now() - start,
-      });
+      await this.settle(ctx, { requestId, error: String(e), latencyMs: Date.now() - start });
       throw e;
     }
   }
@@ -941,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

@@ -230,9 +230,11 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         "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

@@ -161,6 +161,64 @@ describe("server-tool pricing", () => {
   });
 });
 
+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);
@@ -272,6 +330,51 @@ describe("tag-filtered request log", () => {
     expect(acme.length).toBe(1);
     expect(acme[0].userId).toBe("u");
   });
+
+  test("blocked attempts appear in the tag-filtered log", async () => {
+    const t = convexTest(schema, modules);
+    const tags = [{ dimension: "burst", value: "run-1" }];
+    // Cap the tag bucket below one request's reservation so the attempt is
+    // budget-blocked (a persisted rejection).
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "burst",
+      value: "run-1",
+      lifetimeSpendLimitNanos: 1_000,
+    });
+    const r = await start(t, { userId: "u", tags });
+    expect(r.allowed).toBe(false);
+    const log = await t.query(api.lib.listRequests, {
+      dimension: "burst",
+      value: "run-1",
+    });
+    expect(log.length).toBe(1);
+    expect(log[0].status).toBe("blocked");
+  });
+
+  test("persisted blocked attempts don't consume a custom-tag rate limit", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "customer",
+      value: "acme",
+      requestsPerMinute: 1,
+      // also cap spend so attempts get budget-blocked (persisted) first
+      lifetimeSpendLimitNanos: 1_000,
+    });
+    const tags = [{ dimension: "customer", value: "acme" }];
+    // A budget-blocked (persisted) attempt writes a requestTags row…
+    const blocked = await start(t, { userId: "u1", tags });
+    expect(blocked.allowed).toBe(false);
+    expect(blocked.code).toBe("customer_lifetime_spend_limit");
+    // …which must NOT count toward the 1/min rate limit. Lift the spend cap:
+    // with no admitted requests in the window, the next request goes through.
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "customer",
+      value: "acme",
+      lifetimeSpendLimitNanos: 1_000_000_000,
+    });
+    const next = await start(t, { userId: "u2", tags });
+    expect(next.allowed).toBe(true);
+  });
 });
 
 describe("tagged attribution buckets", () => {

package/src/component/lib.ts

@@ -419,6 +419,15 @@ export const startRequest = mutation({
     tags: v.optional(v.array(vTag)),
     model: v.string(),
     messages: v.array(vMessage),
+    // Reserve this exact amount (nanodollars) instead of the token-based
+    // estimate. Use it whenever the cost is known up front — image generation
+    // (n × per-image), audio, per-call APIs — so a hard cap reserves the real
+    // amount rather than a meaningless token guess.
+    estimatedCostNanos: v.optional(v.number()),
+    // Hold the reservation this long (ms) before the reconciler may reap it —
+    // for long async jobs (video) that settle minutes later. Extends the 30-min
+    // default floor.
+    reserveTtlMs: v.optional(v.number()),
     rerunOf: v.optional(v.id("requests")),
   },
   returns: vStartResult,
@@ -431,7 +440,7 @@ export const startRequest = mutation({
     // bloat the 60s rate-limit window read below.
     const reject = async (code: string, reason: string, persist = true) => {
       if (persist) {
-        await ctx.db.insert("requests", {
+        const requestId = await ctx.db.insert("requests", {
           userId: args.userId,
           actionName: args.actionName,
           ...(extraTags.length ? { tags: extraTags } : {}),
@@ -441,6 +450,15 @@ export const startRequest = mutation({
           status: "blocked" as const,
           error: reason,
         });
+        // Reverse-index the blocked attempt too, so tag-filtered request logs
+        // show rejections alongside admitted traffic.
+        for (const t of extraTags) {
+          await ctx.db.insert("requestTags", {
+            dimension: t.dimension,
+            value: t.value,
+            requestId,
+          });
+        }
       }
       return { allowed: false as const, code, reason };
     };
@@ -449,6 +467,11 @@ export const startRequest = mutation({
     const month = monthStamp();
     const priceInfo = await getPrice(ctx, args.model);
     const est = estimateUsage(args.messages, priceInfo);
+    // A caller-supplied known cost (image gen, audio, per-call APIs) reserves
+    // the real amount up front; the token estimate stays as the token reserve.
+    if (args.estimatedCostNanos !== undefined && args.estimatedCostNanos >= 0) {
+      est.cost = Math.round(args.estimatedCostNanos);
+    }
     const warnings: string[] = [];
     const notices: string[] = [];
 
@@ -531,17 +554,19 @@ export const startRequest = mutation({
           .take(limit + 50);
         recentCount = recent.filter((r) => r.status !== "blocked").length;
       } else {
-        recentCount = (
-          await ctx.db
-            .query("requestTags")
-            .withIndex("dim_value", (q) =>
-              q
-                .eq("dimension", b.dimension)
-                .eq("value", b.value)
-                .gt("_creationTime", rateCutoff)
-            )
-            .take(limit)
-        ).length;
+        // Tag rows also cover persisted blocked attempts; fetch each request
+        // to exclude them, matching the user/action paths above.
+        const tagRows = await ctx.db
+          .query("requestTags")
+          .withIndex("dim_value", (q) =>
+            q
+              .eq("dimension", b.dimension)
+              .eq("value", b.value)
+              .gt("_creationTime", rateCutoff)
+          )
+          .take(limit + 50);
+        const recent = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
+        recentCount = recent.filter((r) => r !== null && r.status !== "blocked").length;
       }
 
       if (recentCount >= limit) {
@@ -682,6 +707,7 @@ export const startRequest = mutation({
       model: args.model,
       messages: args.messages,
       rerunOf: args.rerunOf,
+      ...(args.reserveTtlMs !== undefined ? { reserveTtlMs: args.reserveTtlMs } : {}),
       status: "pending",
       estimatedNanos: est.cost,
       estimatedTokens: est.tokens,
@@ -868,14 +894,20 @@ export const reconcile = internalMutation({
       .take(200);
     for (const req of toFold) await foldOne(ctx, req);
 
+    // Reap dead reservations: pending rows older than the default floor, but a
+    // per-request reserveTtlMs (set for long async jobs like video) holds the
+    // reservation until *its* deadline so a still-running job isn't reaped.
     const cutoff = Date.now() - STALE_PENDING_MS;
-    const stale = await ctx.db
+    const candidates = await ctx.db
       .query("requests")
       .withIndex("status", (q) =>
         q.eq("status", "pending").lt("_creationTime", cutoff)
       )
       .take(200);
-    for (const req of stale) {
+    let expired = 0;
+    for (const req of candidates) {
+      const ttl = req.reserveTtlMs ?? STALE_PENDING_MS;
+      if (Date.now() - req._creationTime <= ttl) continue; // still within its window
       await ctx.db.patch(req._id, {
         status: "error",
         error: "Timed out before settling; reservation released",
@@ -883,6 +915,7 @@ export const reconcile = internalMutation({
         settled: false,
       });
       await foldOne(ctx, await ctx.db.get(req._id));
+      expired++;
     }
 
     // Retention: delete terminal, fully-accounted request rows past the window.
@@ -907,7 +940,7 @@ export const reconcile = internalMutation({
         }
       }
     }
-    return { folded: toFold.length, expired: stale.length, purged };
+    return { folded: toFold.length, expired, purged };
   },
 });
 

package/src/component/schema.ts

@@ -135,6 +135,11 @@ export default defineSchema({
     // server-side tool invocations that bill a per-call fee on top of tokens
     // (e.g. { web_search: 3 }). Priced via serverToolPrices at settle.
     serverToolUses: v.optional(v.record(v.string(), v.number())),
+    // How long the reservation may stay held before the reconciler reaps it as
+    // dead (ms). For long async jobs (video generation) set this to the job's
+    // max duration so the hold isn't released mid-flight. Extends the default
+    // 30-min floor; only stored while pending.
+    reserveTtlMs: v.optional(v.number()),
     costNanos: v.optional(v.number()),
     latencyMs: v.optional(v.number()),
     rerunOf: v.optional(v.id("requests")),