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