@convex-dev/ai-budget 0.0.2-alpha.3 → 0.0.2-alpha.5
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 +116 -17
- package/dist/client/index.d.ts +433 -34
- package/dist/client/index.js +141 -33
- package/dist/component/_generated/component.d.ts +54 -24
- package/dist/component/lib.d.ts +130 -40
- package/dist/component/lib.js +550 -290
- package/dist/component/schema.d.ts +103 -53
- package/dist/component/schema.js +81 -38
- package/package.json +1 -1
- package/src/client/index.ts +241 -71
- package/src/component/_generated/component.ts +77 -29
- package/src/component/lib.test.ts +167 -15
- package/src/component/lib.ts +640 -313
- package/src/component/schema.ts +84 -38
package/dist/component/lib.js
CHANGED
|
@@ -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);
|
|
@@ -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
|
-
//
|
|
87
|
-
//
|
|
88
|
-
//
|
|
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
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
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
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
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
|
|
135
|
-
const existing = await ctx
|
|
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("
|
|
142
|
-
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
197
|
-
return reject("blocked", `User "${args.userId}" is blocked`, false);
|
|
198
|
-
}
|
|
314
|
+
const notices = [];
|
|
199
315
|
// Model allow/deny policy (component-wide).
|
|
200
|
-
const
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
const
|
|
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,176 @@ export const startRequest = mutation({
|
|
|
208
325
|
return reject("model_denied", `Model "${args.model}" is denied`);
|
|
209
326
|
}
|
|
210
327
|
}
|
|
211
|
-
|
|
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);
|
|
342
|
+
}
|
|
343
|
+
}
|
|
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
|
+
// Rate limit is enforced on the `user` dimension (the requests table is
|
|
353
|
+
// indexed by userId, so the 60s window is a bounded index read).
|
|
354
|
+
const userBucket = buckets.find((b) => b.dimension === USER_DIM);
|
|
355
|
+
if (userBucket.requestsPerMinute !== undefined) {
|
|
212
356
|
// Bounded read: we only need to know whether non-blocked requests in the
|
|
213
|
-
// window have reached the limit. Reading limit+
|
|
357
|
+
// window have reached the limit. Reading limit+50 non-blocked rows is
|
|
214
358
|
// enough, and .take caps the scan even if the window is flooded. Blocked
|
|
215
359
|
// rows aren't persisted for rate-limit/blocked rejections (see reject),
|
|
216
360
|
// so the window stays small.
|
|
217
361
|
const recent = await ctx.db
|
|
218
362
|
.query("requests")
|
|
219
363
|
.withIndex("userId", (q) => q.eq("userId", args.userId).gt("_creationTime", Date.now() - 60_000))
|
|
220
|
-
.take(
|
|
364
|
+
.take(userBucket.requestsPerMinute + 50);
|
|
221
365
|
if (recent.filter((r) => r.status !== "blocked").length >=
|
|
222
|
-
|
|
223
|
-
return reject("rate_limit", `Rate limit exceeded for "${args.userId}" (${
|
|
366
|
+
userBucket.requestsPerMinute) {
|
|
367
|
+
return reject("rate_limit", `Rate limit exceeded for "${args.userId}" (${userBucket.requestsPerMinute}/min)`, false);
|
|
224
368
|
}
|
|
225
369
|
}
|
|
226
370
|
// Committed + already-reserved in-flight usage + this estimate must fit
|
|
227
|
-
// under
|
|
228
|
-
// serializable isolation makes it a true atomic check-and-reserve
|
|
229
|
-
// enforcement turns a violation into a
|
|
230
|
-
|
|
231
|
-
const
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
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`);
|
|
258
|
-
}
|
|
259
|
-
const aSameDay = action.dayStamp === today;
|
|
260
|
-
const actionEval = evaluateCaps({
|
|
261
|
-
label: "action",
|
|
262
|
-
name: action.name,
|
|
263
|
-
enforcement: action.enforcement ?? "hard",
|
|
371
|
+
// under EACH bucket's cap. Decided and reserved in one transaction, so
|
|
372
|
+
// Convex's serializable isolation makes it a true atomic check-and-reserve
|
|
373
|
+
// across every capped bucket. Soft enforcement turns a violation into a
|
|
374
|
+
// warning instead of a block.
|
|
375
|
+
for (const b of buckets) {
|
|
376
|
+
const sameDay = b.dayStamp === today;
|
|
377
|
+
const sameMonth = b.monthStamp === month;
|
|
378
|
+
const ev = evaluateCaps({
|
|
379
|
+
label: b.dimension,
|
|
380
|
+
name: b.value,
|
|
381
|
+
enforcement: b.enforcement ?? "hard",
|
|
382
|
+
warnAtPct: b.warnAtPct ?? defaultWarnAtPct,
|
|
264
383
|
estCost: est.cost,
|
|
265
384
|
estTokens: est.tokens,
|
|
266
|
-
spendToday:
|
|
267
|
-
reservedSpendToday:
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
385
|
+
spendToday: sameDay ? b.spendTodayNanos : 0,
|
|
386
|
+
reservedSpendToday: sameDay ? b.reservedTodayNanos ?? 0 : 0,
|
|
387
|
+
spendThisMonth: sameMonth ? b.spendThisMonthNanos ?? 0 : 0,
|
|
388
|
+
reservedSpendMonth: sameMonth ? b.reservedMonthNanos ?? 0 : 0,
|
|
389
|
+
totalSpend: b.totalSpendNanos,
|
|
390
|
+
reservedSpendTotal: b.reservedTotalNanos ?? 0,
|
|
391
|
+
tokensToday: sameDay ? b.tokensToday ?? 0 : 0,
|
|
392
|
+
reservedTokensToday: sameDay ? b.reservedTodayTokens ?? 0 : 0,
|
|
393
|
+
tokensThisMonth: sameMonth ? b.tokensThisMonth ?? 0 : 0,
|
|
394
|
+
reservedTokensMonth: sameMonth ? b.reservedMonthTokens ?? 0 : 0,
|
|
395
|
+
totalTokens: b.totalTokens,
|
|
396
|
+
reservedTokensTotal: b.reservedTotalTokens ?? 0,
|
|
397
|
+
dailySpendLimitNanos: withBump(b.dailySpendLimitNanos, b.bumpDayStamp === today ? b.dailyBumpNanos : 0),
|
|
398
|
+
monthlySpendLimitNanos: withBump(b.monthlySpendLimitNanos, b.bumpMonthStamp === month ? b.monthlyBumpNanos : 0),
|
|
399
|
+
lifetimeSpendLimitNanos: withBump(b.lifetimeSpendLimitNanos, b.lifetimeBumpNanos),
|
|
400
|
+
dailyTokenLimit: b.dailyTokenLimit,
|
|
401
|
+
monthlyTokenLimit: b.monthlyTokenLimit,
|
|
402
|
+
lifetimeTokenLimit: b.lifetimeTokenLimit,
|
|
278
403
|
});
|
|
279
|
-
if (
|
|
280
|
-
return reject(
|
|
281
|
-
warnings.push(...
|
|
404
|
+
if (ev.hard)
|
|
405
|
+
return reject(ev.hard.code, ev.hard.reason);
|
|
406
|
+
warnings.push(...ev.warnings);
|
|
407
|
+
notices.push(...ev.notices);
|
|
282
408
|
}
|
|
283
409
|
// Deployment-wide ("global") spend cap — same reserve-then-settle model and
|
|
284
|
-
// the SAME evaluateCaps logic as the per-
|
|
410
|
+
// the SAME evaluateCaps logic as the per-bucket caps above, so the guarantee
|
|
285
411
|
// statement is uniform: a request is admitted only if
|
|
286
412
|
// committed + reserved + estimate <= cap. The ONE difference is the holder:
|
|
287
|
-
// per-
|
|
413
|
+
// per-bucket caps reserve on a single row (an exact atomic
|
|
288
414
|
// check-and-reserve), while the global holder is a sharded counter for
|
|
289
415
|
// throughput — its committed total is read as an eventually-consistent sum
|
|
290
416
|
// with no cross-request reservation, so a hard global cap can overshoot by a
|
|
291
417
|
// bounded amount under burst. That's the deliberate exactness/throughput
|
|
292
418
|
// trade for a deployment-wide killswitch; it's the only approximate scope.
|
|
293
|
-
if (
|
|
294
|
-
(
|
|
295
|
-
|
|
419
|
+
if (settings &&
|
|
420
|
+
(settings.globalDailySpendLimitNanos !== undefined ||
|
|
421
|
+
settings.globalLifetimeSpendLimitNanos !== undefined)) {
|
|
296
422
|
const globalEval = evaluateCaps({
|
|
297
423
|
label: "global",
|
|
298
424
|
name: "deployment",
|
|
299
|
-
enforcement:
|
|
425
|
+
enforcement: settings.globalEnforcement ?? "hard",
|
|
426
|
+
warnAtPct: defaultWarnAtPct,
|
|
300
427
|
estCost: est.cost,
|
|
301
428
|
estTokens: est.tokens,
|
|
302
429
|
spendToday: await globalSpend.count(ctx, globalDayKey(today)),
|
|
303
430
|
reservedSpendToday: 0, // sharded holder: no cross-request reservation
|
|
431
|
+
spendThisMonth: 0, // global tracks daily + lifetime only
|
|
432
|
+
reservedSpendMonth: 0,
|
|
304
433
|
totalSpend: await globalSpend.count(ctx, GLOBAL_TOTAL),
|
|
305
434
|
reservedSpendTotal: 0,
|
|
306
435
|
tokensToday: 0,
|
|
307
436
|
reservedTokensToday: 0,
|
|
437
|
+
tokensThisMonth: 0,
|
|
438
|
+
reservedTokensMonth: 0,
|
|
308
439
|
totalTokens: 0,
|
|
309
440
|
reservedTokensTotal: 0,
|
|
310
|
-
dailySpendLimitNanos: withBump(
|
|
311
|
-
lifetimeSpendLimitNanos: withBump(
|
|
441
|
+
dailySpendLimitNanos: withBump(settings.globalDailySpendLimitNanos, settings.globalBumpDayStamp === today ? settings.globalDailyBumpNanos : 0),
|
|
442
|
+
lifetimeSpendLimitNanos: withBump(settings.globalLifetimeSpendLimitNanos, settings.globalLifetimeBumpNanos),
|
|
312
443
|
});
|
|
313
444
|
if (globalEval.hard)
|
|
314
445
|
return reject(globalEval.hard.code, globalEval.hard.reason);
|
|
315
446
|
warnings.push(...globalEval.warnings);
|
|
447
|
+
notices.push(...globalEval.notices);
|
|
316
448
|
}
|
|
317
|
-
// Passed — reserve, but ONLY on
|
|
318
|
-
//
|
|
319
|
-
// (e.g. all callers of one action); with no cap
|
|
320
|
-
// to consult. Totals are still accrued later,
|
|
321
|
-
|
|
322
|
-
|
|
323
|
-
|
|
324
|
-
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
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, {
|
|
449
|
+
// Passed — reserve, but ONLY on buckets that actually have a cap. Writing an
|
|
450
|
+
// uncapped bucket's row here would serialize every request that shares it
|
|
451
|
+
// (e.g. all callers of one action, or every request in one env); with no cap
|
|
452
|
+
// there's no reserved amount to consult. Totals are still accrued later,
|
|
453
|
+
// asynchronously, in foldOne — for every bucket, capped or not.
|
|
454
|
+
for (const b of buckets) {
|
|
455
|
+
if (!needsReserve(b))
|
|
456
|
+
continue;
|
|
457
|
+
const sameDay = b.dayStamp === today;
|
|
458
|
+
const sameMonth = b.monthStamp === month;
|
|
459
|
+
await ctx.db.patch(b._id, {
|
|
336
460
|
dayStamp: today,
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
461
|
+
monthStamp: month,
|
|
462
|
+
spendTodayNanos: sameDay ? b.spendTodayNanos : 0,
|
|
463
|
+
tokensToday: sameDay ? b.tokensToday ?? 0 : 0,
|
|
464
|
+
spendThisMonthNanos: sameMonth ? b.spendThisMonthNanos ?? 0 : 0,
|
|
465
|
+
tokensThisMonth: sameMonth ? b.tokensThisMonth ?? 0 : 0,
|
|
466
|
+
reservedTodayNanos: (sameDay ? b.reservedTodayNanos ?? 0 : 0) + est.cost,
|
|
467
|
+
reservedMonthNanos: (sameMonth ? b.reservedMonthNanos ?? 0 : 0) + est.cost,
|
|
468
|
+
reservedTotalNanos: (b.reservedTotalNanos ?? 0) + est.cost,
|
|
469
|
+
reservedTodayTokens: (sameDay ? b.reservedTodayTokens ?? 0 : 0) + est.tokens,
|
|
470
|
+
reservedMonthTokens: (sameMonth ? b.reservedMonthTokens ?? 0 : 0) + est.tokens,
|
|
471
|
+
reservedTotalTokens: (b.reservedTotalTokens ?? 0) + est.tokens,
|
|
472
|
+
pendingCount: (b.pendingCount ?? 0) + 1,
|
|
344
473
|
});
|
|
345
474
|
}
|
|
346
475
|
const requestId = await ctx.db.insert("requests", {
|
|
347
|
-
|
|
476
|
+
userId: args.userId,
|
|
477
|
+
actionName: args.actionName,
|
|
478
|
+
...(extraTags.length ? { tags: extraTags } : {}),
|
|
479
|
+
model: args.model,
|
|
480
|
+
messages: args.messages,
|
|
481
|
+
rerunOf: args.rerunOf,
|
|
348
482
|
status: "pending",
|
|
349
483
|
estimatedNanos: est.cost,
|
|
350
484
|
estimatedTokens: est.tokens,
|
|
351
485
|
...(priceInfo.known ? {} : { unpricedModel: true }),
|
|
352
486
|
...(warnings.length > 0 ? { overBudget: true } : {}),
|
|
353
487
|
});
|
|
354
|
-
|
|
488
|
+
// Reverse index so the request log can be filtered by any extra tag
|
|
489
|
+
// dimension (user/action are already indexed on the requests table).
|
|
490
|
+
for (const t of extraTags) {
|
|
491
|
+
await ctx.db.insert("requestTags", {
|
|
492
|
+
dimension: t.dimension,
|
|
493
|
+
value: t.value,
|
|
494
|
+
requestId,
|
|
495
|
+
});
|
|
496
|
+
}
|
|
497
|
+
return { allowed: true, requestId, warnings, notices };
|
|
355
498
|
},
|
|
356
499
|
});
|
|
357
500
|
export const finishRequest = mutation({
|
|
@@ -362,6 +505,10 @@ export const finishRequest = mutation({
|
|
|
362
505
|
promptTokens: v.optional(v.number()),
|
|
363
506
|
completionTokens: v.optional(v.number()),
|
|
364
507
|
cachedTokens: v.optional(v.number()),
|
|
508
|
+
// Authoritative cost from the gateway, if it ever reports one. When present
|
|
509
|
+
// it's recorded verbatim (no token-based estimate); when absent we price
|
|
510
|
+
// from tokens (cache-aware). Wired now so adopting a real cost is one line.
|
|
511
|
+
costNanos: v.optional(v.number()),
|
|
365
512
|
latencyMs: v.optional(v.number()),
|
|
366
513
|
},
|
|
367
514
|
returns: v.object({ costNanos: v.number() }),
|
|
@@ -384,7 +531,11 @@ export const finishRequest = mutation({
|
|
|
384
531
|
const promptTokens = Math.max(0, args.promptTokens ?? 0);
|
|
385
532
|
const completionTokens = Math.max(0, args.completionTokens ?? 0);
|
|
386
533
|
const cachedTokens = Math.min(promptTokens, Math.max(0, args.cachedTokens ?? 0));
|
|
387
|
-
|
|
534
|
+
// Prefer an authoritative gateway cost when supplied; otherwise price from
|
|
535
|
+
// tokens, discounting the cached (prompt-cache-read) slice of the prompt.
|
|
536
|
+
const costNanos = args.costNanos !== undefined && args.costNanos >= 0
|
|
537
|
+
? Math.round(args.costNanos)
|
|
538
|
+
: settleCost(promptTokens, cachedTokens, completionTokens, await getPrice(ctx, request.model));
|
|
388
539
|
// Durable write to the request's OWN row only — uncontended, so it always
|
|
389
540
|
// lands. `settled: false` hands it to the fold step; the row is never left
|
|
390
541
|
// orphaned in "pending" even if the totals update below fails and retries.
|
|
@@ -399,7 +550,7 @@ export const finishRequest = mutation({
|
|
|
399
550
|
latencyMs: args.latencyMs,
|
|
400
551
|
settled: false,
|
|
401
552
|
});
|
|
402
|
-
// Fold into the (hot)
|
|
553
|
+
// Fold into the (hot) per-bucket counters in a separate mutation. If it
|
|
403
554
|
// exhausts retries under contention, the cron reconciler picks it up.
|
|
404
555
|
await ctx.scheduler.runAfter(0, internal.lib.foldTotals, {
|
|
405
556
|
requestId: args.requestId,
|
|
@@ -407,9 +558,9 @@ export const finishRequest = mutation({
|
|
|
407
558
|
return { costNanos };
|
|
408
559
|
},
|
|
409
560
|
});
|
|
410
|
-
// Fold one finished request into
|
|
411
|
-
// reservation. Idempotent: guarded by `settled` so the scheduler
|
|
412
|
-
// reconciler can never double-count.
|
|
561
|
+
// Fold one finished request into every attributed bucket's running totals,
|
|
562
|
+
// releasing its reservation. Idempotent: guarded by `settled` so the scheduler
|
|
563
|
+
// and the cron reconciler can never double-count.
|
|
413
564
|
async function foldOne(ctx, req) {
|
|
414
565
|
if (!req || req.settled !== false)
|
|
415
566
|
return;
|
|
@@ -418,46 +569,42 @@ async function foldOne(ctx, req) {
|
|
|
418
569
|
const tokens = (req.promptTokens ?? 0) + (req.completionTokens ?? 0);
|
|
419
570
|
const estTokens = req.estimatedTokens ?? 0;
|
|
420
571
|
const today = dayStamp();
|
|
421
|
-
const
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
433
|
-
|
|
434
|
-
|
|
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,
|
|
572
|
+
const month = monthStamp();
|
|
573
|
+
// Accrue into every attributed bucket (user, action, and each tag) — capped
|
|
574
|
+
// or not. Buckets that never held a reservation have their reserved fields
|
|
575
|
+
// clamped at 0 by Math.max, so subtracting an estimate is a harmless no-op.
|
|
576
|
+
// Also write the durable per-(bucket, day/month) usage rows that survive
|
|
577
|
+
// request retention, so spend history outlives the raw request log.
|
|
578
|
+
for (const t of requestBuckets(req.userId, req.actionName, req.tags)) {
|
|
579
|
+
const b = await getOrCreateBucket(ctx, t.dimension, t.value);
|
|
580
|
+
const sameDay = b.dayStamp === today;
|
|
581
|
+
const sameMonth = b.monthStamp === month;
|
|
582
|
+
await ctx.db.patch(b._id, {
|
|
583
|
+
totalSpendNanos: b.totalSpendNanos + actual,
|
|
584
|
+
totalRequests: b.totalRequests + 1,
|
|
585
|
+
totalTokens: b.totalTokens + tokens,
|
|
443
586
|
dayStamp: today,
|
|
444
|
-
|
|
445
|
-
|
|
446
|
-
|
|
447
|
-
|
|
448
|
-
|
|
449
|
-
|
|
450
|
-
|
|
587
|
+
monthStamp: month,
|
|
588
|
+
spendTodayNanos: (sameDay ? b.spendTodayNanos : 0) + actual,
|
|
589
|
+
tokensToday: (sameDay ? b.tokensToday ?? 0 : 0) + tokens,
|
|
590
|
+
spendThisMonthNanos: (sameMonth ? b.spendThisMonthNanos ?? 0 : 0) + actual,
|
|
591
|
+
tokensThisMonth: (sameMonth ? b.tokensThisMonth ?? 0 : 0) + tokens,
|
|
592
|
+
reservedTodayNanos: Math.max(0, (sameDay ? b.reservedTodayNanos ?? 0 : 0) - estCost),
|
|
593
|
+
reservedMonthNanos: Math.max(0, (sameMonth ? b.reservedMonthNanos ?? 0 : 0) - estCost),
|
|
594
|
+
reservedTotalNanos: Math.max(0, (b.reservedTotalNanos ?? 0) - estCost),
|
|
595
|
+
reservedTodayTokens: Math.max(0, (sameDay ? b.reservedTodayTokens ?? 0 : 0) - estTokens),
|
|
596
|
+
reservedMonthTokens: Math.max(0, (sameMonth ? b.reservedMonthTokens ?? 0 : 0) - estTokens),
|
|
597
|
+
reservedTotalTokens: Math.max(0, (b.reservedTotalTokens ?? 0) - estTokens),
|
|
598
|
+
pendingCount: Math.max(0, (b.pendingCount ?? 0) - 1),
|
|
451
599
|
});
|
|
600
|
+
await addUsage(ctx, t.dimension, t.value, "day", today, actual, tokens, 1);
|
|
601
|
+
await addUsage(ctx, t.dimension, t.value, "month", month, actual, tokens, 1);
|
|
452
602
|
}
|
|
453
603
|
// Deployment-wide totals via the sharded counter (only when a global cap is
|
|
454
604
|
// configured — otherwise skip the writes entirely). Distributed across shards,
|
|
455
605
|
// so this does not serialize on a single row.
|
|
456
606
|
if (actual > 0) {
|
|
457
|
-
const settings = await ctx
|
|
458
|
-
.query("settings")
|
|
459
|
-
.withIndex("key", (q) => q.eq("key", "singleton"))
|
|
460
|
-
.unique();
|
|
607
|
+
const settings = await getSettings(ctx);
|
|
461
608
|
if (settings &&
|
|
462
609
|
(settings.globalDailySpendLimitNanos !== undefined ||
|
|
463
610
|
settings.globalLifetimeSpendLimitNanos !== undefined)) {
|
|
@@ -520,6 +667,7 @@ export const reconcile = internalMutation({
|
|
|
520
667
|
// Only rows that are done and accounted: folded (settled === true) or a
|
|
521
668
|
// blocked attempt (never needs folding). Never a pending/unfolded row.
|
|
522
669
|
if (req.settled === true || req.status === "blocked") {
|
|
670
|
+
await deleteRequestTags(ctx, req._id);
|
|
523
671
|
await ctx.db.delete(req._id);
|
|
524
672
|
purged++;
|
|
525
673
|
}
|
|
@@ -565,10 +713,19 @@ export const getRequest = query({
|
|
|
565
713
|
handler: async (ctx, args) => ctx.db.get(args.requestId),
|
|
566
714
|
});
|
|
567
715
|
export const listRequests = query({
|
|
568
|
-
args: {
|
|
716
|
+
args: {
|
|
717
|
+
userId: v.optional(v.string()),
|
|
718
|
+
// Filter by any attribution dimension (user/action indexed on the request;
|
|
719
|
+
// custom tag dimensions resolved via the requestTags reverse index).
|
|
720
|
+
dimension: v.optional(v.string()),
|
|
721
|
+
value: v.optional(v.string()),
|
|
722
|
+
limit: v.optional(v.number()),
|
|
723
|
+
},
|
|
569
724
|
handler: async (ctx, args) => {
|
|
570
725
|
const limit = args.limit ?? 50;
|
|
571
|
-
|
|
726
|
+
const dim = args.dimension;
|
|
727
|
+
const val = args.value ?? args.userId;
|
|
728
|
+
if (dim === undefined && args.userId !== undefined) {
|
|
572
729
|
const userId = args.userId;
|
|
573
730
|
return await ctx.db
|
|
574
731
|
.query("requests")
|
|
@@ -576,155 +733,234 @@ export const listRequests = query({
|
|
|
576
733
|
.order("desc")
|
|
577
734
|
.take(limit);
|
|
578
735
|
}
|
|
736
|
+
if (dim !== undefined && val !== undefined) {
|
|
737
|
+
if (dim === USER_DIM) {
|
|
738
|
+
return await ctx.db
|
|
739
|
+
.query("requests")
|
|
740
|
+
.withIndex("userId", (q) => q.eq("userId", val))
|
|
741
|
+
.order("desc")
|
|
742
|
+
.take(limit);
|
|
743
|
+
}
|
|
744
|
+
if (dim === ACTION_DIM) {
|
|
745
|
+
return await ctx.db
|
|
746
|
+
.query("requests")
|
|
747
|
+
.withIndex("actionName", (q) => q.eq("actionName", val))
|
|
748
|
+
.order("desc")
|
|
749
|
+
.take(limit);
|
|
750
|
+
}
|
|
751
|
+
// Custom tag dimension: walk the reverse index, then fetch each request.
|
|
752
|
+
const tagRows = await ctx.db
|
|
753
|
+
.query("requestTags")
|
|
754
|
+
.withIndex("dim_value", (q) => q.eq("dimension", dim).eq("value", val))
|
|
755
|
+
.order("desc")
|
|
756
|
+
.take(limit);
|
|
757
|
+
const rows = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
|
|
758
|
+
return rows.filter((r) => r !== null);
|
|
759
|
+
}
|
|
579
760
|
return await ctx.db.query("requests").order("desc").take(limit);
|
|
580
761
|
},
|
|
581
762
|
});
|
|
582
763
|
const ADMIN_LIST_CAP = 2000;
|
|
583
|
-
|
|
584
|
-
|
|
585
|
-
|
|
764
|
+
// List budget buckets, optionally filtered to one dimension ("user", "action",
|
|
765
|
+
// or any custom tag dimension). Today's spend is zeroed for stale day windows.
|
|
766
|
+
export const listBuckets = query({
|
|
767
|
+
args: { dimension: v.optional(v.string()) },
|
|
768
|
+
handler: async (ctx, args) => {
|
|
586
769
|
// Bounded to avoid an unbounded full-table scan on this reactive query.
|
|
587
|
-
// Paginate (ctx.db.query("
|
|
588
|
-
const
|
|
770
|
+
// Paginate (ctx.db.query("buckets").paginate(...)) for larger deployments.
|
|
771
|
+
const rows = args.dimension !== undefined
|
|
772
|
+
? await ctx.db
|
|
773
|
+
.query("buckets")
|
|
774
|
+
.withIndex("dimension", (q) => q.eq("dimension", args.dimension))
|
|
775
|
+
.take(ADMIN_LIST_CAP)
|
|
776
|
+
: await ctx.db.query("buckets").take(ADMIN_LIST_CAP);
|
|
589
777
|
const today = dayStamp();
|
|
590
|
-
|
|
591
|
-
|
|
592
|
-
|
|
778
|
+
const month = monthStamp();
|
|
779
|
+
return rows.map((b) => ({
|
|
780
|
+
...b,
|
|
781
|
+
spendTodayNanos: b.dayStamp === today ? b.spendTodayNanos : 0,
|
|
782
|
+
spendThisMonthNanos: b.monthStamp === month ? b.spendThisMonthNanos ?? 0 : 0,
|
|
593
783
|
}));
|
|
594
784
|
},
|
|
595
785
|
});
|
|
596
|
-
export const
|
|
786
|
+
export const getBucket = query({
|
|
787
|
+
args: { dimension: v.string(), value: v.string() },
|
|
788
|
+
handler: async (ctx, args) => {
|
|
789
|
+
const b = await getBucketDoc(ctx, args.dimension, args.value);
|
|
790
|
+
if (!b)
|
|
791
|
+
return null;
|
|
792
|
+
const today = dayStamp();
|
|
793
|
+
const month = monthStamp();
|
|
794
|
+
return {
|
|
795
|
+
...b,
|
|
796
|
+
spendTodayNanos: b.dayStamp === today ? b.spendTodayNanos : 0,
|
|
797
|
+
spendThisMonthNanos: b.monthStamp === month ? b.spendThisMonthNanos ?? 0 : 0,
|
|
798
|
+
};
|
|
799
|
+
},
|
|
800
|
+
});
|
|
801
|
+
// Set a bucket's limits/controls. `user` and `action` are just dimensions here;
|
|
802
|
+
// the client's ai.users / ai.actions namespaces are thin wrappers over this.
|
|
803
|
+
export const setBucketLimits = mutation({
|
|
597
804
|
args: {
|
|
598
|
-
|
|
805
|
+
dimension: v.string(),
|
|
806
|
+
value: v.string(),
|
|
599
807
|
requestsPerMinute: v.optional(v.number()),
|
|
808
|
+
maxConcurrent: v.optional(v.number()),
|
|
600
809
|
dailySpendLimitNanos: v.optional(v.number()),
|
|
810
|
+
monthlySpendLimitNanos: v.optional(v.number()),
|
|
601
811
|
lifetimeSpendLimitNanos: v.optional(v.number()),
|
|
602
812
|
dailyTokenLimit: v.optional(v.number()),
|
|
813
|
+
monthlyTokenLimit: v.optional(v.number()),
|
|
603
814
|
lifetimeTokenLimit: v.optional(v.number()),
|
|
815
|
+
warnAtPct: v.optional(v.number()),
|
|
604
816
|
enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
|
|
605
817
|
blocked: v.optional(v.boolean()),
|
|
606
818
|
},
|
|
607
819
|
returns: v.null(),
|
|
608
820
|
handler: async (ctx, args) => {
|
|
609
|
-
const
|
|
610
|
-
const {
|
|
611
|
-
await ctx.db.patch(
|
|
821
|
+
const bucket = await getOrCreateBucket(ctx, args.dimension, args.value);
|
|
822
|
+
const { dimension: _d, value: _v, ...limits } = args;
|
|
823
|
+
await ctx.db.patch(bucket._id, limits);
|
|
612
824
|
return null;
|
|
613
825
|
},
|
|
614
826
|
});
|
|
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
827
|
const vBumpArgs = {
|
|
656
828
|
dailyNanos: v.optional(v.number()),
|
|
829
|
+
monthlyNanos: v.optional(v.number()),
|
|
657
830
|
lifetimeNanos: v.optional(v.number()),
|
|
658
831
|
};
|
|
659
|
-
// One-time "approve another $X" bumps, added on top of
|
|
660
|
-
// changing it. Daily bumps apply to
|
|
661
|
-
|
|
662
|
-
|
|
832
|
+
// One-time "approve another $X" bumps, added on top of a bucket's standing cap
|
|
833
|
+
// without changing it. Daily/monthly bumps apply to the current window only;
|
|
834
|
+
// lifetime bumps persist.
|
|
835
|
+
export const bumpBucket = mutation({
|
|
836
|
+
args: { dimension: v.string(), value: v.string(), ...vBumpArgs },
|
|
663
837
|
returns: v.null(),
|
|
664
|
-
handler: async (ctx, {
|
|
665
|
-
const
|
|
838
|
+
handler: async (ctx, { dimension, value, dailyNanos, monthlyNanos, lifetimeNanos }) => {
|
|
839
|
+
const bucket = await getOrCreateBucket(ctx, dimension, value);
|
|
666
840
|
const today = dayStamp();
|
|
667
|
-
const
|
|
668
|
-
|
|
841
|
+
const month = monthStamp();
|
|
842
|
+
const curDaily = bucket.bumpDayStamp === today ? bucket.dailyBumpNanos ?? 0 : 0;
|
|
843
|
+
const curMonthly = bucket.bumpMonthStamp === month ? bucket.monthlyBumpNanos ?? 0 : 0;
|
|
844
|
+
await ctx.db.patch(bucket._id, {
|
|
669
845
|
bumpDayStamp: today,
|
|
670
846
|
dailyBumpNanos: curDaily + (dailyNanos ?? 0),
|
|
671
|
-
|
|
847
|
+
bumpMonthStamp: month,
|
|
848
|
+
monthlyBumpNanos: curMonthly + (monthlyNanos ?? 0),
|
|
849
|
+
lifetimeBumpNanos: (bucket.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
|
|
672
850
|
});
|
|
673
851
|
return null;
|
|
674
852
|
},
|
|
675
853
|
});
|
|
676
|
-
|
|
677
|
-
|
|
854
|
+
// Manually credit or debit a bucket (comp a user, correct an overcharge).
|
|
855
|
+
// Negative deltaNanos = credit/refund, positive = extra charge. Adjusts the
|
|
856
|
+
// live day/month/lifetime windows AND the durable usage history, and records an
|
|
857
|
+
// audit row. Does not touch the global sharded total or reservations.
|
|
858
|
+
export const adjustBucket = mutation({
|
|
859
|
+
args: {
|
|
860
|
+
dimension: v.string(),
|
|
861
|
+
value: v.string(),
|
|
862
|
+
deltaNanos: v.number(),
|
|
863
|
+
tokens: v.optional(v.number()),
|
|
864
|
+
reason: v.optional(v.string()),
|
|
865
|
+
},
|
|
678
866
|
returns: v.null(),
|
|
679
|
-
handler: async (ctx, {
|
|
680
|
-
const
|
|
867
|
+
handler: async (ctx, { dimension, value, deltaNanos, tokens, reason }) => {
|
|
868
|
+
const b = await getOrCreateBucket(ctx, dimension, value);
|
|
681
869
|
const today = dayStamp();
|
|
682
|
-
const
|
|
683
|
-
|
|
684
|
-
|
|
685
|
-
|
|
686
|
-
|
|
870
|
+
const month = monthStamp();
|
|
871
|
+
const dSame = b.dayStamp === today;
|
|
872
|
+
const mSame = b.monthStamp === month;
|
|
873
|
+
const dt = tokens ?? 0;
|
|
874
|
+
await ctx.db.patch(b._id, {
|
|
875
|
+
totalSpendNanos: Math.max(0, b.totalSpendNanos + deltaNanos),
|
|
876
|
+
totalTokens: Math.max(0, b.totalTokens + dt),
|
|
877
|
+
dayStamp: today,
|
|
878
|
+
monthStamp: month,
|
|
879
|
+
spendTodayNanos: Math.max(0, (dSame ? b.spendTodayNanos : 0) + deltaNanos),
|
|
880
|
+
tokensToday: Math.max(0, (dSame ? b.tokensToday ?? 0 : 0) + dt),
|
|
881
|
+
spendThisMonthNanos: Math.max(0, (mSame ? b.spendThisMonthNanos ?? 0 : 0) + deltaNanos),
|
|
882
|
+
tokensThisMonth: Math.max(0, (mSame ? b.tokensThisMonth ?? 0 : 0) + dt),
|
|
687
883
|
});
|
|
884
|
+
await ctx.db.insert("adjustments", { dimension, value, deltaNanos, tokens: dt, reason });
|
|
885
|
+
await addUsage(ctx, dimension, value, "day", today, deltaNanos, dt, 0);
|
|
886
|
+
await addUsage(ctx, dimension, value, "month", month, deltaNanos, dt, 0);
|
|
688
887
|
return null;
|
|
689
888
|
},
|
|
690
889
|
});
|
|
691
|
-
export const
|
|
692
|
-
args:
|
|
693
|
-
|
|
694
|
-
|
|
695
|
-
|
|
696
|
-
|
|
697
|
-
|
|
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
|
-
},
|
|
890
|
+
export const listAdjustments = query({
|
|
891
|
+
args: { dimension: v.string(), value: v.string(), limit: v.optional(v.number()) },
|
|
892
|
+
handler: async (ctx, { dimension, value, limit }) => ctx.db
|
|
893
|
+
.query("adjustments")
|
|
894
|
+
.withIndex("dim_value", (q) => q.eq("dimension", dimension).eq("value", value))
|
|
895
|
+
.order("desc")
|
|
896
|
+
.take(limit ?? 50),
|
|
709
897
|
});
|
|
710
|
-
|
|
898
|
+
// Durable spend history for a bucket: per-day or per-month rows, newest first.
|
|
899
|
+
// Survives request retention.
|
|
900
|
+
export const usageHistory = query({
|
|
711
901
|
args: {
|
|
712
|
-
|
|
713
|
-
|
|
714
|
-
|
|
715
|
-
|
|
716
|
-
lifetimeTokenLimit: v.optional(v.number()),
|
|
717
|
-
enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
|
|
718
|
-
disabled: v.optional(v.boolean()),
|
|
902
|
+
dimension: v.string(),
|
|
903
|
+
value: v.string(),
|
|
904
|
+
period: v.union(v.literal("day"), v.literal("month")),
|
|
905
|
+
limit: v.optional(v.number()),
|
|
719
906
|
},
|
|
907
|
+
handler: async (ctx, { dimension, value, period, limit }) => ctx.db
|
|
908
|
+
.query("usage")
|
|
909
|
+
.withIndex("bucket_period_stamp", (q) => q.eq("dimension", dimension).eq("value", value).eq("period", period))
|
|
910
|
+
.order("desc")
|
|
911
|
+
.take(limit ?? 90),
|
|
912
|
+
});
|
|
913
|
+
// Deployment-wide default threshold for approaching-limit alerts (fraction of a
|
|
914
|
+
// cap, e.g. 0.8). Buckets can override with their own warnAtPct.
|
|
915
|
+
export const setAlertDefaults = mutation({
|
|
916
|
+
args: { warnAtPct: v.optional(v.number()) },
|
|
720
917
|
returns: v.null(),
|
|
721
|
-
handler: async (ctx,
|
|
722
|
-
const
|
|
723
|
-
|
|
724
|
-
|
|
918
|
+
handler: async (ctx, { warnAtPct }) => {
|
|
919
|
+
const existing = await getSettings(ctx);
|
|
920
|
+
if (existing)
|
|
921
|
+
await ctx.db.patch(existing._id, { defaultWarnAtPct: warnAtPct });
|
|
922
|
+
else
|
|
923
|
+
await ctx.db.insert("settings", { key: "singleton", defaultWarnAtPct: warnAtPct });
|
|
725
924
|
return null;
|
|
726
925
|
},
|
|
727
926
|
});
|
|
927
|
+
// Delete a bucket and (for the `user` dimension) all of that user's request
|
|
928
|
+
// rows — e.g. account deletion / GDPR. Deletes requests in bounded batches and
|
|
929
|
+
// self-reschedules so it never exceeds the per-transaction document limit.
|
|
930
|
+
const DELETE_BATCH = 500;
|
|
931
|
+
export const deleteBucket = mutation({
|
|
932
|
+
args: { dimension: v.string(), value: v.string() },
|
|
933
|
+
returns: v.object({ deletedThisBatch: v.number(), done: v.boolean() }),
|
|
934
|
+
handler: async (ctx, { dimension, value }) => {
|
|
935
|
+
// Only the user dimension owns request rows (indexed by userId). Other
|
|
936
|
+
// dimensions just drop their budget-holder row.
|
|
937
|
+
if (dimension === USER_DIM) {
|
|
938
|
+
const rows = await ctx.db
|
|
939
|
+
.query("requests")
|
|
940
|
+
.withIndex("userId", (q) => q.eq("userId", value))
|
|
941
|
+
.take(DELETE_BATCH);
|
|
942
|
+
for (const r of rows) {
|
|
943
|
+
await deleteRequestTags(ctx, r._id);
|
|
944
|
+
await ctx.db.delete(r._id);
|
|
945
|
+
}
|
|
946
|
+
if (rows.length === DELETE_BATCH) {
|
|
947
|
+
await ctx.scheduler.runAfter(0, api.lib.deleteBucket, {
|
|
948
|
+
dimension,
|
|
949
|
+
value,
|
|
950
|
+
});
|
|
951
|
+
return { deletedThisBatch: rows.length, done: false };
|
|
952
|
+
}
|
|
953
|
+
const bucket = await getBucketDoc(ctx, dimension, value);
|
|
954
|
+
if (bucket)
|
|
955
|
+
await ctx.db.delete(bucket._id);
|
|
956
|
+
return { deletedThisBatch: rows.length + (bucket ? 1 : 0), done: true };
|
|
957
|
+
}
|
|
958
|
+
const bucket = await getBucketDoc(ctx, dimension, value);
|
|
959
|
+
if (bucket)
|
|
960
|
+
await ctx.db.delete(bucket._id);
|
|
961
|
+
return { deletedThisBatch: bucket ? 1 : 0, done: true };
|
|
962
|
+
},
|
|
963
|
+
});
|
|
728
964
|
export const getModelPolicy = query({
|
|
729
965
|
args: {},
|
|
730
966
|
returns: v.object({
|
|
@@ -786,6 +1022,25 @@ export const setGlobalLimits = mutation({
|
|
|
786
1022
|
return null;
|
|
787
1023
|
},
|
|
788
1024
|
});
|
|
1025
|
+
export const bumpGlobal = mutation({
|
|
1026
|
+
args: vBumpArgs,
|
|
1027
|
+
returns: v.null(),
|
|
1028
|
+
handler: async (ctx, { dailyNanos, lifetimeNanos }) => {
|
|
1029
|
+
const today = dayStamp();
|
|
1030
|
+
const s = await getSettings(ctx);
|
|
1031
|
+
const curDaily = s?.globalBumpDayStamp === today ? s?.globalDailyBumpNanos ?? 0 : 0;
|
|
1032
|
+
const patch = {
|
|
1033
|
+
globalBumpDayStamp: today,
|
|
1034
|
+
globalDailyBumpNanos: curDaily + (dailyNanos ?? 0),
|
|
1035
|
+
globalLifetimeBumpNanos: (s?.globalLifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
|
|
1036
|
+
};
|
|
1037
|
+
if (s)
|
|
1038
|
+
await ctx.db.patch(s._id, patch);
|
|
1039
|
+
else
|
|
1040
|
+
await ctx.db.insert("settings", { key: "singleton", ...patch });
|
|
1041
|
+
return null;
|
|
1042
|
+
},
|
|
1043
|
+
});
|
|
789
1044
|
export const setModelPolicy = mutation({
|
|
790
1045
|
args: {
|
|
791
1046
|
mode: v.union(v.literal("open"), v.literal("allowlist"), v.literal("denylist")),
|
|
@@ -815,12 +1070,16 @@ export const setPrice = mutation({
|
|
|
815
1070
|
model: v.string(),
|
|
816
1071
|
inputNanosPerMTok: v.number(),
|
|
817
1072
|
outputNanosPerMTok: v.number(),
|
|
1073
|
+
// optional cache-read rate; if omitted, a default discount off input applies
|
|
1074
|
+
cachedNanosPerMTok: v.optional(v.number()),
|
|
818
1075
|
},
|
|
819
1076
|
returns: v.null(),
|
|
820
1077
|
handler: async (ctx, args) => {
|
|
821
1078
|
// Negative prices would make costOf return a negative cost, which folds
|
|
822
1079
|
// into totals as a spend *refund* — pushing a user back under their cap.
|
|
823
|
-
if (args.inputNanosPerMTok < 0 ||
|
|
1080
|
+
if (args.inputNanosPerMTok < 0 ||
|
|
1081
|
+
args.outputNanosPerMTok < 0 ||
|
|
1082
|
+
(args.cachedNanosPerMTok ?? 0) < 0) {
|
|
824
1083
|
throw new Error("Prices must be non-negative");
|
|
825
1084
|
}
|
|
826
1085
|
const existing = await ctx.db
|
|
@@ -848,6 +1107,7 @@ export const listPrices = query({
|
|
|
848
1107
|
merged[o.model] = {
|
|
849
1108
|
input: o.inputNanosPerMTok,
|
|
850
1109
|
output: o.outputNanosPerMTok,
|
|
1110
|
+
cached: o.cachedNanosPerMTok,
|
|
851
1111
|
overridden: true,
|
|
852
1112
|
};
|
|
853
1113
|
}
|