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

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.
@@ -6,7 +6,7 @@ import {
6
6
  type MutationCtx,
7
7
  } from "./_generated/server";
8
8
  import { api, internal, components } from "./_generated/api";
9
- import { vMessage } from "./schema";
9
+ import { vMessage, vTag } from "./schema";
10
10
  import type { Doc } from "./_generated/dataModel";
11
11
  import { ShardedCounter } from "@convex-dev/sharded-counter";
12
12
 
@@ -19,6 +19,13 @@ import { ShardedCounter } from "@convex-dev/sharded-counter";
19
19
  const NANOS_PER_DOLLAR = 1e9;
20
20
  const fmtUsd = (nanos: number) => `$${(nanos / NANOS_PER_DOLLAR).toFixed(4)}`;
21
21
 
22
+ // Built-in attribution dimensions. `user` and `action` are always populated
23
+ // from a request's userId/actionName; apps can add any other dimensions
24
+ // (team, project, customer, env, …) as tags. These two names are reserved —
25
+ // tags carrying them are ignored in favor of the first-class fields.
26
+ const USER_DIM = "user";
27
+ const ACTION_DIM = "action";
28
+
22
29
  // Deployment-wide spend totals (nanodollars), sharded for high write throughput.
23
30
  // Keyed "total" (lifetime) and "day:<UTC date>" (natural daily reset).
24
31
  const globalSpend = new ShardedCounter(components.shardedCounter);
@@ -36,9 +43,9 @@ const DEFAULT_PRICES: Record<string, { input: number; output: number }> = {
36
43
  "openai/gpt-5-mini": { input: 250_000_000, output: 2_000_000_000 },
37
44
  };
38
45
 
39
- // Pessimistic assumed output length when reserving budget up front. Reservations
40
- // only need to be an upper bound often enough to keep a hard cap honest; the
41
- // actual cost replaces the estimate the moment the request settles.
46
+ // Pessimistic assumed output length when reserving budget up front. This makes
47
+ // concurrent admission atomic against the estimate; a response that exceeds the
48
+ // estimate can still settle above the cap by the estimation delta.
42
49
  const ESTIMATED_OUTPUT_TOKENS = 800;
43
50
  // A request still "pending" after this long is presumed dead (its action
44
51
  // crashed before settling); the reconciler releases its reservation. Set well
@@ -53,7 +60,13 @@ const STALE_PENDING_MS = 30 * 60 * 1000;
53
60
  // Override per-deployment via setRetention.
54
61
  const DEFAULT_RETENTION_MS = 60 * 60 * 1000; // 1 hour
55
62
 
56
- const dayStamp = () => new Date().toISOString().slice(0, 10);
63
+ const dayStamp = () => new Date().toISOString().slice(0, 10); // "2026-09-04"
64
+ const monthStamp = () => new Date().toISOString().slice(0, 7); // "2026-09"
65
+
66
+ // Cached (prompt-cache-read) input tokens are billed far below the normal input
67
+ // rate. When a model's price has no explicit cachedNanosPerMTok, charge this
68
+ // fraction of its input rate (providers commonly discount ~90%).
69
+ const CACHE_DISCOUNT = 0.1;
57
70
 
58
71
  // Conservative fallback for any model not in the price table: the max of every
59
72
  // known price dimension. Falling back to 0 would be fail-open — an unpriced
@@ -78,14 +91,21 @@ async function getPrice(ctx: MutationCtx, model: string) {
78
91
  return {
79
92
  input: override.inputNanosPerMTok,
80
93
  output: override.outputNanosPerMTok,
94
+ cached: override.cachedNanosPerMTok,
81
95
  known: true,
82
96
  };
83
97
  }
84
98
  const known = DEFAULT_PRICES[model];
85
- if (known) return { ...known, known: true };
86
- return { ...CONSERVATIVE_PRICE, known: false };
99
+ if (known) return { ...known, cached: undefined, known: true };
100
+ return { ...CONSERVATIVE_PRICE, cached: undefined, known: false };
87
101
  }
88
102
 
103
+ type Price = { input: number; output: number; cached?: number };
104
+ // The per-Mtok rate for cached (prompt-cache-read) tokens: an explicit override,
105
+ // else a discount off the input rate.
106
+ const cachedRate = (p: Price) =>
107
+ p.cached ?? Math.round(p.input * CACHE_DISCOUNT);
108
+
89
109
  // Integer nanodollars. Divide-before-multiply keeps the intermediate product
90
110
  // within 2^53 even for large token counts × large per-Mtok prices.
91
111
  const costOf = (
@@ -97,6 +117,66 @@ const costOf = (
97
117
  (inputTokens / 1e6) * price.input + (outputTokens / 1e6) * price.output
98
118
  );
99
119
 
120
+ // Cache-aware settle cost: the cached slice of the prompt is billed at the
121
+ // (discounted) cache rate, the rest of the prompt at the input rate, and
122
+ // completions at the output rate. `cachedTokens` is the gateway's real
123
+ // prompt-cache-read count (usage.inputTokenDetails.cacheReadTokens).
124
+ const settleCost = (
125
+ promptTokens: number,
126
+ cachedTokens: number,
127
+ completionTokens: number,
128
+ price: Price
129
+ ) => {
130
+ const cached = Math.min(Math.max(0, cachedTokens), Math.max(0, promptTokens));
131
+ const fresh = Math.max(0, promptTokens - cached);
132
+ return Math.round(
133
+ (fresh / 1e6) * price.input +
134
+ (cached / 1e6) * cachedRate(price) +
135
+ (completionTokens / 1e6) * price.output
136
+ );
137
+ };
138
+
139
+ // Upsert-add a settled amount into the durable per-(bucket, period) usage row.
140
+ // These rows are never swept by request retention, so spend history survives.
141
+ async function addUsage(
142
+ ctx: MutationCtx,
143
+ dimension: string,
144
+ value: string,
145
+ period: "day" | "month",
146
+ stamp: string,
147
+ spendNanos: number,
148
+ tokens: number,
149
+ requests: number
150
+ ) {
151
+ const existing = await ctx.db
152
+ .query("usage")
153
+ .withIndex("bucket_period_stamp", (q) =>
154
+ q
155
+ .eq("dimension", dimension)
156
+ .eq("value", value)
157
+ .eq("period", period)
158
+ .eq("stamp", stamp)
159
+ )
160
+ .unique();
161
+ if (existing) {
162
+ await ctx.db.patch(existing._id, {
163
+ spendNanos: existing.spendNanos + spendNanos,
164
+ tokens: existing.tokens + tokens,
165
+ requests: existing.requests + requests,
166
+ });
167
+ } else {
168
+ await ctx.db.insert("usage", {
169
+ dimension,
170
+ value,
171
+ period,
172
+ stamp,
173
+ spendNanos,
174
+ tokens,
175
+ requests,
176
+ });
177
+ }
178
+ }
179
+
100
180
  // Up-front estimate of a request's cost and token count, reserved before the
101
181
  // call so concurrent in-flight requests are visible to each other's caps.
102
182
  function estimateUsage(
@@ -111,58 +191,116 @@ function estimateUsage(
111
191
  };
112
192
  }
113
193
 
114
- // Evaluate an entity's (user or action) spend + token budgets against
115
- // committed + reserved + this request's estimate. Returns a hard rejection
116
- // (block) or a list of soft warnings (allow), per the entity's enforcement.
194
+ // The full set of attribution buckets a request touches: the built-in `user`
195
+ // and `action` dimensions plus any extra tags. Reserved dimensions in `extra`
196
+ // are dropped (userId/actionName own them), and (dimension, value) pairs are
197
+ // de-duplicated. Used identically at reserve time (startRequest) and settle
198
+ // time (foldOne), so a request always settles exactly the buckets it reserved.
199
+ function requestBuckets(
200
+ userId: string,
201
+ actionName: string | undefined,
202
+ extra: { dimension: string; value: string }[] | undefined
203
+ ): { dimension: string; value: string }[] {
204
+ const out = [{ dimension: USER_DIM, value: userId }];
205
+ if (actionName !== undefined)
206
+ out.push({ dimension: ACTION_DIM, value: actionName });
207
+ for (const t of extra ?? []) {
208
+ if (t.dimension === USER_DIM || t.dimension === ACTION_DIM) continue;
209
+ if (!t.dimension || !t.value) continue;
210
+ if (out.some((x) => x.dimension === t.dimension && x.value === t.value))
211
+ continue;
212
+ out.push({ dimension: t.dimension, value: t.value });
213
+ }
214
+ return out;
215
+ }
216
+
217
+ // Drop reserved/empty/duplicate tags from a caller-supplied list, leaving the
218
+ // "extra" dimensions stored on the request row.
219
+ function sanitizeExtraTags(
220
+ extra: { dimension: string; value: string }[] | undefined
221
+ ): { dimension: string; value: string }[] {
222
+ const out: { dimension: string; value: string }[] = [];
223
+ for (const t of extra ?? []) {
224
+ if (t.dimension === USER_DIM || t.dimension === ACTION_DIM) continue;
225
+ if (!t.dimension || !t.value) continue;
226
+ if (out.some((x) => x.dimension === t.dimension && x.value === t.value))
227
+ continue;
228
+ out.push({ dimension: t.dimension, value: t.value });
229
+ }
230
+ return out;
231
+ }
232
+
233
+ // Evaluate a bucket's spend + token budgets against committed + reserved + this
234
+ // request's estimate. Returns a hard rejection (block), soft warnings (allow),
235
+ // and threshold notices (approaching a cap — see warnAtPct). Each window (daily,
236
+ // monthly, lifetime) × kind (spend, token) is one check.
117
237
  function evaluateCaps(o: {
118
238
  label: string;
119
239
  name: string;
120
240
  enforcement: "hard" | "soft";
241
+ warnAtPct?: number;
121
242
  estCost: number;
122
243
  estTokens: number;
123
244
  spendToday: number;
124
245
  reservedSpendToday: number;
246
+ spendThisMonth: number;
247
+ reservedSpendMonth: number;
125
248
  totalSpend: number;
126
249
  reservedSpendTotal: number;
127
250
  tokensToday: number;
128
251
  reservedTokensToday: number;
252
+ tokensThisMonth: number;
253
+ reservedTokensMonth: number;
129
254
  totalTokens: number;
130
255
  reservedTokensTotal: number;
131
256
  dailySpendLimitNanos?: number;
257
+ monthlySpendLimitNanos?: number;
132
258
  lifetimeSpendLimitNanos?: number;
133
259
  dailyTokenLimit?: number;
260
+ monthlyTokenLimit?: number;
134
261
  lifetimeTokenLimit?: number;
135
- }): { hard?: { code: string; reason: string }; warnings: string[] } {
262
+ }): {
263
+ hard?: { code: string; reason: string };
264
+ warnings: string[];
265
+ notices: string[];
266
+ } {
267
+ // { code, projected usage (incl. this estimate), cap, human window label,
268
+ // whether it's a money cap (formatted as $), spend? for notices }
269
+ const checks = [
270
+ { w: "daily_spend_limit", used: o.spendToday + o.reservedSpendToday + o.estCost, cap: o.dailySpendLimitNanos, label: "daily spend limit", money: true },
271
+ { w: "monthly_spend_limit", used: o.spendThisMonth + o.reservedSpendMonth + o.estCost, cap: o.monthlySpendLimitNanos, label: "monthly spend limit", money: true },
272
+ { w: "lifetime_spend_limit", used: o.totalSpend + o.reservedSpendTotal + o.estCost, cap: o.lifetimeSpendLimitNanos, label: "lifetime spend limit", money: true },
273
+ { w: "daily_token_limit", used: o.tokensToday + o.reservedTokensToday + o.estTokens, cap: o.dailyTokenLimit, label: "daily token limit", money: false },
274
+ { w: "monthly_token_limit", used: o.tokensThisMonth + o.reservedTokensMonth + o.estTokens, cap: o.monthlyTokenLimit, label: "monthly token limit", money: false },
275
+ { w: "lifetime_token_limit", used: o.totalTokens + o.reservedTokensTotal + o.estTokens, cap: o.lifetimeTokenLimit, label: "lifetime token limit", money: false },
276
+ ];
277
+
136
278
  const violations: { code: string; reason: string }[] = [];
137
- const push = (code: string, reason: string) =>
138
- violations.push({ code: `${o.label}_${code}`, reason });
139
-
140
- if (
141
- o.dailySpendLimitNanos !== undefined &&
142
- o.spendToday + o.reservedSpendToday + o.estCost > o.dailySpendLimitNanos
143
- )
144
- push("daily_spend_limit", `Daily spend limit reached for ${o.label} "${o.name}" (${fmtUsd(o.dailySpendLimitNanos)}/day)`);
145
- if (
146
- o.lifetimeSpendLimitNanos !== undefined &&
147
- o.totalSpend + o.reservedSpendTotal + o.estCost > o.lifetimeSpendLimitNanos
148
- )
149
- push("lifetime_spend_limit", `Lifetime spend limit reached for ${o.label} "${o.name}"`);
150
- if (
151
- o.dailyTokenLimit !== undefined &&
152
- o.tokensToday + o.reservedTokensToday + o.estTokens > o.dailyTokenLimit
153
- )
154
- push("daily_token_limit", `Daily token limit reached for ${o.label} "${o.name}" (${o.dailyTokenLimit}/day)`);
155
- if (
156
- o.lifetimeTokenLimit !== undefined &&
157
- o.totalTokens + o.reservedTokensTotal + o.estTokens > o.lifetimeTokenLimit
158
- )
159
- push("lifetime_token_limit", `Lifetime token limit reached for ${o.label} "${o.name}"`);
160
-
161
- if (violations.length === 0) return { warnings: [] };
162
- if (o.enforcement === "soft") return { warnings: violations.map((v) => v.reason) };
163
- return { hard: violations[0], warnings: [] };
279
+ const notices: string[] = [];
280
+ const pct = o.warnAtPct;
281
+ for (const c of checks) {
282
+ if (c.cap === undefined) continue;
283
+ const capStr = c.money ? `${fmtUsd(c.cap)}` : `${c.cap} tokens`;
284
+ if (c.used > c.cap) {
285
+ violations.push({
286
+ code: `${o.label}_${c.w}`,
287
+ reason: `${cap(c.label)} reached for ${o.label} "${o.name}" (${capStr})`,
288
+ });
289
+ } else if (pct !== undefined && pct > 0 && pct < 1 && c.used >= pct * c.cap) {
290
+ notices.push(
291
+ `${o.label} "${o.name}" at ${Math.round((c.used / c.cap) * 100)}% of ${c.label} (${capStr})`
292
+ );
293
+ }
294
+ }
295
+
296
+ if (violations.length === 0) return { warnings: [], notices };
297
+ if (o.enforcement === "soft")
298
+ return { warnings: violations.map((v) => v.reason), notices };
299
+ return { hard: violations[0], warnings: [], notices };
164
300
  }
165
301
 
302
+ const cap = (s: string) => s.charAt(0).toUpperCase() + s.slice(1);
303
+
166
304
  // A cap plus any one-time bump. Returns undefined when there's no base cap
167
305
  // (a bump alone never creates a cap).
168
306
  const withBump = (base: number | undefined, bump: number | undefined) =>
@@ -170,40 +308,43 @@ const withBump = (base: number | undefined, bump: number | undefined) =>
170
308
 
171
309
  const hasAnyCap = (e: {
172
310
  dailySpendLimitNanos?: number;
311
+ monthlySpendLimitNanos?: number;
173
312
  lifetimeSpendLimitNanos?: number;
174
313
  dailyTokenLimit?: number;
314
+ monthlyTokenLimit?: number;
175
315
  lifetimeTokenLimit?: number;
176
316
  }) =>
177
317
  e.dailySpendLimitNanos !== undefined ||
318
+ e.monthlySpendLimitNanos !== undefined ||
178
319
  e.lifetimeSpendLimitNanos !== undefined ||
179
320
  e.dailyTokenLimit !== undefined ||
321
+ e.monthlyTokenLimit !== undefined ||
180
322
  e.lifetimeTokenLimit !== undefined;
181
323
 
182
- async function getOrCreateUser(ctx: MutationCtx, userId: string) {
183
- const existing = await ctx.db
184
- .query("users")
185
- .withIndex("userId", (q) => q.eq("userId", userId))
324
+ // A bucket needs a reservation row-write if it has any spend/token cap OR a
325
+ // concurrency cap (which reads pendingCount, incremented at reserve time).
326
+ const needsReserve = (e: Parameters<typeof hasAnyCap>[0] & { maxConcurrent?: number }) =>
327
+ hasAnyCap(e) || e.maxConcurrent !== undefined;
328
+
329
+ async function getBucketDoc(ctx: MutationCtx, dimension: string, value: string) {
330
+ return await ctx.db
331
+ .query("buckets")
332
+ .withIndex("dim_value", (q) =>
333
+ q.eq("dimension", dimension).eq("value", value)
334
+ )
186
335
  .unique();
187
- if (existing) return existing;
188
- const id = await ctx.db.insert("users", {
189
- userId,
190
- totalSpendNanos: 0,
191
- totalRequests: 0,
192
- totalTokens: 0,
193
- dayStamp: dayStamp(),
194
- spendTodayNanos: 0,
195
- });
196
- return (await ctx.db.get(id))!;
197
336
  }
198
337
 
199
- async function getOrCreateAction(ctx: MutationCtx, name: string) {
200
- const existing = await ctx.db
201
- .query("actions")
202
- .withIndex("name", (q) => q.eq("name", name))
203
- .unique();
338
+ async function getOrCreateBucket(
339
+ ctx: MutationCtx,
340
+ dimension: string,
341
+ value: string
342
+ ) {
343
+ const existing = await getBucketDoc(ctx, dimension, value);
204
344
  if (existing) return existing;
205
- const id = await ctx.db.insert("actions", {
206
- name,
345
+ const id = await ctx.db.insert("buckets", {
346
+ dimension,
347
+ value,
207
348
  totalSpendNanos: 0,
208
349
  totalRequests: 0,
209
350
  totalTokens: 0,
@@ -220,11 +361,22 @@ async function getSettings(ctx: MutationCtx) {
220
361
  .unique();
221
362
  }
222
363
 
364
+ // Delete the reverse-index rows for a request (called when the request row is
365
+ // deleted, so the tag index never outlives the request it points at).
366
+ async function deleteRequestTags(ctx: MutationCtx, requestId: Doc<"requests">["_id"]) {
367
+ const tags = await ctx.db
368
+ .query("requestTags")
369
+ .withIndex("requestId", (q) => q.eq("requestId", requestId))
370
+ .collect();
371
+ for (const t of tags) await ctx.db.delete(t._id);
372
+ }
373
+
223
374
  const vStartResult = v.union(
224
375
  v.object({
225
376
  allowed: v.literal(true),
226
377
  requestId: v.id("requests"),
227
- warnings: v.array(v.string()),
378
+ warnings: v.array(v.string()), // soft caps exceeded (allowed with a warning)
379
+ notices: v.array(v.string()), // approaching a cap (warnAtPct threshold)
228
380
  }),
229
381
  v.object({
230
382
  allowed: v.literal(false),
@@ -237,13 +389,16 @@ export const startRequest = mutation({
237
389
  args: {
238
390
  userId: v.string(),
239
391
  actionName: v.optional(v.string()),
392
+ // Extra attribution dimensions to bill/limit (team, customer, env, …).
393
+ // `user` and `action` are reserved (owned by userId/actionName).
394
+ tags: v.optional(v.array(vTag)),
240
395
  model: v.string(),
241
396
  messages: v.array(vMessage),
242
397
  rerunOf: v.optional(v.id("requests")),
243
398
  },
244
399
  returns: vStartResult,
245
400
  handler: async (ctx, args) => {
246
- const user = await getOrCreateUser(ctx, args.userId);
401
+ const extraTags = sanitizeExtraTags(args.tags);
247
402
  // Record the blocked attempt and return a rejection (throwing would roll
248
403
  // back the record). `persist` is false for the high-frequency-by-design
249
404
  // rejections (rate limit, blocked user) that a client retries in a tight
@@ -252,7 +407,12 @@ export const startRequest = mutation({
252
407
  const reject = async (code: string, reason: string, persist = true) => {
253
408
  if (persist) {
254
409
  await ctx.db.insert("requests", {
255
- ...args,
410
+ userId: args.userId,
411
+ actionName: args.actionName,
412
+ ...(extraTags.length ? { tags: extraTags } : {}),
413
+ model: args.model,
414
+ messages: args.messages,
415
+ rerunOf: args.rerunOf,
256
416
  status: "blocked" as const,
257
417
  error: reason,
258
418
  });
@@ -261,18 +421,18 @@ export const startRequest = mutation({
261
421
  };
262
422
 
263
423
  const today = dayStamp();
424
+ const month = monthStamp();
264
425
  const priceInfo = await getPrice(ctx, args.model);
265
426
  const est = estimateUsage(args.messages, priceInfo);
266
427
  const warnings: string[] = [];
428
+ const notices: string[] = [];
267
429
 
268
- if (user.blocked) {
269
- return reject("blocked", `User "${args.userId}" is blocked`, false);
270
- }
271
430
  // Model allow/deny policy (component-wide).
272
- const policy = await getSettings(ctx);
273
- if (policy) {
274
- const mode = policy.modelMode ?? "open";
275
- const list = policy.models ?? [];
431
+ const settings = await getSettings(ctx);
432
+ const defaultWarnAtPct = settings?.defaultWarnAtPct;
433
+ if (settings) {
434
+ const mode = settings.modelMode ?? "open";
435
+ const list = settings.models ?? [];
276
436
  if (mode === "allowlist" && !list.includes(args.model)) {
277
437
  return reject(
278
438
  "model_not_allowed",
@@ -283,178 +443,236 @@ export const startRequest = mutation({
283
443
  return reject("model_denied", `Model "${args.model}" is denied`);
284
444
  }
285
445
  }
286
- if (user.requestsPerMinute !== undefined) {
287
- // Bounded read: we only need to know whether non-blocked requests in the
288
- // window have reached the limit. Reading limit+1 non-blocked rows is
289
- // enough, and .take caps the scan even if the window is flooded. Blocked
290
- // rows aren't persisted for rate-limit/blocked rejections (see reject),
291
- // so the window stays small.
292
- const recent = await ctx.db
293
- .query("requests")
294
- .withIndex("userId", (q) =>
295
- q.eq("userId", args.userId).gt("_creationTime", Date.now() - 60_000)
296
- )
297
- .take(user.requestsPerMinute + 50);
298
- if (
299
- recent.filter((r) => r.status !== "blocked").length >=
300
- user.requestsPerMinute
301
- ) {
446
+
447
+ // Fetch/create every bucket this request is attributed to (user, action,
448
+ // and any extra tags). Each may carry its own budget.
449
+ const bucketTags = requestBuckets(args.userId, args.actionName, extraTags);
450
+ const buckets: Doc<"buckets">[] = [];
451
+ for (const t of bucketTags) {
452
+ buckets.push(await getOrCreateBucket(ctx, t.dimension, t.value));
453
+ }
454
+
455
+ // A hard block on ANY bucket rejects the request. The user dimension's block
456
+ // isn't persisted (retried in a loop); config-level blocks on other
457
+ // dimensions are rarer, so they persist for the audit log.
458
+ for (const b of buckets) {
459
+ if (b.blocked) {
460
+ const label = b.dimension === USER_DIM ? "User" : b.dimension;
302
461
  return reject(
303
- "rate_limit",
304
- `Rate limit exceeded for "${args.userId}" (${user.requestsPerMinute}/min)`,
462
+ `${b.dimension}_blocked`,
463
+ `${label} "${b.value}" is blocked`,
464
+ b.dimension !== USER_DIM
465
+ );
466
+ }
467
+ }
468
+
469
+ // Concurrency cap: reject when a bucket already has maxConcurrent requests
470
+ // in flight (pendingCount). A transient limit like rate-limiting, so it
471
+ // isn't persisted (the caller retries once something settles).
472
+ for (const b of buckets) {
473
+ if (b.maxConcurrent !== undefined && (b.pendingCount ?? 0) >= b.maxConcurrent) {
474
+ return reject(
475
+ `${b.dimension}_max_concurrent`,
476
+ `Too many concurrent requests for ${b.dimension} "${b.value}" (max ${b.maxConcurrent})`,
305
477
  false
306
478
  );
307
479
  }
308
480
  }
309
- // Committed + already-reserved in-flight usage + this estimate must fit
310
- // under each cap. Decided and reserved in one transaction, so Convex's
311
- // serializable isolation makes it a true atomic check-and-reserve. Soft
312
- // enforcement turns a violation into a warning instead of a block.
313
- const userSameDay = user.dayStamp === today;
314
- const userEval = evaluateCaps({
315
- label: "user",
316
- name: args.userId,
317
- enforcement: user.enforcement ?? "hard",
318
- estCost: est.cost,
319
- estTokens: est.tokens,
320
- spendToday: userSameDay ? user.spendTodayNanos : 0,
321
- reservedSpendToday: userSameDay ? user.reservedTodayNanos ?? 0 : 0,
322
- totalSpend: user.totalSpendNanos,
323
- reservedSpendTotal: user.reservedTotalNanos ?? 0,
324
- tokensToday: userSameDay ? user.tokensToday ?? 0 : 0,
325
- reservedTokensToday: userSameDay ? user.reservedTodayTokens ?? 0 : 0,
326
- totalTokens: user.totalTokens,
327
- reservedTokensTotal: user.reservedTotalTokens ?? 0,
328
- dailySpendLimitNanos: withBump(
329
- user.dailySpendLimitNanos,
330
- user.bumpDayStamp === today ? user.dailyBumpNanos : 0
331
- ),
332
- lifetimeSpendLimitNanos: withBump(
333
- user.lifetimeSpendLimitNanos,
334
- user.lifetimeBumpNanos
335
- ),
336
- dailyTokenLimit: user.dailyTokenLimit,
337
- lifetimeTokenLimit: user.lifetimeTokenLimit,
338
- });
339
- if (userEval.hard) return reject(userEval.hard.code, userEval.hard.reason);
340
- warnings.push(...userEval.warnings);
341
-
342
- let action: Doc<"actions"> | null = null;
343
- if (args.actionName !== undefined) {
344
- action = await getOrCreateAction(ctx, args.actionName);
345
- if (action.disabled) {
346
- return reject("action_disabled", `Action "${action.name}" is disabled`);
481
+
482
+ // Enforce the rolling 60-second rate limit on every configured dimension.
483
+ // User/action requests use their first-class indexes; custom dimensions use
484
+ // the requestTags reverse index. All reads are bounded by the configured
485
+ // limit (plus a small allowance for persisted blocked attempts).
486
+ const rateCutoff = Date.now() - 60_000;
487
+ for (const b of buckets) {
488
+ const limit = b.requestsPerMinute;
489
+ if (limit === undefined) continue;
490
+
491
+ let recentCount: number;
492
+ if (b.dimension === USER_DIM) {
493
+ const recent = await ctx.db
494
+ .query("requests")
495
+ .withIndex("userId", (q) =>
496
+ q.eq("userId", b.value).gt("_creationTime", rateCutoff)
497
+ )
498
+ .take(limit + 50);
499
+ recentCount = recent.filter((r) => r.status !== "blocked").length;
500
+ } else if (b.dimension === ACTION_DIM) {
501
+ const recent = await ctx.db
502
+ .query("requests")
503
+ .withIndex("actionName", (q) =>
504
+ q.eq("actionName", b.value).gt("_creationTime", rateCutoff)
505
+ )
506
+ .take(limit + 50);
507
+ recentCount = recent.filter((r) => r.status !== "blocked").length;
508
+ } else {
509
+ recentCount = (
510
+ await ctx.db
511
+ .query("requestTags")
512
+ .withIndex("dim_value", (q) =>
513
+ q
514
+ .eq("dimension", b.dimension)
515
+ .eq("value", b.value)
516
+ .gt("_creationTime", rateCutoff)
517
+ )
518
+ .take(limit)
519
+ ).length;
347
520
  }
348
- const aSameDay = action.dayStamp === today;
349
- const actionEval = evaluateCaps({
350
- label: "action",
351
- name: action.name,
352
- enforcement: action.enforcement ?? "hard",
521
+
522
+ if (recentCount >= limit) {
523
+ const code = b.dimension === USER_DIM ? "rate_limit" : `${b.dimension}_rate_limit`;
524
+ return reject(
525
+ code,
526
+ `Rate limit exceeded for ${b.dimension} "${b.value}" (${limit}/min)`,
527
+ false
528
+ );
529
+ }
530
+ }
531
+
532
+ // Committed + already-reserved in-flight usage + this estimate must fit
533
+ // under EACH bucket's cap. Decided and reserved in one transaction, so
534
+ // Convex's serializable isolation makes concurrent admission atomic across
535
+ // every capped bucket. Final usage can exceed the estimate; settlement then
536
+ // records the actual amount. Soft enforcement turns a violation into a
537
+ // warning instead of a block.
538
+ for (const b of buckets) {
539
+ const sameDay = b.dayStamp === today;
540
+ const sameMonth = b.monthStamp === month;
541
+ const ev = evaluateCaps({
542
+ label: b.dimension,
543
+ name: b.value,
544
+ enforcement: b.enforcement ?? "hard",
545
+ warnAtPct: b.warnAtPct ?? defaultWarnAtPct,
353
546
  estCost: est.cost,
354
547
  estTokens: est.tokens,
355
- spendToday: aSameDay ? action.spendTodayNanos : 0,
356
- reservedSpendToday: aSameDay ? action.reservedTodayNanos ?? 0 : 0,
357
- totalSpend: action.totalSpendNanos,
358
- reservedSpendTotal: action.reservedTotalNanos ?? 0,
359
- tokensToday: aSameDay ? action.tokensToday ?? 0 : 0,
360
- reservedTokensToday: aSameDay ? action.reservedTodayTokens ?? 0 : 0,
361
- totalTokens: action.totalTokens,
362
- reservedTokensTotal: action.reservedTotalTokens ?? 0,
548
+ spendToday: sameDay ? b.spendTodayNanos : 0,
549
+ reservedSpendToday: sameDay ? b.reservedTodayNanos ?? 0 : 0,
550
+ spendThisMonth: sameMonth ? b.spendThisMonthNanos ?? 0 : 0,
551
+ reservedSpendMonth: sameMonth ? b.reservedMonthNanos ?? 0 : 0,
552
+ totalSpend: b.totalSpendNanos,
553
+ reservedSpendTotal: b.reservedTotalNanos ?? 0,
554
+ tokensToday: sameDay ? b.tokensToday ?? 0 : 0,
555
+ reservedTokensToday: sameDay ? b.reservedTodayTokens ?? 0 : 0,
556
+ tokensThisMonth: sameMonth ? b.tokensThisMonth ?? 0 : 0,
557
+ reservedTokensMonth: sameMonth ? b.reservedMonthTokens ?? 0 : 0,
558
+ totalTokens: b.totalTokens,
559
+ reservedTokensTotal: b.reservedTotalTokens ?? 0,
363
560
  dailySpendLimitNanos: withBump(
364
- action.dailySpendLimitNanos,
365
- action.bumpDayStamp === today ? action.dailyBumpNanos : 0
561
+ b.dailySpendLimitNanos,
562
+ b.bumpDayStamp === today ? b.dailyBumpNanos : 0
563
+ ),
564
+ monthlySpendLimitNanos: withBump(
565
+ b.monthlySpendLimitNanos,
566
+ b.bumpMonthStamp === month ? b.monthlyBumpNanos : 0
366
567
  ),
367
568
  lifetimeSpendLimitNanos: withBump(
368
- action.lifetimeSpendLimitNanos,
369
- action.lifetimeBumpNanos
569
+ b.lifetimeSpendLimitNanos,
570
+ b.lifetimeBumpNanos
370
571
  ),
371
- dailyTokenLimit: action.dailyTokenLimit,
372
- lifetimeTokenLimit: action.lifetimeTokenLimit,
572
+ dailyTokenLimit: b.dailyTokenLimit,
573
+ monthlyTokenLimit: b.monthlyTokenLimit,
574
+ lifetimeTokenLimit: b.lifetimeTokenLimit,
373
575
  });
374
- if (actionEval.hard) return reject(actionEval.hard.code, actionEval.hard.reason);
375
- warnings.push(...actionEval.warnings);
576
+ if (ev.hard) return reject(ev.hard.code, ev.hard.reason);
577
+ warnings.push(...ev.warnings);
578
+ notices.push(...ev.notices);
376
579
  }
377
580
 
378
- // Deployment-wide ("global") spend cap — same reserve-then-settle model and
379
- // the SAME evaluateCaps logic as the per-entity caps above, so the guarantee
380
- // statement is uniform: a request is admitted only if
581
+ // Deployment-wide ("global") spend cap — the same estimated-usage admission
582
+ // rule as the per-bucket caps above: a request is admitted only if
381
583
  // committed + reserved + estimate <= cap. The ONE difference is the holder:
382
- // per-user/per-action reserve on a single row (an exact atomic
383
- // check-and-reserve), while the global holder is a sharded counter for
584
+ // per-bucket caps reserve on a single row (an atomic check-and-reserve),
585
+ // while the global holder is a sharded counter for
384
586
  // throughput — its committed total is read as an eventually-consistent sum
385
587
  // with no cross-request reservation, so a hard global cap can overshoot by a
386
- // bounded amount under burst. That's the deliberate exactness/throughput
588
+ // bounded amount under burst. That's the deliberate consistency/throughput
387
589
  // trade for a deployment-wide killswitch; it's the only approximate scope.
388
590
  if (
389
- policy &&
390
- (policy.globalDailySpendLimitNanos !== undefined ||
391
- policy.globalLifetimeSpendLimitNanos !== undefined)
591
+ settings &&
592
+ (settings.globalDailySpendLimitNanos !== undefined ||
593
+ settings.globalLifetimeSpendLimitNanos !== undefined)
392
594
  ) {
393
595
  const globalEval = evaluateCaps({
394
596
  label: "global",
395
597
  name: "deployment",
396
- enforcement: policy.globalEnforcement ?? "hard",
598
+ enforcement: settings.globalEnforcement ?? "hard",
599
+ warnAtPct: defaultWarnAtPct,
397
600
  estCost: est.cost,
398
601
  estTokens: est.tokens,
399
602
  spendToday: await globalSpend.count(ctx, globalDayKey(today)),
400
603
  reservedSpendToday: 0, // sharded holder: no cross-request reservation
604
+ spendThisMonth: 0, // global tracks daily + lifetime only
605
+ reservedSpendMonth: 0,
401
606
  totalSpend: await globalSpend.count(ctx, GLOBAL_TOTAL),
402
607
  reservedSpendTotal: 0,
403
608
  tokensToday: 0,
404
609
  reservedTokensToday: 0,
610
+ tokensThisMonth: 0,
611
+ reservedTokensMonth: 0,
405
612
  totalTokens: 0,
406
613
  reservedTokensTotal: 0,
407
614
  dailySpendLimitNanos: withBump(
408
- policy.globalDailySpendLimitNanos,
409
- policy.globalBumpDayStamp === today ? policy.globalDailyBumpNanos : 0
615
+ settings.globalDailySpendLimitNanos,
616
+ settings.globalBumpDayStamp === today ? settings.globalDailyBumpNanos : 0
410
617
  ),
411
618
  lifetimeSpendLimitNanos: withBump(
412
- policy.globalLifetimeSpendLimitNanos,
413
- policy.globalLifetimeBumpNanos
619
+ settings.globalLifetimeSpendLimitNanos,
620
+ settings.globalLifetimeBumpNanos
414
621
  ),
415
622
  });
416
623
  if (globalEval.hard) return reject(globalEval.hard.code, globalEval.hard.reason);
417
624
  warnings.push(...globalEval.warnings);
625
+ notices.push(...globalEval.notices);
418
626
  }
419
627
 
420
- // Passed — reserve, but ONLY on entities that actually have a cap. Writing
421
- // an uncapped entity's row here would serialize every request that shares it
422
- // (e.g. all callers of one action); with no cap there's no reserved amount
423
- // to consult. Totals are still accrued later, asynchronously, in foldOne.
424
- if (hasAnyCap(user)) {
425
- await ctx.db.patch(user._id, {
426
- dayStamp: today,
427
- spendTodayNanos: userSameDay ? user.spendTodayNanos : 0,
428
- tokensToday: userSameDay ? user.tokensToday ?? 0 : 0,
429
- reservedTodayNanos: (userSameDay ? user.reservedTodayNanos ?? 0 : 0) + est.cost,
430
- reservedTotalNanos: (user.reservedTotalNanos ?? 0) + est.cost,
431
- reservedTodayTokens: (userSameDay ? user.reservedTodayTokens ?? 0 : 0) + est.tokens,
432
- reservedTotalTokens: (user.reservedTotalTokens ?? 0) + est.tokens,
433
- pendingCount: (user.pendingCount ?? 0) + 1,
434
- });
435
- }
436
- if (action && hasAnyCap(action)) {
437
- const aSameDay = action.dayStamp === today;
438
- await ctx.db.patch(action._id, {
628
+ // Passed — reserve, but ONLY on buckets that actually have a cap. Writing an
629
+ // uncapped bucket's row here would serialize every request that shares it
630
+ // (e.g. all callers of one action, or every request in one env); with no cap
631
+ // there's no reserved amount to consult. Totals are still accrued later,
632
+ // asynchronously, in foldOne — for every bucket, capped or not.
633
+ for (const b of buckets) {
634
+ if (!needsReserve(b)) continue;
635
+ const sameDay = b.dayStamp === today;
636
+ const sameMonth = b.monthStamp === month;
637
+ await ctx.db.patch(b._id, {
439
638
  dayStamp: today,
440
- spendTodayNanos: aSameDay ? action.spendTodayNanos : 0,
441
- tokensToday: aSameDay ? action.tokensToday ?? 0 : 0,
442
- reservedTodayNanos: (aSameDay ? action.reservedTodayNanos ?? 0 : 0) + est.cost,
443
- reservedTotalNanos: (action.reservedTotalNanos ?? 0) + est.cost,
444
- reservedTodayTokens: (aSameDay ? action.reservedTodayTokens ?? 0 : 0) + est.tokens,
445
- reservedTotalTokens: (action.reservedTotalTokens ?? 0) + est.tokens,
446
- pendingCount: (action.pendingCount ?? 0) + 1,
639
+ monthStamp: month,
640
+ spendTodayNanos: sameDay ? b.spendTodayNanos : 0,
641
+ tokensToday: sameDay ? b.tokensToday ?? 0 : 0,
642
+ spendThisMonthNanos: sameMonth ? b.spendThisMonthNanos ?? 0 : 0,
643
+ tokensThisMonth: sameMonth ? b.tokensThisMonth ?? 0 : 0,
644
+ reservedTodayNanos: (sameDay ? b.reservedTodayNanos ?? 0 : 0) + est.cost,
645
+ reservedMonthNanos: (sameMonth ? b.reservedMonthNanos ?? 0 : 0) + est.cost,
646
+ reservedTotalNanos: (b.reservedTotalNanos ?? 0) + est.cost,
647
+ reservedTodayTokens: (sameDay ? b.reservedTodayTokens ?? 0 : 0) + est.tokens,
648
+ reservedMonthTokens: (sameMonth ? b.reservedMonthTokens ?? 0 : 0) + est.tokens,
649
+ reservedTotalTokens: (b.reservedTotalTokens ?? 0) + est.tokens,
650
+ pendingCount: (b.pendingCount ?? 0) + 1,
447
651
  });
448
652
  }
449
653
  const requestId = await ctx.db.insert("requests", {
450
- ...args,
654
+ userId: args.userId,
655
+ actionName: args.actionName,
656
+ ...(extraTags.length ? { tags: extraTags } : {}),
657
+ model: args.model,
658
+ messages: args.messages,
659
+ rerunOf: args.rerunOf,
451
660
  status: "pending",
452
661
  estimatedNanos: est.cost,
453
662
  estimatedTokens: est.tokens,
454
663
  ...(priceInfo.known ? {} : { unpricedModel: true }),
455
664
  ...(warnings.length > 0 ? { overBudget: true } : {}),
456
665
  });
457
- return { allowed: true as const, requestId, warnings };
666
+ // Reverse index so the request log can be filtered by any extra tag
667
+ // dimension (user/action are already indexed on the requests table).
668
+ for (const t of extraTags) {
669
+ await ctx.db.insert("requestTags", {
670
+ dimension: t.dimension,
671
+ value: t.value,
672
+ requestId,
673
+ });
674
+ }
675
+ return { allowed: true as const, requestId, warnings, notices };
458
676
  },
459
677
  });
460
678
 
@@ -466,6 +684,10 @@ export const finishRequest = mutation({
466
684
  promptTokens: v.optional(v.number()),
467
685
  completionTokens: v.optional(v.number()),
468
686
  cachedTokens: v.optional(v.number()),
687
+ // Authoritative cost from the gateway, if it ever reports one. When present
688
+ // it's recorded verbatim (no token-based estimate); when absent we price
689
+ // from tokens (cache-aware). Wired now so adopting a real cost is one line.
690
+ costNanos: v.optional(v.number()),
469
691
  latencyMs: v.optional(v.number()),
470
692
  },
471
693
  returns: v.object({ costNanos: v.number() }),
@@ -489,10 +711,17 @@ export const finishRequest = mutation({
489
711
  const promptTokens = Math.max(0, args.promptTokens ?? 0);
490
712
  const completionTokens = Math.max(0, args.completionTokens ?? 0);
491
713
  const cachedTokens = Math.min(promptTokens, Math.max(0, args.cachedTokens ?? 0));
492
- const costNanos = Math.max(
493
- 0,
494
- costOf(promptTokens, completionTokens, await getPrice(ctx, request.model))
495
- );
714
+ // Prefer an authoritative gateway cost when supplied; otherwise price from
715
+ // tokens, discounting the cached (prompt-cache-read) slice of the prompt.
716
+ const costNanos =
717
+ args.costNanos !== undefined && args.costNanos >= 0
718
+ ? Math.round(args.costNanos)
719
+ : settleCost(
720
+ promptTokens,
721
+ cachedTokens,
722
+ completionTokens,
723
+ await getPrice(ctx, request.model)
724
+ );
496
725
 
497
726
  // Durable write to the request's OWN row only — uncontended, so it always
498
727
  // lands. `settled: false` hands it to the fold step; the row is never left
@@ -509,7 +738,7 @@ export const finishRequest = mutation({
509
738
  settled: false,
510
739
  });
511
740
 
512
- // Fold into the (hot) user/action counters in a separate mutation. If it
741
+ // Fold into the (hot) per-bucket counters in a separate mutation. If it
513
742
  // exhausts retries under contention, the cron reconciler picks it up.
514
743
  await ctx.scheduler.runAfter(0, internal.lib.foldTotals, {
515
744
  requestId: args.requestId,
@@ -518,9 +747,9 @@ export const finishRequest = mutation({
518
747
  },
519
748
  });
520
749
 
521
- // Fold one finished request into the user/action running totals, releasing its
522
- // reservation. Idempotent: guarded by `settled` so the scheduler and the cron
523
- // reconciler can never double-count.
750
+ // Fold one finished request into every attributed bucket's running totals,
751
+ // releasing its reservation. Idempotent: guarded by `settled` so the scheduler
752
+ // and the cron reconciler can never double-count.
524
753
  async function foldOne(ctx: MutationCtx, req: Doc<"requests"> | null) {
525
754
  if (!req || req.settled !== false) return;
526
755
  const actual = req.costNanos ?? 0;
@@ -528,48 +757,44 @@ async function foldOne(ctx: MutationCtx, req: Doc<"requests"> | null) {
528
757
  const tokens = (req.promptTokens ?? 0) + (req.completionTokens ?? 0);
529
758
  const estTokens = req.estimatedTokens ?? 0;
530
759
  const today = dayStamp();
760
+ const month = monthStamp();
531
761
 
532
- const user = await getOrCreateUser(ctx, req.userId);
533
- const uSameDay = user.dayStamp === today;
534
- await ctx.db.patch(user._id, {
535
- totalSpendNanos: user.totalSpendNanos + actual,
536
- totalRequests: user.totalRequests + 1,
537
- totalTokens: user.totalTokens + tokens,
538
- dayStamp: today,
539
- spendTodayNanos: (uSameDay ? user.spendTodayNanos : 0) + actual,
540
- tokensToday: (uSameDay ? user.tokensToday ?? 0 : 0) + tokens,
541
- reservedTodayNanos: Math.max(0, (uSameDay ? user.reservedTodayNanos ?? 0 : 0) - estCost),
542
- reservedTotalNanos: Math.max(0, (user.reservedTotalNanos ?? 0) - estCost),
543
- reservedTodayTokens: Math.max(0, (uSameDay ? user.reservedTodayTokens ?? 0 : 0) - estTokens),
544
- reservedTotalTokens: Math.max(0, (user.reservedTotalTokens ?? 0) - estTokens),
545
- pendingCount: Math.max(0, (user.pendingCount ?? 0) - 1),
546
- });
547
-
548
- if (req.actionName !== undefined) {
549
- const action = await getOrCreateAction(ctx, req.actionName);
550
- const aSameDay = action.dayStamp === today;
551
- await ctx.db.patch(action._id, {
552
- totalSpendNanos: action.totalSpendNanos + actual,
553
- totalRequests: action.totalRequests + 1,
554
- totalTokens: action.totalTokens + tokens,
762
+ // Accrue into every attributed bucket (user, action, and each tag) — capped
763
+ // or not. Buckets that never held a reservation have their reserved fields
764
+ // clamped at 0 by Math.max, so subtracting an estimate is a harmless no-op.
765
+ // Also write the durable per-(bucket, day/month) usage rows that survive
766
+ // request retention, so spend history outlives the raw request log.
767
+ for (const t of requestBuckets(req.userId, req.actionName, req.tags)) {
768
+ const b = await getOrCreateBucket(ctx, t.dimension, t.value);
769
+ const sameDay = b.dayStamp === today;
770
+ const sameMonth = b.monthStamp === month;
771
+ await ctx.db.patch(b._id, {
772
+ totalSpendNanos: b.totalSpendNanos + actual,
773
+ totalRequests: b.totalRequests + 1,
774
+ totalTokens: b.totalTokens + tokens,
555
775
  dayStamp: today,
556
- spendTodayNanos: (aSameDay ? action.spendTodayNanos : 0) + actual,
557
- tokensToday: (aSameDay ? action.tokensToday ?? 0 : 0) + tokens,
558
- reservedTodayNanos: Math.max(0, (aSameDay ? action.reservedTodayNanos ?? 0 : 0) - estCost),
559
- reservedTotalNanos: Math.max(0, (action.reservedTotalNanos ?? 0) - estCost),
560
- reservedTodayTokens: Math.max(0, (aSameDay ? action.reservedTodayTokens ?? 0 : 0) - estTokens),
561
- reservedTotalTokens: Math.max(0, (action.reservedTotalTokens ?? 0) - estTokens),
562
- pendingCount: Math.max(0, (action.pendingCount ?? 0) - 1),
776
+ monthStamp: month,
777
+ spendTodayNanos: (sameDay ? b.spendTodayNanos : 0) + actual,
778
+ tokensToday: (sameDay ? b.tokensToday ?? 0 : 0) + tokens,
779
+ spendThisMonthNanos: (sameMonth ? b.spendThisMonthNanos ?? 0 : 0) + actual,
780
+ tokensThisMonth: (sameMonth ? b.tokensThisMonth ?? 0 : 0) + tokens,
781
+ reservedTodayNanos: Math.max(0, (sameDay ? b.reservedTodayNanos ?? 0 : 0) - estCost),
782
+ reservedMonthNanos: Math.max(0, (sameMonth ? b.reservedMonthNanos ?? 0 : 0) - estCost),
783
+ reservedTotalNanos: Math.max(0, (b.reservedTotalNanos ?? 0) - estCost),
784
+ reservedTodayTokens: Math.max(0, (sameDay ? b.reservedTodayTokens ?? 0 : 0) - estTokens),
785
+ reservedMonthTokens: Math.max(0, (sameMonth ? b.reservedMonthTokens ?? 0 : 0) - estTokens),
786
+ reservedTotalTokens: Math.max(0, (b.reservedTotalTokens ?? 0) - estTokens),
787
+ pendingCount: Math.max(0, (b.pendingCount ?? 0) - 1),
563
788
  });
789
+ await addUsage(ctx, t.dimension, t.value, "day", today, actual, tokens, 1);
790
+ await addUsage(ctx, t.dimension, t.value, "month", month, actual, tokens, 1);
564
791
  }
792
+
565
793
  // Deployment-wide totals via the sharded counter (only when a global cap is
566
794
  // configured — otherwise skip the writes entirely). Distributed across shards,
567
795
  // so this does not serialize on a single row.
568
796
  if (actual > 0) {
569
- const settings = await ctx.db
570
- .query("settings")
571
- .withIndex("key", (q) => q.eq("key", "singleton"))
572
- .unique();
797
+ const settings = await getSettings(ctx);
573
798
  if (
574
799
  settings &&
575
800
  (settings.globalDailySpendLimitNanos !== undefined ||
@@ -641,6 +866,7 @@ export const reconcile = internalMutation({
641
866
  // Only rows that are done and accounted: folded (settled === true) or a
642
867
  // blocked attempt (never needs folding). Never a pending/unfolded row.
643
868
  if (req.settled === true || req.status === "blocked") {
869
+ await deleteRequestTags(ctx, req._id);
644
870
  await ctx.db.delete(req._id);
645
871
  purged++;
646
872
  }
@@ -687,10 +913,19 @@ export const getRequest = query({
687
913
  });
688
914
 
689
915
  export const listRequests = query({
690
- args: { userId: v.optional(v.string()), limit: v.optional(v.number()) },
916
+ args: {
917
+ userId: v.optional(v.string()),
918
+ // Filter by any attribution dimension (user/action indexed on the request;
919
+ // custom tag dimensions resolved via the requestTags reverse index).
920
+ dimension: v.optional(v.string()),
921
+ value: v.optional(v.string()),
922
+ limit: v.optional(v.number()),
923
+ },
691
924
  handler: async (ctx, args) => {
692
925
  const limit = args.limit ?? 50;
693
- if (args.userId !== undefined) {
926
+ const dim = args.dimension;
927
+ const val = args.value ?? args.userId;
928
+ if (dim === undefined && args.userId !== undefined) {
694
929
  const userId = args.userId;
695
930
  return await ctx.db
696
931
  .query("requests")
@@ -698,161 +933,248 @@ export const listRequests = query({
698
933
  .order("desc")
699
934
  .take(limit);
700
935
  }
936
+ if (dim !== undefined && val !== undefined) {
937
+ if (dim === USER_DIM) {
938
+ return await ctx.db
939
+ .query("requests")
940
+ .withIndex("userId", (q) => q.eq("userId", val))
941
+ .order("desc")
942
+ .take(limit);
943
+ }
944
+ if (dim === ACTION_DIM) {
945
+ return await ctx.db
946
+ .query("requests")
947
+ .withIndex("actionName", (q) => q.eq("actionName", val))
948
+ .order("desc")
949
+ .take(limit);
950
+ }
951
+ // Custom tag dimension: walk the reverse index, then fetch each request.
952
+ const tagRows = await ctx.db
953
+ .query("requestTags")
954
+ .withIndex("dim_value", (q) => q.eq("dimension", dim).eq("value", val))
955
+ .order("desc")
956
+ .take(limit);
957
+ const rows = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
958
+ return rows.filter((r): r is Doc<"requests"> => r !== null);
959
+ }
701
960
  return await ctx.db.query("requests").order("desc").take(limit);
702
961
  },
703
962
  });
704
963
 
705
964
  const ADMIN_LIST_CAP = 2000;
706
- export const listUsers = query({
707
- args: {},
708
- handler: async (ctx) => {
965
+
966
+ // List budget buckets, optionally filtered to one dimension ("user", "action",
967
+ // or any custom tag dimension). Today's spend is zeroed for stale day windows.
968
+ export const listBuckets = query({
969
+ args: { dimension: v.optional(v.string()) },
970
+ handler: async (ctx, args) => {
709
971
  // Bounded to avoid an unbounded full-table scan on this reactive query.
710
- // Paginate (ctx.db.query("users").paginate(...)) for larger deployments.
711
- const users = await ctx.db.query("users").take(ADMIN_LIST_CAP);
972
+ // Paginate (ctx.db.query("buckets").paginate(...)) for larger deployments.
973
+ const rows =
974
+ args.dimension !== undefined
975
+ ? await ctx.db
976
+ .query("buckets")
977
+ .withIndex("dimension", (q) => q.eq("dimension", args.dimension!))
978
+ .take(ADMIN_LIST_CAP)
979
+ : await ctx.db.query("buckets").take(ADMIN_LIST_CAP);
712
980
  const today = dayStamp();
713
- return users.map((u) => ({
714
- ...u,
715
- spendTodayNanos: u.dayStamp === today ? u.spendTodayNanos : 0,
981
+ const month = monthStamp();
982
+ return rows.map((b) => ({
983
+ ...b,
984
+ spendTodayNanos: b.dayStamp === today ? b.spendTodayNanos : 0,
985
+ spendThisMonthNanos: b.monthStamp === month ? b.spendThisMonthNanos ?? 0 : 0,
716
986
  }));
717
987
  },
718
988
  });
719
989
 
720
- export const setLimits = mutation({
990
+ export const getBucket = query({
991
+ args: { dimension: v.string(), value: v.string() },
992
+ handler: async (ctx, args) => {
993
+ const b = await getBucketDoc(ctx as any, args.dimension, args.value);
994
+ if (!b) return null;
995
+ const today = dayStamp();
996
+ const month = monthStamp();
997
+ return {
998
+ ...b,
999
+ spendTodayNanos: b.dayStamp === today ? b.spendTodayNanos : 0,
1000
+ spendThisMonthNanos: b.monthStamp === month ? b.spendThisMonthNanos ?? 0 : 0,
1001
+ };
1002
+ },
1003
+ });
1004
+
1005
+ // Set a bucket's limits/controls. `user` and `action` are just dimensions here;
1006
+ // the client's ai.users / ai.actions namespaces are thin wrappers over this.
1007
+ export const setBucketLimits = mutation({
721
1008
  args: {
722
- userId: v.string(),
1009
+ dimension: v.string(),
1010
+ value: v.string(),
723
1011
  requestsPerMinute: v.optional(v.number()),
1012
+ maxConcurrent: v.optional(v.number()),
724
1013
  dailySpendLimitNanos: v.optional(v.number()),
1014
+ monthlySpendLimitNanos: v.optional(v.number()),
725
1015
  lifetimeSpendLimitNanos: v.optional(v.number()),
726
1016
  dailyTokenLimit: v.optional(v.number()),
1017
+ monthlyTokenLimit: v.optional(v.number()),
727
1018
  lifetimeTokenLimit: v.optional(v.number()),
1019
+ warnAtPct: v.optional(v.number()),
728
1020
  enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
729
1021
  blocked: v.optional(v.boolean()),
730
1022
  },
731
1023
  returns: v.null(),
732
1024
  handler: async (ctx, args) => {
733
- const user = await getOrCreateUser(ctx, args.userId);
734
- const { userId: _userId, ...limits } = args;
735
- await ctx.db.patch(user._id, limits);
1025
+ const bucket = await getOrCreateBucket(ctx, args.dimension, args.value);
1026
+ const { dimension: _d, value: _v, ...limits } = args;
1027
+ await ctx.db.patch(bucket._id, limits);
736
1028
  return null;
737
1029
  },
738
1030
  });
739
1031
 
740
- // Delete a user and all their request rows (e.g. account deletion / GDPR).
741
- // Deletes in bounded batches and self-reschedules so it never exceeds the
742
- // per-transaction document limit — a user with millions of rows still deletes.
743
- const DELETE_BATCH = 500;
744
- export const deleteUser = mutation({
745
- args: { userId: v.string() },
746
- returns: v.object({ deletedThisBatch: v.number(), done: v.boolean() }),
747
- handler: async (ctx, { userId }) => {
748
- const rows = await ctx.db
749
- .query("requests")
750
- .withIndex("userId", (q) => q.eq("userId", userId))
751
- .take(DELETE_BATCH);
752
- for (const r of rows) await ctx.db.delete(r._id);
753
- if (rows.length === DELETE_BATCH) {
754
- // More to go — continue in a fresh transaction.
755
- await ctx.scheduler.runAfter(0, api.lib.deleteUser, { userId });
756
- return { deletedThisBatch: rows.length, done: false };
757
- }
758
- // Last batch: remove the user row itself.
759
- const user = await ctx.db
760
- .query("users")
761
- .withIndex("userId", (q) => q.eq("userId", userId))
762
- .unique();
763
- if (user) await ctx.db.delete(user._id);
764
- return { deletedThisBatch: rows.length + (user ? 1 : 0), done: true };
765
- },
766
- });
767
-
768
- export const listActions = query({
769
- args: {},
770
- handler: async (ctx) => {
771
- const actions = await ctx.db.query("actions").take(ADMIN_LIST_CAP);
772
- const today = dayStamp();
773
- return actions.map((a) => ({
774
- ...a,
775
- spendTodayNanos: a.dayStamp === today ? a.spendTodayNanos : 0,
776
- }));
777
- },
778
- });
779
-
780
1032
  const vBumpArgs = {
781
1033
  dailyNanos: v.optional(v.number()),
1034
+ monthlyNanos: v.optional(v.number()),
782
1035
  lifetimeNanos: v.optional(v.number()),
783
1036
  };
784
1037
 
785
- // One-time "approve another $X" bumps, added on top of the standing cap without
786
- // changing it. Daily bumps apply to today only; lifetime bumps are permanent.
787
- export const bumpUser = mutation({
788
- args: { userId: v.string(), ...vBumpArgs },
1038
+ // One-time "approve another $X" bumps, added on top of a bucket's standing cap
1039
+ // without changing it. Daily/monthly bumps apply to the current window only;
1040
+ // lifetime bumps persist.
1041
+ export const bumpBucket = mutation({
1042
+ args: { dimension: v.string(), value: v.string(), ...vBumpArgs },
789
1043
  returns: v.null(),
790
- handler: async (ctx, { userId, dailyNanos, lifetimeNanos }) => {
791
- const user = await getOrCreateUser(ctx, userId);
1044
+ handler: async (ctx, { dimension, value, dailyNanos, monthlyNanos, lifetimeNanos }) => {
1045
+ const bucket = await getOrCreateBucket(ctx, dimension, value);
792
1046
  const today = dayStamp();
793
- const curDaily = user.bumpDayStamp === today ? user.dailyBumpNanos ?? 0 : 0;
794
- await ctx.db.patch(user._id, {
1047
+ const month = monthStamp();
1048
+ const curDaily =
1049
+ bucket.bumpDayStamp === today ? bucket.dailyBumpNanos ?? 0 : 0;
1050
+ const curMonthly =
1051
+ bucket.bumpMonthStamp === month ? bucket.monthlyBumpNanos ?? 0 : 0;
1052
+ await ctx.db.patch(bucket._id, {
795
1053
  bumpDayStamp: today,
796
1054
  dailyBumpNanos: curDaily + (dailyNanos ?? 0),
797
- lifetimeBumpNanos: (user.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
1055
+ bumpMonthStamp: month,
1056
+ monthlyBumpNanos: curMonthly + (monthlyNanos ?? 0),
1057
+ lifetimeBumpNanos: (bucket.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
798
1058
  });
799
1059
  return null;
800
1060
  },
801
1061
  });
802
1062
 
803
- export const bumpAction = mutation({
804
- args: { name: v.string(), ...vBumpArgs },
1063
+ // Manually credit or debit a bucket (comp a user, correct an overcharge).
1064
+ // Negative deltaNanos = credit/refund, positive = extra charge. Adjusts the
1065
+ // live day/month/lifetime windows AND the durable usage history, and records an
1066
+ // audit row. Does not touch the global sharded total or reservations.
1067
+ export const adjustBucket = mutation({
1068
+ args: {
1069
+ dimension: v.string(),
1070
+ value: v.string(),
1071
+ deltaNanos: v.number(),
1072
+ tokens: v.optional(v.number()),
1073
+ reason: v.optional(v.string()),
1074
+ },
805
1075
  returns: v.null(),
806
- handler: async (ctx, { name, dailyNanos, lifetimeNanos }) => {
807
- const action = await getOrCreateAction(ctx, name);
1076
+ handler: async (ctx, { dimension, value, deltaNanos, tokens, reason }) => {
1077
+ const b = await getOrCreateBucket(ctx, dimension, value);
808
1078
  const today = dayStamp();
809
- const curDaily = action.bumpDayStamp === today ? action.dailyBumpNanos ?? 0 : 0;
810
- await ctx.db.patch(action._id, {
811
- bumpDayStamp: today,
812
- dailyBumpNanos: curDaily + (dailyNanos ?? 0),
813
- lifetimeBumpNanos: (action.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
1079
+ const month = monthStamp();
1080
+ const dSame = b.dayStamp === today;
1081
+ const mSame = b.monthStamp === month;
1082
+ const dt = tokens ?? 0;
1083
+ await ctx.db.patch(b._id, {
1084
+ totalSpendNanos: Math.max(0, b.totalSpendNanos + deltaNanos),
1085
+ totalTokens: Math.max(0, b.totalTokens + dt),
1086
+ dayStamp: today,
1087
+ monthStamp: month,
1088
+ spendTodayNanos: Math.max(0, (dSame ? b.spendTodayNanos : 0) + deltaNanos),
1089
+ tokensToday: Math.max(0, (dSame ? b.tokensToday ?? 0 : 0) + dt),
1090
+ spendThisMonthNanos: Math.max(0, (mSame ? b.spendThisMonthNanos ?? 0 : 0) + deltaNanos),
1091
+ tokensThisMonth: Math.max(0, (mSame ? b.tokensThisMonth ?? 0 : 0) + dt),
814
1092
  });
1093
+ await ctx.db.insert("adjustments", { dimension, value, deltaNanos, tokens: dt, reason });
1094
+ await addUsage(ctx, dimension, value, "day", today, deltaNanos, dt, 0);
1095
+ await addUsage(ctx, dimension, value, "month", month, deltaNanos, dt, 0);
815
1096
  return null;
816
1097
  },
817
1098
  });
818
1099
 
819
- export const bumpGlobal = mutation({
820
- args: vBumpArgs,
821
- returns: v.null(),
822
- handler: async (ctx, { dailyNanos, lifetimeNanos }) => {
823
- const today = dayStamp();
824
- const s = await getSettings(ctx);
825
- const curDaily = s?.globalBumpDayStamp === today ? s?.globalDailyBumpNanos ?? 0 : 0;
826
- const patch = {
827
- globalBumpDayStamp: today,
828
- globalDailyBumpNanos: curDaily + (dailyNanos ?? 0),
829
- globalLifetimeBumpNanos: (s?.globalLifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
830
- };
831
- if (s) await ctx.db.patch(s._id, patch);
832
- else await ctx.db.insert("settings", { key: "singleton", ...patch });
833
- return null;
834
- },
1100
+ export const listAdjustments = query({
1101
+ args: { dimension: v.string(), value: v.string(), limit: v.optional(v.number()) },
1102
+ handler: async (ctx, { dimension, value, limit }) =>
1103
+ ctx.db
1104
+ .query("adjustments")
1105
+ .withIndex("dim_value", (q) => q.eq("dimension", dimension).eq("value", value))
1106
+ .order("desc")
1107
+ .take(limit ?? 50),
835
1108
  });
836
1109
 
837
- export const setActionLimits = mutation({
1110
+ // Durable spend history for a bucket: per-day or per-month rows, newest first.
1111
+ // Survives request retention.
1112
+ export const usageHistory = query({
838
1113
  args: {
839
- name: v.string(),
840
- dailySpendLimitNanos: v.optional(v.number()),
841
- lifetimeSpendLimitNanos: v.optional(v.number()),
842
- dailyTokenLimit: v.optional(v.number()),
843
- lifetimeTokenLimit: v.optional(v.number()),
844
- enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
845
- disabled: v.optional(v.boolean()),
1114
+ dimension: v.string(),
1115
+ value: v.string(),
1116
+ period: v.union(v.literal("day"), v.literal("month")),
1117
+ limit: v.optional(v.number()),
846
1118
  },
1119
+ handler: async (ctx, { dimension, value, period, limit }) =>
1120
+ ctx.db
1121
+ .query("usage")
1122
+ .withIndex("bucket_period_stamp", (q) =>
1123
+ q.eq("dimension", dimension).eq("value", value).eq("period", period)
1124
+ )
1125
+ .order("desc")
1126
+ .take(limit ?? 90),
1127
+ });
1128
+
1129
+ // Deployment-wide default threshold for approaching-limit alerts (fraction of a
1130
+ // cap, e.g. 0.8). Buckets can override with their own warnAtPct.
1131
+ export const setAlertDefaults = mutation({
1132
+ args: { warnAtPct: v.optional(v.number()) },
847
1133
  returns: v.null(),
848
- handler: async (ctx, args) => {
849
- const action = await getOrCreateAction(ctx, args.name);
850
- const { name: _name, ...limits } = args;
851
- await ctx.db.patch(action._id, limits);
1134
+ handler: async (ctx, { warnAtPct }) => {
1135
+ const existing = await getSettings(ctx);
1136
+ if (existing) await ctx.db.patch(existing._id, { defaultWarnAtPct: warnAtPct });
1137
+ else await ctx.db.insert("settings", { key: "singleton", defaultWarnAtPct: warnAtPct });
852
1138
  return null;
853
1139
  },
854
1140
  });
855
1141
 
1142
+ // Delete a bucket and (for the `user` dimension) all of that user's request
1143
+ // rows — e.g. account deletion / GDPR. Deletes requests in bounded batches and
1144
+ // self-reschedules so it never exceeds the per-transaction document limit.
1145
+ const DELETE_BATCH = 500;
1146
+ export const deleteBucket = mutation({
1147
+ args: { dimension: v.string(), value: v.string() },
1148
+ returns: v.object({ deletedThisBatch: v.number(), done: v.boolean() }),
1149
+ handler: async (ctx, { dimension, value }) => {
1150
+ // Only the user dimension owns request rows (indexed by userId). Other
1151
+ // dimensions just drop their budget-holder row.
1152
+ if (dimension === USER_DIM) {
1153
+ const rows = await ctx.db
1154
+ .query("requests")
1155
+ .withIndex("userId", (q) => q.eq("userId", value))
1156
+ .take(DELETE_BATCH);
1157
+ for (const r of rows) {
1158
+ await deleteRequestTags(ctx, r._id);
1159
+ await ctx.db.delete(r._id);
1160
+ }
1161
+ if (rows.length === DELETE_BATCH) {
1162
+ await ctx.scheduler.runAfter(0, api.lib.deleteBucket, {
1163
+ dimension,
1164
+ value,
1165
+ });
1166
+ return { deletedThisBatch: rows.length, done: false };
1167
+ }
1168
+ const bucket = await getBucketDoc(ctx, dimension, value);
1169
+ if (bucket) await ctx.db.delete(bucket._id);
1170
+ return { deletedThisBatch: rows.length + (bucket ? 1 : 0), done: true };
1171
+ }
1172
+ const bucket = await getBucketDoc(ctx, dimension, value);
1173
+ if (bucket) await ctx.db.delete(bucket._id);
1174
+ return { deletedThisBatch: bucket ? 1 : 0, done: true };
1175
+ },
1176
+ });
1177
+
856
1178
  export const getModelPolicy = query({
857
1179
  args: {},
858
1180
  returns: v.object({
@@ -880,6 +1202,9 @@ export const getGlobalStatus = query({
880
1202
  enforcement: v.union(v.literal("hard"), v.literal("soft")),
881
1203
  spentTodayNanos: v.number(),
882
1204
  spentTotalNanos: v.number(),
1205
+ // deployment-wide config (surfaced for the admin dashboard)
1206
+ retentionMs: v.union(v.number(), v.null()),
1207
+ defaultWarnAtPct: v.union(v.number(), v.null()),
883
1208
  }),
884
1209
  handler: async (ctx) => {
885
1210
  const s = await ctx.db
@@ -892,6 +1217,8 @@ export const getGlobalStatus = query({
892
1217
  enforcement: s?.globalEnforcement ?? "hard",
893
1218
  spentTodayNanos: await globalSpend.count(ctx, globalDayKey(dayStamp())),
894
1219
  spentTotalNanos: await globalSpend.count(ctx, GLOBAL_TOTAL),
1220
+ retentionMs: s?.retentionMs ?? null,
1221
+ defaultWarnAtPct: s?.defaultWarnAtPct ?? null,
895
1222
  };
896
1223
  },
897
1224
  });
@@ -920,6 +1247,24 @@ export const setGlobalLimits = mutation({
920
1247
  },
921
1248
  });
922
1249
 
1250
+ export const bumpGlobal = mutation({
1251
+ args: vBumpArgs,
1252
+ returns: v.null(),
1253
+ handler: async (ctx, { dailyNanos, lifetimeNanos }) => {
1254
+ const today = dayStamp();
1255
+ const s = await getSettings(ctx);
1256
+ const curDaily = s?.globalBumpDayStamp === today ? s?.globalDailyBumpNanos ?? 0 : 0;
1257
+ const patch = {
1258
+ globalBumpDayStamp: today,
1259
+ globalDailyBumpNanos: curDaily + (dailyNanos ?? 0),
1260
+ globalLifetimeBumpNanos: (s?.globalLifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
1261
+ };
1262
+ if (s) await ctx.db.patch(s._id, patch);
1263
+ else await ctx.db.insert("settings", { key: "singleton", ...patch });
1264
+ return null;
1265
+ },
1266
+ });
1267
+
923
1268
  export const setModelPolicy = mutation({
924
1269
  args: {
925
1270
  mode: v.union(
@@ -953,12 +1298,18 @@ export const setPrice = mutation({
953
1298
  model: v.string(),
954
1299
  inputNanosPerMTok: v.number(),
955
1300
  outputNanosPerMTok: v.number(),
1301
+ // optional cache-read rate; if omitted, a default discount off input applies
1302
+ cachedNanosPerMTok: v.optional(v.number()),
956
1303
  },
957
1304
  returns: v.null(),
958
1305
  handler: async (ctx, args) => {
959
1306
  // Negative prices would make costOf return a negative cost, which folds
960
1307
  // into totals as a spend *refund* — pushing a user back under their cap.
961
- if (args.inputNanosPerMTok < 0 || args.outputNanosPerMTok < 0) {
1308
+ if (
1309
+ args.inputNanosPerMTok < 0 ||
1310
+ args.outputNanosPerMTok < 0 ||
1311
+ (args.cachedNanosPerMTok ?? 0) < 0
1312
+ ) {
962
1313
  throw new Error("Prices must be non-negative");
963
1314
  }
964
1315
  const existing = await ctx.db
@@ -978,7 +1329,10 @@ export const listPrices = query({
978
1329
  args: {},
979
1330
  handler: async (ctx) => {
980
1331
  const overrides = await ctx.db.query("prices").collect();
981
- const merged: Record<string, { input: number; output: number; overridden: boolean }> = {};
1332
+ const merged: Record<
1333
+ string,
1334
+ { input: number; output: number; cached?: number; overridden: boolean }
1335
+ > = {};
982
1336
  for (const [model, p] of Object.entries(DEFAULT_PRICES)) {
983
1337
  merged[model] = { ...p, overridden: false };
984
1338
  }
@@ -986,6 +1340,7 @@ export const listPrices = query({
986
1340
  merged[o.model] = {
987
1341
  input: o.inputNanosPerMTok,
988
1342
  output: o.outputNanosPerMTok,
1343
+ cached: o.cachedNanosPerMTok,
989
1344
  overridden: true,
990
1345
  };
991
1346
  }