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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md CHANGED
@@ -234,6 +234,55 @@ await ai.meter(ctx, { userId, model: "anthropic/claude-…", messages }, async (
234
234
  normalized) or explicit `promptTokens`/`completionTokens`/`cachedTokens`, plus
235
235
  optional `serverToolUses` (see [pricing](#pricing--cost)) and `costNanos`.
236
236
 
237
+ **Cost known up front (image gen, per-call APIs).** When you know the price
238
+ before the call — image generation (`n` × per-image), audio, anything per-call —
239
+ pass `estimatedCostNanos` so the *reservation* holds the real amount and a hard
240
+ cap is exact (the token estimate is meaningless for these). Price the units with
241
+ `serverToolUses` + [`setServerTool`](#pricing--cost):
242
+
243
+ ```ts
244
+ const IMG = 130_000_000; // $0.13/image
245
+ await ai.meter(ctx,
246
+ { userId, model: "openai/gpt-image-1", messages: [{ role: "user", content: prompt }],
247
+ estimatedCostNanos: IMG * n },
248
+ async () => {
249
+ const res = await openrouter.images.generate({ model: "openai/gpt-image-1", prompt, n });
250
+ return { serverToolUses: { image: n } }; // priced via setServerTool({ tool: "image", … })
251
+ });
252
+ ```
253
+
254
+ ### `ai.begin` / `ai.settle` — long async jobs (video)
255
+
256
+ A video job is submit → wait minutes → poll/webhook → done, spanning multiple
257
+ Convex functions, so the synchronous `meter` bracket doesn't fit. Reserve with
258
+ `begin` at submit, `settle` from the later context. Set `reserveTtlMs` to the
259
+ job's max duration so the reconciler doesn't reap the hold mid-flight:
260
+
261
+ ```ts
262
+ // submit (action): reserve, then kick off the job
263
+ const started = await ai.begin(ctx, {
264
+ userId, model: "openai/sora",
265
+ estimatedCostNanos: perSecond * seconds,
266
+ reserveTtlMs: 20 * 60 * 1000, // hold up to 20 min
267
+ });
268
+ if (!started.allowed) throw new ConvexError(started.reason);
269
+ const job = await sora.videos.create({ … });
270
+ await ctx.db.insert("videoJobs", { requestId: started.requestId, jobId: job.id });
271
+
272
+ // later — settle from a poll, or a provider webhook:
273
+ ai.registerWebhook(http, {
274
+ path: "/aibudget/video-done",
275
+ resolve: async (ctx, request, body) => {
276
+ if (!verifySignature(request, body)) return null; // 202, ignored
277
+ const job = await lookupJob(ctx, body.id); // → { requestId }
278
+ return { requestId: job.requestId, serverToolUses: { video_seconds: body.seconds } };
279
+ },
280
+ });
281
+ ```
282
+
283
+ `begin` returns the admission result (it doesn't throw — check `allowed`).
284
+ `settle` is idempotent (exactly-once server-side), so a retried webhook is safe.
285
+
237
286
  ### Replay
238
287
 
239
288
  ```ts
@@ -113,15 +113,65 @@ export declare class AIBudget {
113
113
  private fireBudgetEvents;
114
114
  private fireLimitReached;
115
115
  /**
116
- * Meter ANY LLM call gateway, a provider SDK, a raw fetch — with the same
117
- * budgets, audit log, and cost tracking. Reserves before your `run` (throwing
118
- * a ConvexError over a hard cap), runs it, then records the actual usage/cost.
119
- * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
116
+ * Reserve budget for a call WITHOUT running it the async half of the
117
+ * lifecycle. Use for long-running jobs (video generation) where the result
118
+ * arrives minutes later via a poll or webhook: `begin` here, then `settle`
119
+ * from that later context with the request's `requestId`. Returns the
120
+ * admission result (does NOT throw over a cap — check `allowed`). Set
121
+ * `reserveTtlMs` to the job's max duration so the hold isn't reaped mid-flight.
122
+ */
123
+ begin(ctx: RunMutationCtx, opts: {
124
+ model: string;
125
+ messages?: Message[];
126
+ userId?: string;
127
+ action?: string;
128
+ tags?: Tag[];
129
+ /** Reserve this exact amount (nanodollars) when the cost is known up front. */
130
+ estimatedCostNanos?: number;
131
+ /** Hold the reservation up to this long (ms) for long async jobs. */
132
+ reserveTtlMs?: number;
133
+ rerunOf?: string;
134
+ }): Promise<{
135
+ allowed: true;
136
+ requestId: string;
137
+ warnings: string[];
138
+ notices: string[];
139
+ } | {
140
+ allowed: false;
141
+ code: string;
142
+ reason: string;
143
+ }>;
144
+ /**
145
+ * Record the actual usage/cost of a `begin`-reserved request and release its
146
+ * reservation. Idempotent (exactly-once server-side). Pass a raw provider
147
+ * `usage` (auto-normalized) or explicit token counts, plus optional
148
+ * `serverToolUses` and an authoritative `costNanos`.
149
+ */
150
+ settle(ctx: RunMutationCtx, args: {
151
+ requestId: string;
152
+ responseText?: string;
153
+ error?: string;
154
+ usage?: any;
155
+ promptTokens?: number;
156
+ completionTokens?: number;
157
+ cachedTokens?: number;
158
+ serverToolUses?: Record<string, number>;
159
+ costNanos?: number;
160
+ latencyMs?: number;
161
+ }): Promise<{
162
+ costNanos: number;
163
+ }>;
164
+ /**
165
+ * Meter ANY synchronous LLM call — gateway, a provider SDK, a raw fetch — with
166
+ * the same budgets, audit log, and cost tracking. Reserves before your `run`
167
+ * (throwing a ConvexError over a hard cap), runs it, then records the actual
168
+ * usage/cost. The provider-agnostic core; `chat` is sugar over it. (For async
169
+ * jobs that settle later, use `begin`/`settle` instead.)
120
170
  *
121
171
  * `run` returns what happened. Pass a raw provider `usage` object (auto-
122
172
  * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
123
- * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`, priced on top of
124
- * tokens) and an authoritative `costNanos` (used verbatim if present).
173
+ * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`) and an
174
+ * authoritative `costNanos` (used verbatim if present).
125
175
  */
126
176
  meter(ctx: RunMutationCtx, opts: {
127
177
  model: string;
@@ -130,6 +180,8 @@ export declare class AIBudget {
130
180
  action?: string;
131
181
  tags?: Tag[];
132
182
  rerunOf?: string;
183
+ /** Reserve this exact amount (nanodollars) instead of the token estimate. */
184
+ estimatedCostNanos?: number;
133
185
  }, run: () => Promise<{
134
186
  text?: string;
135
187
  usage?: any;
@@ -198,6 +250,7 @@ export declare class AIBudget {
198
250
  serverToolUses?: {
199
251
  [x: string]: number;
200
252
  } | undefined;
253
+ reserveTtlMs?: number | undefined;
201
254
  costNanos?: number | undefined;
202
255
  latencyMs?: number | undefined;
203
256
  rerunOf?: string | undefined;
@@ -233,6 +286,7 @@ export declare class AIBudget {
233
286
  serverToolUses?: {
234
287
  [x: string]: number;
235
288
  } | undefined;
289
+ reserveTtlMs?: number | undefined;
236
290
  costNanos?: number | undefined;
237
291
  latencyMs?: number | undefined;
238
292
  rerunOf?: string | undefined;
@@ -269,6 +323,7 @@ export declare class AIBudget {
269
323
  serverToolUses?: {
270
324
  [x: string]: number;
271
325
  } | undefined;
326
+ reserveTtlMs?: number | undefined;
272
327
  costNanos?: number | undefined;
273
328
  latencyMs?: number | undefined;
274
329
  rerunOf?: string | undefined;
@@ -301,6 +356,7 @@ export declare class AIBudget {
301
356
  serverToolUses?: {
302
357
  [x: string]: number;
303
358
  } | undefined;
359
+ reserveTtlMs?: number | undefined;
304
360
  costNanos?: number | undefined;
305
361
  latencyMs?: number | undefined;
306
362
  rerunOf?: string | undefined;
@@ -834,6 +890,37 @@ export declare class AIBudget {
834
890
  /** Return true to allow the request. Runs on the HTML page and every API call. */
835
891
  authorize?: (ctx: any, request: Request) => boolean | Promise<boolean>;
836
892
  }): void;
893
+ /**
894
+ * Mount a POST webhook that settles a `begin`-reserved request from a
895
+ * provider's completion callback (async image/video jobs). You supply
896
+ * `resolve` — verify the payload's signature and map the provider's job id to
897
+ * your stored `requestId` + final usage/cost; the helper calls `settle`.
898
+ * Return `null` to ignore an unrecognized/duplicate callback (HTTP 202).
899
+ *
900
+ * ai.registerWebhook(http, {
901
+ * path: "/aibudget/video-done",
902
+ * resolve: async (ctx, request, body) => {
903
+ * if (!verifySignature(request, body)) return null;
904
+ * const job = await lookupJob(ctx, body.id); // your table: { requestId }
905
+ * return { requestId: job.requestId, serverToolUses: { video_seconds: body.seconds } };
906
+ * },
907
+ * });
908
+ */
909
+ registerWebhook(http: HttpRouter, opts: {
910
+ path?: string;
911
+ resolve: (ctx: any, request: Request, body: any) => Promise<({
912
+ requestId: string;
913
+ } & {
914
+ responseText?: string;
915
+ error?: string;
916
+ usage?: any;
917
+ promptTokens?: number;
918
+ completionTokens?: number;
919
+ cachedTokens?: number;
920
+ serverToolUses?: Record<string, number>;
921
+ costNanos?: number;
922
+ }) | null>;
923
+ }): void;
837
924
  }
838
925
  /** @deprecated Renamed to `AIBudget`. */
839
926
  export declare const WorryFreeAI: typeof AIBudget;
@@ -155,17 +155,14 @@ export class AIBudget {
155
155
  }
156
156
  }
157
157
  /**
158
- * Meter ANY LLM call gateway, a provider SDK, a raw fetch — with the same
159
- * budgets, audit log, and cost tracking. Reserves before your `run` (throwing
160
- * a ConvexError over a hard cap), runs it, then records the actual usage/cost.
161
- * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
162
- *
163
- * `run` returns what happened. Pass a raw provider `usage` object (auto-
164
- * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
165
- * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`, priced on top of
166
- * tokens) and an authoritative `costNanos` (used verbatim if present).
158
+ * Reserve budget for a call WITHOUT running it the async half of the
159
+ * lifecycle. Use for long-running jobs (video generation) where the result
160
+ * arrives minutes later via a poll or webhook: `begin` here, then `settle`
161
+ * from that later context with the request's `requestId`. Returns the
162
+ * admission result (does NOT throw over a cap — check `allowed`). Set
163
+ * `reserveTtlMs` to the job's max duration so the hold isn't reaped mid-flight.
167
164
  */
168
- async meter(ctx, opts, run) {
165
+ async begin(ctx, opts) {
169
166
  const userId = await resolveUserId(ctx, opts.userId);
170
167
  const actionName = await resolveActionName(ctx, opts.action);
171
168
  const started = await ctx.runMutation(this.component.lib.startRequest, {
@@ -173,10 +170,15 @@ export class AIBudget {
173
170
  actionName,
174
171
  tags: opts.tags,
175
172
  model: opts.model,
176
- messages: opts.messages,
173
+ messages: opts.messages ?? [],
174
+ estimatedCostNanos: opts.estimatedCostNanos,
175
+ reserveTtlMs: opts.reserveTtlMs,
177
176
  rerunOf: opts.rerunOf,
178
177
  });
179
- if (!started.allowed) {
178
+ if (started.allowed) {
179
+ await this.fireBudgetEvents({ userId, action: actionName, tags: opts.tags, requestId: started.requestId }, started.warnings, started.notices);
180
+ }
181
+ else {
180
182
  await this.fireLimitReached({
181
183
  userId,
182
184
  action: actionName,
@@ -185,19 +187,71 @@ export class AIBudget {
185
187
  code: started.code,
186
188
  reason: started.reason,
187
189
  });
190
+ }
191
+ return started;
192
+ }
193
+ /**
194
+ * Record the actual usage/cost of a `begin`-reserved request and release its
195
+ * reservation. Idempotent (exactly-once server-side). Pass a raw provider
196
+ * `usage` (auto-normalized) or explicit token counts, plus optional
197
+ * `serverToolUses` and an authoritative `costNanos`.
198
+ */
199
+ async settle(ctx, args) {
200
+ const { requestId, usage, promptTokens, completionTokens, cachedTokens, ...rest } = args;
201
+ const tokens = promptTokens !== undefined ||
202
+ completionTokens !== undefined ||
203
+ cachedTokens !== undefined
204
+ ? {
205
+ promptTokens: promptTokens ?? 0,
206
+ completionTokens: completionTokens ?? 0,
207
+ cachedTokens: cachedTokens ?? 0,
208
+ }
209
+ : usage !== undefined
210
+ ? extractUsage(usage)
211
+ : {};
212
+ return ctx.runMutation(this.component.lib.finishRequest, {
213
+ requestId: requestId,
214
+ ...tokens,
215
+ ...rest,
216
+ });
217
+ }
218
+ /**
219
+ * Meter ANY synchronous LLM call — gateway, a provider SDK, a raw fetch — with
220
+ * the same budgets, audit log, and cost tracking. Reserves before your `run`
221
+ * (throwing a ConvexError over a hard cap), runs it, then records the actual
222
+ * usage/cost. The provider-agnostic core; `chat` is sugar over it. (For async
223
+ * jobs that settle later, use `begin`/`settle` instead.)
224
+ *
225
+ * `run` returns what happened. Pass a raw provider `usage` object (auto-
226
+ * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
227
+ * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`) and an
228
+ * authoritative `costNanos` (used verbatim if present).
229
+ */
230
+ async meter(ctx, opts, run) {
231
+ const started = await this.begin(ctx, opts);
232
+ if (!started.allowed) {
188
233
  throw new ConvexError({
189
234
  kind: "AIBudgetLimit",
190
235
  code: started.code,
191
236
  reason: started.reason,
192
237
  });
193
238
  }
194
- const requestId = started.requestId;
195
- const { warnings, notices } = started;
196
- await this.fireBudgetEvents({ userId, action: actionName, tags: opts.tags, requestId }, warnings, notices);
239
+ const { requestId, warnings, notices } = started;
197
240
  const start = Date.now();
198
241
  try {
199
242
  const out = await run();
200
- // Explicit token fields win; otherwise normalize a raw provider usage.
243
+ const { costNanos } = await this.settle(ctx, {
244
+ requestId,
245
+ responseText: out.text,
246
+ usage: out.usage,
247
+ promptTokens: out.promptTokens,
248
+ completionTokens: out.completionTokens,
249
+ cachedTokens: out.cachedTokens,
250
+ serverToolUses: out.serverToolUses,
251
+ costNanos: out.costNanos,
252
+ latencyMs: Date.now() - start,
253
+ });
254
+ // Re-derive the recorded usage for the return value.
201
255
  const usage = out.promptTokens !== undefined ||
202
256
  out.completionTokens !== undefined ||
203
257
  out.cachedTokens !== undefined
@@ -207,22 +261,10 @@ export class AIBudget {
207
261
  cachedTokens: out.cachedTokens ?? 0,
208
262
  }
209
263
  : extractUsage(out.usage);
210
- const { costNanos } = await ctx.runMutation(this.component.lib.finishRequest, {
211
- requestId,
212
- responseText: out.text,
213
- ...usage,
214
- serverToolUses: out.serverToolUses,
215
- costNanos: out.costNanos,
216
- latencyMs: Date.now() - start,
217
- });
218
264
  return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
219
265
  }
220
266
  catch (e) {
221
- await ctx.runMutation(this.component.lib.finishRequest, {
222
- requestId,
223
- error: String(e),
224
- latencyMs: Date.now() - start,
225
- });
267
+ await this.settle(ctx, { requestId, error: String(e), latencyMs: Date.now() - start });
226
268
  throw e;
227
269
  }
228
270
  }
@@ -627,6 +669,38 @@ export class AIBudget {
627
669
  http.route({ pathPrefix: `${prefix}/`, method: "GET", handler });
628
670
  http.route({ pathPrefix: `${prefix}/`, method: "POST", handler });
629
671
  }
672
+ /**
673
+ * Mount a POST webhook that settles a `begin`-reserved request from a
674
+ * provider's completion callback (async image/video jobs). You supply
675
+ * `resolve` — verify the payload's signature and map the provider's job id to
676
+ * your stored `requestId` + final usage/cost; the helper calls `settle`.
677
+ * Return `null` to ignore an unrecognized/duplicate callback (HTTP 202).
678
+ *
679
+ * ai.registerWebhook(http, {
680
+ * path: "/aibudget/video-done",
681
+ * resolve: async (ctx, request, body) => {
682
+ * if (!verifySignature(request, body)) return null;
683
+ * const job = await lookupJob(ctx, body.id); // your table: { requestId }
684
+ * return { requestId: job.requestId, serverToolUses: { video_seconds: body.seconds } };
685
+ * },
686
+ * });
687
+ */
688
+ registerWebhook(http, opts) {
689
+ const path = opts.path ?? "/aibudget/webhook";
690
+ const self = this;
691
+ http.route({
692
+ path,
693
+ method: "POST",
694
+ handler: httpActionGeneric(async (ctx, request) => {
695
+ const body = await request.json().catch(() => ({}));
696
+ const settle = await opts.resolve(ctx, request, body);
697
+ if (!settle)
698
+ return new Response("ignored", { status: 202 });
699
+ await self.settle(ctx, settle);
700
+ return new Response("ok");
701
+ }),
702
+ });
703
+ }
630
704
  }
631
705
  /** @deprecated Renamed to `AIBudget`. */
632
706
  export const WorryFreeAI = AIBudget;
@@ -140,12 +140,14 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
140
140
  }, null, Name>;
141
141
  startRequest: FunctionReference<"mutation", "internal", {
142
142
  actionName?: string;
143
+ estimatedCostNanos?: number;
143
144
  messages: Array<{
144
145
  content: string;
145
146
  role: string;
146
147
  }>;
147
148
  model: string;
148
149
  rerunOf?: string;
150
+ reserveTtlMs?: number;
149
151
  tags?: Array<{
150
152
  dimension: string;
151
153
  value: string;
@@ -4,7 +4,9 @@ export declare const startRequest: import("convex/server").RegisteredMutation<"p
4
4
  dimension: string;
5
5
  value: string;
6
6
  }[] | undefined;
7
+ reserveTtlMs?: number | undefined;
7
8
  rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
9
+ estimatedCostNanos?: number | undefined;
8
10
  userId: string;
9
11
  model: string;
10
12
  messages: {
@@ -67,6 +69,7 @@ export declare const lineage: import("convex/server").RegisteredQuery<"public",
67
69
  completionTokens?: number | undefined;
68
70
  cachedTokens?: number | undefined;
69
71
  serverToolUses?: Record<string, number> | undefined;
72
+ reserveTtlMs?: number | undefined;
70
73
  costNanos?: number | undefined;
71
74
  latencyMs?: number | undefined;
72
75
  rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -97,6 +100,7 @@ export declare const lineage: import("convex/server").RegisteredQuery<"public",
97
100
  completionTokens?: number | undefined;
98
101
  cachedTokens?: number | undefined;
99
102
  serverToolUses?: Record<string, number> | undefined;
103
+ reserveTtlMs?: number | undefined;
100
104
  costNanos?: number | undefined;
101
105
  latencyMs?: number | undefined;
102
106
  rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -130,6 +134,7 @@ export declare const getRequest: import("convex/server").RegisteredQuery<"public
130
134
  completionTokens?: number | undefined;
131
135
  cachedTokens?: number | undefined;
132
136
  serverToolUses?: Record<string, number> | undefined;
137
+ reserveTtlMs?: number | undefined;
133
138
  costNanos?: number | undefined;
134
139
  latencyMs?: number | undefined;
135
140
  rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -165,6 +170,7 @@ export declare const listRequests: import("convex/server").RegisteredQuery<"publ
165
170
  completionTokens?: number | undefined;
166
171
  cachedTokens?: number | undefined;
167
172
  serverToolUses?: Record<string, number> | undefined;
173
+ reserveTtlMs?: number | undefined;
168
174
  costNanos?: number | undefined;
169
175
  latencyMs?: number | undefined;
170
176
  rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -303,6 +303,15 @@ export const startRequest = mutation({
303
303
  tags: v.optional(v.array(vTag)),
304
304
  model: v.string(),
305
305
  messages: v.array(vMessage),
306
+ // Reserve this exact amount (nanodollars) instead of the token-based
307
+ // estimate. Use it whenever the cost is known up front — image generation
308
+ // (n × per-image), audio, per-call APIs — so a hard cap reserves the real
309
+ // amount rather than a meaningless token guess.
310
+ estimatedCostNanos: v.optional(v.number()),
311
+ // Hold the reservation this long (ms) before the reconciler may reap it —
312
+ // for long async jobs (video) that settle minutes later. Extends the 30-min
313
+ // default floor.
314
+ reserveTtlMs: v.optional(v.number()),
306
315
  rerunOf: v.optional(v.id("requests")),
307
316
  },
308
317
  returns: vStartResult,
@@ -332,6 +341,11 @@ export const startRequest = mutation({
332
341
  const month = monthStamp();
333
342
  const priceInfo = await getPrice(ctx, args.model);
334
343
  const est = estimateUsage(args.messages, priceInfo);
344
+ // A caller-supplied known cost (image gen, audio, per-call APIs) reserves
345
+ // the real amount up front; the token estimate stays as the token reserve.
346
+ if (args.estimatedCostNanos !== undefined && args.estimatedCostNanos >= 0) {
347
+ est.cost = Math.round(args.estimatedCostNanos);
348
+ }
335
349
  const warnings = [];
336
350
  const notices = [];
337
351
  // Model allow/deny policy (component-wide).
@@ -521,6 +535,7 @@ export const startRequest = mutation({
521
535
  model: args.model,
522
536
  messages: args.messages,
523
537
  rerunOf: args.rerunOf,
538
+ ...(args.reserveTtlMs !== undefined ? { reserveTtlMs: args.reserveTtlMs } : {}),
524
539
  status: "pending",
525
540
  estimatedNanos: est.cost,
526
541
  estimatedTokens: est.tokens,
@@ -693,12 +708,19 @@ export const reconcile = internalMutation({
693
708
  .take(200);
694
709
  for (const req of toFold)
695
710
  await foldOne(ctx, req);
711
+ // Reap dead reservations: pending rows older than the default floor, but a
712
+ // per-request reserveTtlMs (set for long async jobs like video) holds the
713
+ // reservation until *its* deadline so a still-running job isn't reaped.
696
714
  const cutoff = Date.now() - STALE_PENDING_MS;
697
- const stale = await ctx.db
715
+ const candidates = await ctx.db
698
716
  .query("requests")
699
717
  .withIndex("status", (q) => q.eq("status", "pending").lt("_creationTime", cutoff))
700
718
  .take(200);
701
- for (const req of stale) {
719
+ let expired = 0;
720
+ for (const req of candidates) {
721
+ const ttl = req.reserveTtlMs ?? STALE_PENDING_MS;
722
+ if (Date.now() - req._creationTime <= ttl)
723
+ continue; // still within its window
702
724
  await ctx.db.patch(req._id, {
703
725
  status: "error",
704
726
  error: "Timed out before settling; reservation released",
@@ -706,6 +728,7 @@ export const reconcile = internalMutation({
706
728
  settled: false,
707
729
  });
708
730
  await foldOne(ctx, await ctx.db.get(req._id));
731
+ expired++;
709
732
  }
710
733
  // Retention: delete terminal, fully-accounted request rows past the window.
711
734
  const settings = await getSettings(ctx);
@@ -727,7 +750,7 @@ export const reconcile = internalMutation({
727
750
  }
728
751
  }
729
752
  }
730
- return { folded: toFold.length, expired: stale.length, purged };
753
+ return { folded: toFold.length, expired, purged };
731
754
  },
732
755
  });
733
756
  export const setRetention = mutation({
@@ -151,6 +151,7 @@ declare const _default: import("convex/server").SchemaDefinition<{
151
151
  completionTokens?: number | undefined;
152
152
  cachedTokens?: number | undefined;
153
153
  serverToolUses?: Record<string, number> | undefined;
154
+ reserveTtlMs?: number | undefined;
154
155
  costNanos?: number | undefined;
155
156
  latencyMs?: number | undefined;
156
157
  rerunOf?: import("convex/values").GenericId<"requests"> | undefined;
@@ -197,10 +198,11 @@ declare const _default: import("convex/server").SchemaDefinition<{
197
198
  completionTokens: import("convex/values").VFloat64<number | undefined, "optional">;
198
199
  cachedTokens: import("convex/values").VFloat64<number | undefined, "optional">;
199
200
  serverToolUses: import("convex/values").VRecord<Record<string, number> | undefined, import("convex/values").VString<string, "required">, import("convex/values").VFloat64<number, "required">, "optional", string>;
201
+ reserveTtlMs: import("convex/values").VFloat64<number | undefined, "optional">;
200
202
  costNanos: import("convex/values").VFloat64<number | undefined, "optional">;
201
203
  latencyMs: import("convex/values").VFloat64<number | undefined, "optional">;
202
204
  rerunOf: import("convex/values").VId<import("convex/values").GenericId<"requests"> | undefined, "optional">;
203
- }, "required", "userId" | "actionName" | "tags" | "model" | "estimatedNanos" | "estimatedTokens" | "unpricedModel" | "overBudget" | "settled" | "messages" | "status" | "error" | "responseText" | "promptTokens" | "completionTokens" | "cachedTokens" | "serverToolUses" | "costNanos" | "latencyMs" | "rerunOf" | `serverToolUses.${string}`>, {
205
+ }, "required", "userId" | "actionName" | "tags" | "model" | "estimatedNanos" | "estimatedTokens" | "unpricedModel" | "overBudget" | "settled" | "messages" | "status" | "error" | "responseText" | "promptTokens" | "completionTokens" | "cachedTokens" | "serverToolUses" | "reserveTtlMs" | "costNanos" | "latencyMs" | "rerunOf" | `serverToolUses.${string}`>, {
204
206
  userId: ["userId", "_creationTime"];
205
207
  status: ["status", "_creationTime"];
206
208
  rerunOf: ["rerunOf", "_creationTime"];
@@ -123,6 +123,11 @@ export default defineSchema({
123
123
  // server-side tool invocations that bill a per-call fee on top of tokens
124
124
  // (e.g. { web_search: 3 }). Priced via serverToolPrices at settle.
125
125
  serverToolUses: v.optional(v.record(v.string(), v.number())),
126
+ // How long the reservation may stay held before the reconciler reaps it as
127
+ // dead (ms). For long async jobs (video generation) set this to the job's
128
+ // max duration so the hold isn't released mid-flight. Extends the default
129
+ // 30-min floor; only stored while pending.
130
+ reserveTtlMs: v.optional(v.number()),
126
131
  costNanos: v.optional(v.number()),
127
132
  latencyMs: v.optional(v.number()),
128
133
  rerunOf: v.optional(v.id("requests")),
package/package.json CHANGED
@@ -7,7 +7,7 @@
7
7
  "email": "support@convex.dev",
8
8
  "url": "https://github.com/get-convex/ai-budget/issues"
9
9
  },
10
- "version": "0.0.2-alpha.12",
10
+ "version": "0.0.2-alpha.13",
11
11
  "license": "Apache-2.0",
12
12
  "type": "module",
13
13
  "keywords": [
@@ -306,36 +306,31 @@ export class AIBudget {
306
306
  }
307
307
 
308
308
  /**
309
- * Meter ANY LLM call gateway, a provider SDK, a raw fetch — with the same
310
- * budgets, audit log, and cost tracking. Reserves before your `run` (throwing
311
- * a ConvexError over a hard cap), runs it, then records the actual usage/cost.
312
- * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
313
- *
314
- * `run` returns what happened. Pass a raw provider `usage` object (auto-
315
- * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
316
- * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`, priced on top of
317
- * tokens) and an authoritative `costNanos` (used verbatim if present).
309
+ * Reserve budget for a call WITHOUT running it the async half of the
310
+ * lifecycle. Use for long-running jobs (video generation) where the result
311
+ * arrives minutes later via a poll or webhook: `begin` here, then `settle`
312
+ * from that later context with the request's `requestId`. Returns the
313
+ * admission result (does NOT throw over a cap — check `allowed`). Set
314
+ * `reserveTtlMs` to the job's max duration so the hold isn't reaped mid-flight.
318
315
  */
319
- async meter(
316
+ async begin(
320
317
  ctx: RunMutationCtx,
321
318
  opts: {
322
319
  model: string;
323
- messages: Message[];
320
+ messages?: Message[];
324
321
  userId?: string;
325
322
  action?: string;
326
323
  tags?: Tag[];
324
+ /** Reserve this exact amount (nanodollars) when the cost is known up front. */
325
+ estimatedCostNanos?: number;
326
+ /** Hold the reservation up to this long (ms) for long async jobs. */
327
+ reserveTtlMs?: number;
327
328
  rerunOf?: string;
328
- },
329
- run: () => Promise<{
330
- text?: string;
331
- usage?: any;
332
- promptTokens?: number;
333
- completionTokens?: number;
334
- cachedTokens?: number;
335
- serverToolUses?: Record<string, number>;
336
- costNanos?: number;
337
- }>
338
- ): Promise<ChatResult> {
329
+ }
330
+ ): Promise<
331
+ | { allowed: true; requestId: string; warnings: string[]; notices: string[] }
332
+ | { allowed: false; code: string; reason: string }
333
+ > {
339
334
  const userId = await resolveUserId(ctx, opts.userId);
340
335
  const actionName = await resolveActionName(ctx, opts.action);
341
336
  const started = await ctx.runMutation(this.component.lib.startRequest, {
@@ -343,10 +338,18 @@ export class AIBudget {
343
338
  actionName,
344
339
  tags: opts.tags,
345
340
  model: opts.model,
346
- messages: opts.messages,
341
+ messages: opts.messages ?? [],
342
+ estimatedCostNanos: opts.estimatedCostNanos,
343
+ reserveTtlMs: opts.reserveTtlMs,
347
344
  rerunOf: opts.rerunOf as any,
348
345
  });
349
- if (!started.allowed) {
346
+ if (started.allowed) {
347
+ await this.fireBudgetEvents(
348
+ { userId, action: actionName, tags: opts.tags, requestId: started.requestId },
349
+ started.warnings,
350
+ started.notices
351
+ );
352
+ } else {
350
353
  await this.fireLimitReached({
351
354
  userId,
352
355
  action: actionName,
@@ -355,23 +358,109 @@ export class AIBudget {
355
358
  code: started.code,
356
359
  reason: started.reason,
357
360
  });
361
+ }
362
+ return started;
363
+ }
364
+
365
+ /**
366
+ * Record the actual usage/cost of a `begin`-reserved request and release its
367
+ * reservation. Idempotent (exactly-once server-side). Pass a raw provider
368
+ * `usage` (auto-normalized) or explicit token counts, plus optional
369
+ * `serverToolUses` and an authoritative `costNanos`.
370
+ */
371
+ async settle(
372
+ ctx: RunMutationCtx,
373
+ args: {
374
+ requestId: string;
375
+ responseText?: string;
376
+ error?: string;
377
+ usage?: any;
378
+ promptTokens?: number;
379
+ completionTokens?: number;
380
+ cachedTokens?: number;
381
+ serverToolUses?: Record<string, number>;
382
+ costNanos?: number;
383
+ latencyMs?: number;
384
+ }
385
+ ): Promise<{ costNanos: number }> {
386
+ const { requestId, usage, promptTokens, completionTokens, cachedTokens, ...rest } = args;
387
+ const tokens =
388
+ promptTokens !== undefined ||
389
+ completionTokens !== undefined ||
390
+ cachedTokens !== undefined
391
+ ? {
392
+ promptTokens: promptTokens ?? 0,
393
+ completionTokens: completionTokens ?? 0,
394
+ cachedTokens: cachedTokens ?? 0,
395
+ }
396
+ : usage !== undefined
397
+ ? extractUsage(usage)
398
+ : {};
399
+ return ctx.runMutation(this.component.lib.finishRequest, {
400
+ requestId: requestId as any,
401
+ ...tokens,
402
+ ...rest,
403
+ });
404
+ }
405
+
406
+ /**
407
+ * Meter ANY synchronous LLM call — gateway, a provider SDK, a raw fetch — with
408
+ * the same budgets, audit log, and cost tracking. Reserves before your `run`
409
+ * (throwing a ConvexError over a hard cap), runs it, then records the actual
410
+ * usage/cost. The provider-agnostic core; `chat` is sugar over it. (For async
411
+ * jobs that settle later, use `begin`/`settle` instead.)
412
+ *
413
+ * `run` returns what happened. Pass a raw provider `usage` object (auto-
414
+ * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
415
+ * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`) and an
416
+ * authoritative `costNanos` (used verbatim if present).
417
+ */
418
+ async meter(
419
+ ctx: RunMutationCtx,
420
+ opts: {
421
+ model: string;
422
+ messages: Message[];
423
+ userId?: string;
424
+ action?: string;
425
+ tags?: Tag[];
426
+ rerunOf?: string;
427
+ /** Reserve this exact amount (nanodollars) instead of the token estimate. */
428
+ estimatedCostNanos?: number;
429
+ },
430
+ run: () => Promise<{
431
+ text?: string;
432
+ usage?: any;
433
+ promptTokens?: number;
434
+ completionTokens?: number;
435
+ cachedTokens?: number;
436
+ serverToolUses?: Record<string, number>;
437
+ costNanos?: number;
438
+ }>
439
+ ): Promise<ChatResult> {
440
+ const started = await this.begin(ctx, opts);
441
+ if (!started.allowed) {
358
442
  throw new ConvexError({
359
443
  kind: "AIBudgetLimit",
360
444
  code: started.code,
361
445
  reason: started.reason,
362
446
  });
363
447
  }
364
- const requestId = started.requestId;
365
- const { warnings, notices } = started;
366
- await this.fireBudgetEvents(
367
- { userId, action: actionName, tags: opts.tags, requestId },
368
- warnings,
369
- notices
370
- );
448
+ const { requestId, warnings, notices } = started;
371
449
  const start = Date.now();
372
450
  try {
373
451
  const out = await run();
374
- // Explicit token fields win; otherwise normalize a raw provider usage.
452
+ const { costNanos } = await this.settle(ctx, {
453
+ requestId,
454
+ responseText: out.text,
455
+ usage: out.usage,
456
+ promptTokens: out.promptTokens,
457
+ completionTokens: out.completionTokens,
458
+ cachedTokens: out.cachedTokens,
459
+ serverToolUses: out.serverToolUses,
460
+ costNanos: out.costNanos,
461
+ latencyMs: Date.now() - start,
462
+ });
463
+ // Re-derive the recorded usage for the return value.
375
464
  const usage =
376
465
  out.promptTokens !== undefined ||
377
466
  out.completionTokens !== undefined ||
@@ -382,24 +471,9 @@ export class AIBudget {
382
471
  cachedTokens: out.cachedTokens ?? 0,
383
472
  }
384
473
  : extractUsage(out.usage);
385
- const { costNanos } = await ctx.runMutation(
386
- this.component.lib.finishRequest,
387
- {
388
- requestId,
389
- responseText: out.text,
390
- ...usage,
391
- serverToolUses: out.serverToolUses,
392
- costNanos: out.costNanos,
393
- latencyMs: Date.now() - start,
394
- }
395
- );
396
474
  return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
397
475
  } catch (e) {
398
- await ctx.runMutation(this.component.lib.finishRequest, {
399
- requestId,
400
- error: String(e),
401
- latencyMs: Date.now() - start,
402
- });
476
+ await this.settle(ctx, { requestId, error: String(e), latencyMs: Date.now() - start });
403
477
  throw e;
404
478
  }
405
479
  }
@@ -941,6 +1015,60 @@ export class AIBudget {
941
1015
  http.route({ pathPrefix: `${prefix}/`, method: "GET", handler });
942
1016
  http.route({ pathPrefix: `${prefix}/`, method: "POST", handler });
943
1017
  }
1018
+
1019
+ /**
1020
+ * Mount a POST webhook that settles a `begin`-reserved request from a
1021
+ * provider's completion callback (async image/video jobs). You supply
1022
+ * `resolve` — verify the payload's signature and map the provider's job id to
1023
+ * your stored `requestId` + final usage/cost; the helper calls `settle`.
1024
+ * Return `null` to ignore an unrecognized/duplicate callback (HTTP 202).
1025
+ *
1026
+ * ai.registerWebhook(http, {
1027
+ * path: "/aibudget/video-done",
1028
+ * resolve: async (ctx, request, body) => {
1029
+ * if (!verifySignature(request, body)) return null;
1030
+ * const job = await lookupJob(ctx, body.id); // your table: { requestId }
1031
+ * return { requestId: job.requestId, serverToolUses: { video_seconds: body.seconds } };
1032
+ * },
1033
+ * });
1034
+ */
1035
+ registerWebhook(
1036
+ http: HttpRouter,
1037
+ opts: {
1038
+ path?: string;
1039
+ resolve: (
1040
+ ctx: any,
1041
+ request: Request,
1042
+ body: any
1043
+ ) => Promise<
1044
+ | ({ requestId: string } & {
1045
+ responseText?: string;
1046
+ error?: string;
1047
+ usage?: any;
1048
+ promptTokens?: number;
1049
+ completionTokens?: number;
1050
+ cachedTokens?: number;
1051
+ serverToolUses?: Record<string, number>;
1052
+ costNanos?: number;
1053
+ })
1054
+ | null
1055
+ >;
1056
+ }
1057
+ ) {
1058
+ const path = opts.path ?? "/aibudget/webhook";
1059
+ const self = this;
1060
+ http.route({
1061
+ path,
1062
+ method: "POST",
1063
+ handler: httpActionGeneric(async (ctx: any, request: Request) => {
1064
+ const body = await request.json().catch(() => ({}));
1065
+ const settle = await opts.resolve(ctx, request, body);
1066
+ if (!settle) return new Response("ignored", { status: 202 });
1067
+ await self.settle(ctx, settle);
1068
+ return new Response("ok");
1069
+ }),
1070
+ });
1071
+ }
944
1072
  }
945
1073
 
946
1074
  /** @deprecated Renamed to `AIBudget`. */
@@ -230,9 +230,11 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
230
230
  "internal",
231
231
  {
232
232
  actionName?: string;
233
+ estimatedCostNanos?: number;
233
234
  messages: Array<{ content: string; role: string }>;
234
235
  model: string;
235
236
  rerunOf?: string;
237
+ reserveTtlMs?: number;
236
238
  tags?: Array<{ dimension: string; value: string }>;
237
239
  userId: string;
238
240
  },
@@ -161,6 +161,64 @@ describe("server-tool pricing", () => {
161
161
  });
162
162
  });
163
163
 
164
+ describe("cost known up front (image gen, per-call APIs)", () => {
165
+ test("estimatedCostNanos drives the reservation for a hard cap", async () => {
166
+ const t = convexTest(schema, modules);
167
+ await setUserLimits(t, "u", { dailySpendLimitNanos: 100_000_000 }); // $0.10
168
+ // A $0.13 image is known before the call; reserving it exceeds the cap,
169
+ // even though the token estimate for the prompt alone would pass.
170
+ const r = await start(t, {
171
+ userId: "u",
172
+ model: "openai/gpt-image-1",
173
+ estimatedCostNanos: 130_000_000,
174
+ });
175
+ expect(r.allowed).toBe(false);
176
+ expect(r.code).toBe("user_daily_spend_limit");
177
+ });
178
+
179
+ test("admits when it fits, then settles to the real per-image cost", async () => {
180
+ const t = convexTest(schema, modules);
181
+ await setUserLimits(t, "u", { dailySpendLimitNanos: 500_000_000 });
182
+ const r = await start(t, {
183
+ userId: "u",
184
+ model: "openai/gpt-image-1",
185
+ estimatedCostNanos: 130_000_000,
186
+ });
187
+ expect(r.allowed).toBe(true);
188
+ await settleWith(t, r.requestId, { costNanos: 130_000_000 });
189
+ const u = await userOf(t, "u");
190
+ expect(u.totalSpendNanos).toBe(130_000_000);
191
+ expect(u.reservedTotalNanos ?? 0).toBe(0);
192
+ });
193
+ });
194
+
195
+ describe("async lifecycle (video jobs): begin now, settle later", () => {
196
+ test("reserveTtlMs is stored, and settle records the real cost", async () => {
197
+ const t = convexTest(schema, modules);
198
+ await setUserLimits(t, "u", { dailySpendLimitNanos: 5_000_000_000 });
199
+ // Reserve $2 for a long job that will settle minutes later.
200
+ const r = await start(t, {
201
+ userId: "u",
202
+ model: "openai/sora",
203
+ estimatedCostNanos: 2_000_000_000,
204
+ reserveTtlMs: 30 * 60 * 1000,
205
+ });
206
+ expect(r.allowed).toBe(true);
207
+ const pending = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
208
+ expect(pending.status).toBe("pending");
209
+ expect(pending.reserveTtlMs).toBe(30 * 60 * 1000);
210
+ // held reservation
211
+ let u = await userOf(t, "u");
212
+ expect(u.reservedTotalNanos).toBe(2_000_000_000);
213
+
214
+ // …later: the webhook fires and settles the actual cost.
215
+ await settleWith(t, r.requestId, { costNanos: 1_800_000_000 });
216
+ u = await userOf(t, "u");
217
+ expect(u.totalSpendNanos).toBe(1_800_000_000);
218
+ expect(u.reservedTotalNanos ?? 0).toBe(0);
219
+ });
220
+ });
221
+
164
222
  describe("durable usage history", () => {
165
223
  test("settled spend lands in a per-day usage row", async () => {
166
224
  const t = convexTest(schema, modules);
@@ -419,6 +419,15 @@ export const startRequest = mutation({
419
419
  tags: v.optional(v.array(vTag)),
420
420
  model: v.string(),
421
421
  messages: v.array(vMessage),
422
+ // Reserve this exact amount (nanodollars) instead of the token-based
423
+ // estimate. Use it whenever the cost is known up front — image generation
424
+ // (n × per-image), audio, per-call APIs — so a hard cap reserves the real
425
+ // amount rather than a meaningless token guess.
426
+ estimatedCostNanos: v.optional(v.number()),
427
+ // Hold the reservation this long (ms) before the reconciler may reap it —
428
+ // for long async jobs (video) that settle minutes later. Extends the 30-min
429
+ // default floor.
430
+ reserveTtlMs: v.optional(v.number()),
422
431
  rerunOf: v.optional(v.id("requests")),
423
432
  },
424
433
  returns: vStartResult,
@@ -449,6 +458,11 @@ export const startRequest = mutation({
449
458
  const month = monthStamp();
450
459
  const priceInfo = await getPrice(ctx, args.model);
451
460
  const est = estimateUsage(args.messages, priceInfo);
461
+ // A caller-supplied known cost (image gen, audio, per-call APIs) reserves
462
+ // the real amount up front; the token estimate stays as the token reserve.
463
+ if (args.estimatedCostNanos !== undefined && args.estimatedCostNanos >= 0) {
464
+ est.cost = Math.round(args.estimatedCostNanos);
465
+ }
452
466
  const warnings: string[] = [];
453
467
  const notices: string[] = [];
454
468
 
@@ -682,6 +696,7 @@ export const startRequest = mutation({
682
696
  model: args.model,
683
697
  messages: args.messages,
684
698
  rerunOf: args.rerunOf,
699
+ ...(args.reserveTtlMs !== undefined ? { reserveTtlMs: args.reserveTtlMs } : {}),
685
700
  status: "pending",
686
701
  estimatedNanos: est.cost,
687
702
  estimatedTokens: est.tokens,
@@ -868,14 +883,20 @@ export const reconcile = internalMutation({
868
883
  .take(200);
869
884
  for (const req of toFold) await foldOne(ctx, req);
870
885
 
886
+ // Reap dead reservations: pending rows older than the default floor, but a
887
+ // per-request reserveTtlMs (set for long async jobs like video) holds the
888
+ // reservation until *its* deadline so a still-running job isn't reaped.
871
889
  const cutoff = Date.now() - STALE_PENDING_MS;
872
- const stale = await ctx.db
890
+ const candidates = await ctx.db
873
891
  .query("requests")
874
892
  .withIndex("status", (q) =>
875
893
  q.eq("status", "pending").lt("_creationTime", cutoff)
876
894
  )
877
895
  .take(200);
878
- for (const req of stale) {
896
+ let expired = 0;
897
+ for (const req of candidates) {
898
+ const ttl = req.reserveTtlMs ?? STALE_PENDING_MS;
899
+ if (Date.now() - req._creationTime <= ttl) continue; // still within its window
879
900
  await ctx.db.patch(req._id, {
880
901
  status: "error",
881
902
  error: "Timed out before settling; reservation released",
@@ -883,6 +904,7 @@ export const reconcile = internalMutation({
883
904
  settled: false,
884
905
  });
885
906
  await foldOne(ctx, await ctx.db.get(req._id));
907
+ expired++;
886
908
  }
887
909
 
888
910
  // Retention: delete terminal, fully-accounted request rows past the window.
@@ -907,7 +929,7 @@ export const reconcile = internalMutation({
907
929
  }
908
930
  }
909
931
  }
910
- return { folded: toFold.length, expired: stale.length, purged };
932
+ return { folded: toFold.length, expired, purged };
911
933
  },
912
934
  });
913
935
 
@@ -135,6 +135,11 @@ export default defineSchema({
135
135
  // server-side tool invocations that bill a per-call fee on top of tokens
136
136
  // (e.g. { web_search: 3 }). Priced via serverToolPrices at settle.
137
137
  serverToolUses: v.optional(v.record(v.string(), v.number())),
138
+ // How long the reservation may stay held before the reconciler reaps it as
139
+ // dead (ms). For long async jobs (video generation) set this to the job's
140
+ // max duration so the hold isn't released mid-flight. Extends the default
141
+ // 30-min floor; only stored while pending.
142
+ reserveTtlMs: v.optional(v.number()),
138
143
  costNanos: v.optional(v.number()),
139
144
  latencyMs: v.optional(v.number()),
140
145
  rerunOf: v.optional(v.id("requests")),