@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.
- package/README.md +335 -183
- package/dist/client/dashboard.d.ts +1 -0
- package/dist/client/dashboard.js +223 -0
- package/dist/client/index.d.ts +494 -35
- package/dist/client/index.js +276 -33
- package/dist/component/_generated/component.d.ts +56 -24
- package/dist/component/lib.d.ts +132 -40
- package/dist/component/lib.js +590 -305
- package/dist/component/schema.d.ts +103 -53
- package/dist/component/schema.js +81 -38
- package/package.json +5 -5
- package/src/client/dashboard.ts +223 -0
- package/src/client/index.ts +407 -73
- package/src/component/_generated/component.ts +79 -29
- package/src/component/lib.test.ts +208 -15
- package/src/component/lib.ts +688 -333
- package/src/component/schema.ts +84 -38
package/src/component/lib.ts
CHANGED
|
@@ -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.
|
|
40
|
-
//
|
|
41
|
-
//
|
|
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
|
-
//
|
|
115
|
-
//
|
|
116
|
-
//
|
|
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
|
-
}): {
|
|
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
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
if (
|
|
156
|
-
|
|
157
|
-
|
|
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
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
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
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
|
|
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("
|
|
206
|
-
|
|
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
|
|
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
|
-
|
|
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
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
const
|
|
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
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
|
|
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
|
-
|
|
304
|
-
|
|
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
|
-
|
|
310
|
-
//
|
|
311
|
-
//
|
|
312
|
-
//
|
|
313
|
-
|
|
314
|
-
const
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
320
|
-
|
|
321
|
-
|
|
322
|
-
|
|
323
|
-
|
|
324
|
-
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
|
|
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
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
352
|
-
|
|
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:
|
|
356
|
-
reservedSpendToday:
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
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
|
-
|
|
365
|
-
|
|
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
|
-
|
|
369
|
-
|
|
569
|
+
b.lifetimeSpendLimitNanos,
|
|
570
|
+
b.lifetimeBumpNanos
|
|
370
571
|
),
|
|
371
|
-
dailyTokenLimit:
|
|
372
|
-
|
|
572
|
+
dailyTokenLimit: b.dailyTokenLimit,
|
|
573
|
+
monthlyTokenLimit: b.monthlyTokenLimit,
|
|
574
|
+
lifetimeTokenLimit: b.lifetimeTokenLimit,
|
|
373
575
|
});
|
|
374
|
-
if (
|
|
375
|
-
warnings.push(...
|
|
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
|
|
379
|
-
//
|
|
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-
|
|
383
|
-
//
|
|
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
|
|
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
|
-
|
|
390
|
-
(
|
|
391
|
-
|
|
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:
|
|
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
|
-
|
|
409
|
-
|
|
615
|
+
settings.globalDailySpendLimitNanos,
|
|
616
|
+
settings.globalBumpDayStamp === today ? settings.globalDailyBumpNanos : 0
|
|
410
617
|
),
|
|
411
618
|
lifetimeSpendLimitNanos: withBump(
|
|
412
|
-
|
|
413
|
-
|
|
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
|
|
421
|
-
//
|
|
422
|
-
// (e.g. all callers of one action); with no cap
|
|
423
|
-
// to consult. Totals are still accrued later,
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
|
|
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
|
-
|
|
441
|
-
|
|
442
|
-
|
|
443
|
-
|
|
444
|
-
|
|
445
|
-
|
|
446
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
493
|
-
|
|
494
|
-
|
|
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)
|
|
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
|
|
522
|
-
// reservation. Idempotent: guarded by `settled` so the scheduler
|
|
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
|
-
|
|
533
|
-
|
|
534
|
-
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
|
|
538
|
-
|
|
539
|
-
|
|
540
|
-
|
|
541
|
-
|
|
542
|
-
|
|
543
|
-
|
|
544
|
-
|
|
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
|
-
|
|
557
|
-
|
|
558
|
-
|
|
559
|
-
|
|
560
|
-
|
|
561
|
-
|
|
562
|
-
|
|
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
|
|
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: {
|
|
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
|
-
|
|
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
|
-
|
|
707
|
-
|
|
708
|
-
|
|
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("
|
|
711
|
-
const
|
|
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
|
-
|
|
714
|
-
|
|
715
|
-
|
|
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
|
|
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
|
-
|
|
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
|
|
734
|
-
const {
|
|
735
|
-
await ctx.db.patch(
|
|
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
|
|
786
|
-
// changing it. Daily bumps apply to
|
|
787
|
-
|
|
788
|
-
|
|
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, {
|
|
791
|
-
const
|
|
1044
|
+
handler: async (ctx, { dimension, value, dailyNanos, monthlyNanos, lifetimeNanos }) => {
|
|
1045
|
+
const bucket = await getOrCreateBucket(ctx, dimension, value);
|
|
792
1046
|
const today = dayStamp();
|
|
793
|
-
const
|
|
794
|
-
|
|
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
|
-
|
|
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
|
-
|
|
804
|
-
|
|
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, {
|
|
807
|
-
const
|
|
1076
|
+
handler: async (ctx, { dimension, value, deltaNanos, tokens, reason }) => {
|
|
1077
|
+
const b = await getOrCreateBucket(ctx, dimension, value);
|
|
808
1078
|
const today = dayStamp();
|
|
809
|
-
const
|
|
810
|
-
|
|
811
|
-
|
|
812
|
-
|
|
813
|
-
|
|
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
|
|
820
|
-
args:
|
|
821
|
-
|
|
822
|
-
|
|
823
|
-
|
|
824
|
-
|
|
825
|
-
|
|
826
|
-
|
|
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
|
-
|
|
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
|
-
|
|
840
|
-
|
|
841
|
-
|
|
842
|
-
|
|
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,
|
|
849
|
-
const
|
|
850
|
-
|
|
851
|
-
await ctx.db.
|
|
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 (
|
|
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<
|
|
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
|
}
|