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

package/dist/component/lib.js

@@ -1,7 +1,7 @@
 import { v } from "convex/values";
 import { mutation, internalMutation, query, } from "./_generated/server";
 import { api, internal, components } from "./_generated/api";
-import { vMessage } from "./schema";
+import { vMessage, vTag } from "./schema";
 import { ShardedCounter } from "@convex-dev/sharded-counter";
 // All money is integer **nanodollars** (1 USD = 1e9 nano). Integers avoid the
 // rounding drift that floating-point cents accumulate over millions of
@@ -11,6 +11,12 @@ import { ShardedCounter } from "@convex-dev/sharded-counter";
 // currency code alongside these amounts and convert here, at the one boundary.
 const NANOS_PER_DOLLAR = 1e9;
 const fmtUsd = (nanos) => `$${(nanos / NANOS_PER_DOLLAR).toFixed(4)}`;
+// Built-in attribution dimensions. `user` and `action` are always populated
+// from a request's userId/actionName; apps can add any other dimensions
+// (team, project, customer, env, …) as tags. These two names are reserved —
+// tags carrying them are ignored in favor of the first-class fields.
+const USER_DIM = "user";
+const ACTION_DIM = "action";
 // Deployment-wide spend totals (nanodollars), sharded for high write throughput.
 // Keyed "total" (lifetime) and "day:<UTC date>" (natural daily reset).
 const globalSpend = new ShardedCounter(components.shardedCounter);
@@ -26,9 +32,17 @@ const DEFAULT_PRICES = {
     "openai/gpt-5": { input: 1_250_000_000, output: 10_000_000_000 },
     "openai/gpt-5-mini": { input: 250_000_000, output: 2_000_000_000 },
 };
-// Pessimistic assumed output length when reserving budget up front. Reservations
-// only need to be an upper bound often enough to keep a hard cap honest; the
-// actual cost replaces the estimate the moment the request settles.
+// Per-call price (nanodollars) for provider server-side tools that bill a fee on
+// top of tokens — e.g. Anthropic web search at ~$0.01/call. Keyed by the tool
+// name the caller reports in `serverToolUses` (e.g. { web_search: 3 }). Used
+// only when a request settles WITHOUT an authoritative gateway cost; if you pass
+// `costNanos`, that already includes tool fees. Override via setServerToolPrice.
+const DEFAULT_SERVER_TOOL_PRICES = {
+    web_search: 10_000_000, // $0.01 per search
+};
+// Pessimistic assumed output length when reserving budget up front. This makes
+// concurrent admission atomic against the estimate; a response that exceeds the
+// estimate can still settle above the cap by the estimation delta.
 const ESTIMATED_OUTPUT_TOKENS = 800;
 // A request still "pending" after this long is presumed dead (its action
 // crashed before settling); the reconciler releases its reservation. Set well
@@ -42,7 +56,12 @@ const STALE_PENDING_MS = 30 * 60 * 1000;
 // audit table — and the sensitive content in it — from growing without bound.
 // Override per-deployment via setRetention.
 const DEFAULT_RETENTION_MS = 60 * 60 * 1000; // 1 hour
-const dayStamp = () => new Date().toISOString().slice(0, 10);
+const dayStamp = () => new Date().toISOString().slice(0, 10); // "2026-09-04"
+const monthStamp = () => new Date().toISOString().slice(0, 7); // "2026-09"
+// Cached (prompt-cache-read) input tokens are billed far below the normal input
+// rate. When a model's price has no explicit cachedNanosPerMTok, charge this
+// fraction of its input rate (providers commonly discount ~90%).
+const CACHE_DISCOUNT = 0.1;
 // Conservative fallback for any model not in the price table: the max of every
 // known price dimension. Falling back to 0 would be fail-open — an unpriced
 // model would reserve 0, pass every cap, and log 0¢ while the AI Gateway still
@@ -62,17 +81,76 @@ async function getPrice(ctx, model) {
         return {
             input: override.inputNanosPerMTok,
             output: override.outputNanosPerMTok,
+            cached: override.cachedNanosPerMTok,
             known: true,
         };
     }
     const known = DEFAULT_PRICES[model];
     if (known)
-        return { ...known, known: true };
-    return { ...CONSERVATIVE_PRICE, known: false };
+        return { ...known, cached: undefined, known: true };
+    return { ...CONSERVATIVE_PRICE, cached: undefined, known: false };
 }
+// The per-Mtok rate for cached (prompt-cache-read) tokens: an explicit override,
+// else a discount off the input rate.
+const cachedRate = (p) => p.cached ?? Math.round(p.input * CACHE_DISCOUNT);
 // Integer nanodollars. Divide-before-multiply keeps the intermediate product
 // within 2^53 even for large token counts × large per-Mtok prices.
 const costOf = (inputTokens, outputTokens, price) => Math.round((inputTokens / 1e6) * price.input + (outputTokens / 1e6) * price.output);
+// Cache-aware settle cost: the cached slice of the prompt is billed at the
+// (discounted) cache rate, the rest of the prompt at the input rate, and
+// completions at the output rate. `cachedTokens` is the gateway's real
+// prompt-cache-read count (usage.inputTokenDetails.cacheReadTokens).
+const settleCost = (promptTokens, cachedTokens, completionTokens, price) => {
+    const cached = Math.min(Math.max(0, cachedTokens), Math.max(0, promptTokens));
+    const fresh = Math.max(0, promptTokens - cached);
+    return Math.round((fresh / 1e6) * price.input +
+        (cached / 1e6) * cachedRate(price) +
+        (completionTokens / 1e6) * price.output);
+};
+// Per-call fees for provider server tools (web search, etc.), merging the
+// defaults with any deployment overrides. Unknown tools price at 0 (recorded
+// but not charged) rather than guessing.
+const serverToolCost = (uses, overrides) => {
+    if (!uses)
+        return 0;
+    const prices = { ...DEFAULT_SERVER_TOOL_PRICES, ...(overrides ?? {}) };
+    let total = 0;
+    for (const [tool, count] of Object.entries(uses)) {
+        if (count > 0 && prices[tool] > 0)
+            total += Math.round(count * prices[tool]);
+    }
+    return total;
+};
+// Upsert-add a settled amount into the durable per-(bucket, period) usage row.
+// These rows are never swept by request retention, so spend history survives.
+async function addUsage(ctx, dimension, value, period, stamp, spendNanos, tokens, requests) {
+    const existing = await ctx.db
+        .query("usage")
+        .withIndex("bucket_period_stamp", (q) => q
+        .eq("dimension", dimension)
+        .eq("value", value)
+        .eq("period", period)
+        .eq("stamp", stamp))
+        .unique();
+    if (existing) {
+        await ctx.db.patch(existing._id, {
+            spendNanos: existing.spendNanos + spendNanos,
+            tokens: existing.tokens + tokens,
+            requests: existing.requests + requests,
+        });
+    }
+    else {
+        await ctx.db.insert("usage", {
+            dimension,
+            value,
+            period,
+            stamp,
+            spendNanos,
+            tokens,
+            requests,
+        });
+    }
+}
 // Up-front estimate of a request's cost and token count, reserved before the
 // call so concurrent in-flight requests are visible to each other's caps.
 function estimateUsage(messages, price) {
@@ -83,63 +161,105 @@ function estimateUsage(messages, price) {
         tokens: inputTokens + ESTIMATED_OUTPUT_TOKENS,
     };
 }
-// Evaluate an entity's (user or action) spend + token budgets against
-// committed + reserved + this request's estimate. Returns a hard rejection
-// (block) or a list of soft warnings (allow), per the entity's enforcement.
+// The full set of attribution buckets a request touches: the built-in `user`
+// and `action` dimensions plus any extra tags. Reserved dimensions in `extra`
+// are dropped (userId/actionName own them), and (dimension, value) pairs are
+// de-duplicated. Used identically at reserve time (startRequest) and settle
+// time (foldOne), so a request always settles exactly the buckets it reserved.
+function requestBuckets(userId, actionName, extra) {
+    const out = [{ dimension: USER_DIM, value: userId }];
+    if (actionName !== undefined)
+        out.push({ dimension: ACTION_DIM, value: actionName });
+    for (const t of extra ?? []) {
+        if (t.dimension === USER_DIM || t.dimension === ACTION_DIM)
+            continue;
+        if (!t.dimension || !t.value)
+            continue;
+        if (out.some((x) => x.dimension === t.dimension && x.value === t.value))
+            continue;
+        out.push({ dimension: t.dimension, value: t.value });
+    }
+    return out;
+}
+// Drop reserved/empty/duplicate tags from a caller-supplied list, leaving the
+// "extra" dimensions stored on the request row.
+function sanitizeExtraTags(extra) {
+    const out = [];
+    for (const t of extra ?? []) {
+        if (t.dimension === USER_DIM || t.dimension === ACTION_DIM)
+            continue;
+        if (!t.dimension || !t.value)
+            continue;
+        if (out.some((x) => x.dimension === t.dimension && x.value === t.value))
+            continue;
+        out.push({ dimension: t.dimension, value: t.value });
+    }
+    return out;
+}
+// Evaluate a bucket's spend + token budgets against committed + reserved + this
+// request's estimate. Returns a hard rejection (block), soft warnings (allow),
+// and threshold notices (approaching a cap — see warnAtPct). Each window (daily,
+// monthly, lifetime) × kind (spend, token) is one check.
 function evaluateCaps(o) {
+    // { code, projected usage (incl. this estimate), cap, human window label,
+    //   whether it's a money cap (formatted as $), spend? for notices }
+    const checks = [
+        { w: "daily_spend_limit", used: o.spendToday + o.reservedSpendToday + o.estCost, cap: o.dailySpendLimitNanos, label: "daily spend limit", money: true },
+        { w: "monthly_spend_limit", used: o.spendThisMonth + o.reservedSpendMonth + o.estCost, cap: o.monthlySpendLimitNanos, label: "monthly spend limit", money: true },
+        { w: "lifetime_spend_limit", used: o.totalSpend + o.reservedSpendTotal + o.estCost, cap: o.lifetimeSpendLimitNanos, label: "lifetime spend limit", money: true },
+        { w: "daily_token_limit", used: o.tokensToday + o.reservedTokensToday + o.estTokens, cap: o.dailyTokenLimit, label: "daily token limit", money: false },
+        { w: "monthly_token_limit", used: o.tokensThisMonth + o.reservedTokensMonth + o.estTokens, cap: o.monthlyTokenLimit, label: "monthly token limit", money: false },
+        { w: "lifetime_token_limit", used: o.totalTokens + o.reservedTokensTotal + o.estTokens, cap: o.lifetimeTokenLimit, label: "lifetime token limit", money: false },
+    ];
     const violations = [];
-    const push = (code, reason) => violations.push({ code: `${o.label}_${code}`, reason });
-    if (o.dailySpendLimitNanos !== undefined &&
-        o.spendToday + o.reservedSpendToday + o.estCost > o.dailySpendLimitNanos)
-        push("daily_spend_limit", `Daily spend limit reached for ${o.label} "${o.name}" (${fmtUsd(o.dailySpendLimitNanos)}/day)`);
-    if (o.lifetimeSpendLimitNanos !== undefined &&
-        o.totalSpend + o.reservedSpendTotal + o.estCost > o.lifetimeSpendLimitNanos)
-        push("lifetime_spend_limit", `Lifetime spend limit reached for ${o.label} "${o.name}"`);
-    if (o.dailyTokenLimit !== undefined &&
-        o.tokensToday + o.reservedTokensToday + o.estTokens > o.dailyTokenLimit)
-        push("daily_token_limit", `Daily token limit reached for ${o.label} "${o.name}" (${o.dailyTokenLimit}/day)`);
-    if (o.lifetimeTokenLimit !== undefined &&
-        o.totalTokens + o.reservedTokensTotal + o.estTokens > o.lifetimeTokenLimit)
-        push("lifetime_token_limit", `Lifetime token limit reached for ${o.label} "${o.name}"`);
+    const notices = [];
+    const pct = o.warnAtPct;
+    for (const c of checks) {
+        if (c.cap === undefined)
+            continue;
+        const capStr = c.money ? `${fmtUsd(c.cap)}` : `${c.cap} tokens`;
+        if (c.used > c.cap) {
+            violations.push({
+                code: `${o.label}_${c.w}`,
+                reason: `${cap(c.label)} reached for ${o.label} "${o.name}" (${capStr})`,
+            });
+        }
+        else if (pct !== undefined && pct > 0 && pct < 1 && c.used >= pct * c.cap) {
+            notices.push(`${o.label} "${o.name}" at ${Math.round((c.used / c.cap) * 100)}% of ${c.label} (${capStr})`);
+        }
+    }
     if (violations.length === 0)
-        return { warnings: [] };
+        return { warnings: [], notices };
     if (o.enforcement === "soft")
-        return { warnings: violations.map((v) => v.reason) };
-    return { hard: violations[0], warnings: [] };
+        return { warnings: violations.map((v) => v.reason), notices };
+    return { hard: violations[0], warnings: [], notices };
 }
+const cap = (s) => s.charAt(0).toUpperCase() + s.slice(1);
 // A cap plus any one-time bump. Returns undefined when there's no base cap
 // (a bump alone never creates a cap).
 const withBump = (base, bump) => base === undefined ? undefined : base + (bump ?? 0);
 const hasAnyCap = (e) => e.dailySpendLimitNanos !== undefined ||
+    e.monthlySpendLimitNanos !== undefined ||
     e.lifetimeSpendLimitNanos !== undefined ||
     e.dailyTokenLimit !== undefined ||
+    e.monthlyTokenLimit !== undefined ||
     e.lifetimeTokenLimit !== undefined;
-async function getOrCreateUser(ctx, userId) {
-    const existing = await ctx.db
-        .query("users")
-        .withIndex("userId", (q) => q.eq("userId", userId))
+// A bucket needs a reservation row-write if it has any spend/token cap OR a
+// concurrency cap (which reads pendingCount, incremented at reserve time).
+const needsReserve = (e) => hasAnyCap(e) || e.maxConcurrent !== undefined;
+async function getBucketDoc(ctx, dimension, value) {
+    return await ctx.db
+        .query("buckets")
+        .withIndex("dim_value", (q) => q.eq("dimension", dimension).eq("value", value))
         .unique();
-    if (existing)
-        return existing;
-    const id = await ctx.db.insert("users", {
-        userId,
-        totalSpendNanos: 0,
-        totalRequests: 0,
-        totalTokens: 0,
-        dayStamp: dayStamp(),
-        spendTodayNanos: 0,
-    });
-    return (await ctx.db.get(id));
 }
-async function getOrCreateAction(ctx, name) {
-    const existing = await ctx.db
-        .query("actions")
-        .withIndex("name", (q) => q.eq("name", name))
-        .unique();
+async function getOrCreateBucket(ctx, dimension, value) {
+    const existing = await getBucketDoc(ctx, dimension, value);
     if (existing)
         return existing;
-    const id = await ctx.db.insert("actions", {
-        name,
+    const id = await ctx.db.insert("buckets", {
+        dimension,
+        value,
         totalSpendNanos: 0,
         totalRequests: 0,
         totalTokens: 0,
@@ -154,10 +274,21 @@ async function getSettings(ctx) {
         .withIndex("key", (q) => q.eq("key", "singleton"))
         .unique();
 }
+// Delete the reverse-index rows for a request (called when the request row is
+// deleted, so the tag index never outlives the request it points at).
+async function deleteRequestTags(ctx, requestId) {
+    const tags = await ctx.db
+        .query("requestTags")
+        .withIndex("requestId", (q) => q.eq("requestId", requestId))
+        .collect();
+    for (const t of tags)
+        await ctx.db.delete(t._id);
+}
 const vStartResult = v.union(v.object({
     allowed: v.literal(true),
     requestId: v.id("requests"),
-    warnings: v.array(v.string()),
+    warnings: v.array(v.string()), // soft caps exceeded (allowed with a warning)
+    notices: v.array(v.string()), // approaching a cap (warnAtPct threshold)
 }), v.object({
     allowed: v.literal(false),
     code: v.string(),
@@ -167,13 +298,16 @@ export const startRequest = mutation({
     args: {
         userId: v.string(),
         actionName: v.optional(v.string()),
+        // Extra attribution dimensions to bill/limit (team, customer, env, …).
+        // `user` and `action` are reserved (owned by userId/actionName).
+        tags: v.optional(v.array(vTag)),
         model: v.string(),
         messages: v.array(vMessage),
         rerunOf: v.optional(v.id("requests")),
     },
     returns: vStartResult,
     handler: async (ctx, args) => {
-        const user = await getOrCreateUser(ctx, args.userId);
+        const extraTags = sanitizeExtraTags(args.tags);
         // Record the blocked attempt and return a rejection (throwing would roll
         // back the record). `persist` is false for the high-frequency-by-design
         // rejections (rate limit, blocked user) that a client retries in a tight
@@ -182,7 +316,12 @@ export const startRequest = mutation({
         const reject = async (code, reason, persist = true) => {
             if (persist) {
                 await ctx.db.insert("requests", {
-                    ...args,
+                    userId: args.userId,
+                    actionName: args.actionName,
+                    ...(extraTags.length ? { tags: extraTags } : {}),
+                    model: args.model,
+                    messages: args.messages,
+                    rerunOf: args.rerunOf,
                     status: "blocked",
                     error: reason,
                 });
@@ -190,17 +329,17 @@ export const startRequest = mutation({
             return { allowed: false, code, reason };
         };
         const today = dayStamp();
+        const month = monthStamp();
         const priceInfo = await getPrice(ctx, args.model);
         const est = estimateUsage(args.messages, priceInfo);
         const warnings = [];
-        if (user.blocked) {
-            return reject("blocked", `User "${args.userId}" is blocked`, false);
-        }
+        const notices = [];
         // Model allow/deny policy (component-wide).
-        const policy = await getSettings(ctx);
-        if (policy) {
-            const mode = policy.modelMode ?? "open";
-            const list = policy.models ?? [];
+        const settings = await getSettings(ctx);
+        const defaultWarnAtPct = settings?.defaultWarnAtPct;
+        if (settings) {
+            const mode = settings.modelMode ?? "open";
+            const list = settings.models ?? [];
             if (mode === "allowlist" && !list.includes(args.model)) {
                 return reject("model_not_allowed", `Model "${args.model}" is not on the allowlist`);
             }
@@ -208,150 +347,196 @@ export const startRequest = mutation({
                 return reject("model_denied", `Model "${args.model}" is denied`);
             }
         }
-        if (user.requestsPerMinute !== undefined) {
-            // Bounded read: we only need to know whether non-blocked requests in the
-            // window have reached the limit. Reading limit+1 non-blocked rows is
-            // enough, and .take caps the scan even if the window is flooded. Blocked
-            // rows aren't persisted for rate-limit/blocked rejections (see reject),
-            // so the window stays small.
-            const recent = await ctx.db
-                .query("requests")
-                .withIndex("userId", (q) => q.eq("userId", args.userId).gt("_creationTime", Date.now() - 60_000))
-                .take(user.requestsPerMinute + 50);
-            if (recent.filter((r) => r.status !== "blocked").length >=
-                user.requestsPerMinute) {
-                return reject("rate_limit", `Rate limit exceeded for "${args.userId}" (${user.requestsPerMinute}/min)`, false);
+        // Fetch/create every bucket this request is attributed to (user, action,
+        // and any extra tags). Each may carry its own budget.
+        const bucketTags = requestBuckets(args.userId, args.actionName, extraTags);
+        const buckets = [];
+        for (const t of bucketTags) {
+            buckets.push(await getOrCreateBucket(ctx, t.dimension, t.value));
+        }
+        // A hard block on ANY bucket rejects the request. The user dimension's block
+        // isn't persisted (retried in a loop); config-level blocks on other
+        // dimensions are rarer, so they persist for the audit log.
+        for (const b of buckets) {
+            if (b.blocked) {
+                const label = b.dimension === USER_DIM ? "User" : b.dimension;
+                return reject(`${b.dimension}_blocked`, `${label} "${b.value}" is blocked`, b.dimension !== USER_DIM);
             }
         }
-        // Committed + already-reserved in-flight usage + this estimate must fit
-        // under each cap. Decided and reserved in one transaction, so Convex's
-        // serializable isolation makes it a true atomic check-and-reserve. Soft
-        // enforcement turns a violation into a warning instead of a block.
-        const userSameDay = user.dayStamp === today;
-        const userEval = evaluateCaps({
-            label: "user",
-            name: args.userId,
-            enforcement: user.enforcement ?? "hard",
-            estCost: est.cost,
-            estTokens: est.tokens,
-            spendToday: userSameDay ? user.spendTodayNanos : 0,
-            reservedSpendToday: userSameDay ? user.reservedTodayNanos ?? 0 : 0,
-            totalSpend: user.totalSpendNanos,
-            reservedSpendTotal: user.reservedTotalNanos ?? 0,
-            tokensToday: userSameDay ? user.tokensToday ?? 0 : 0,
-            reservedTokensToday: userSameDay ? user.reservedTodayTokens ?? 0 : 0,
-            totalTokens: user.totalTokens,
-            reservedTokensTotal: user.reservedTotalTokens ?? 0,
-            dailySpendLimitNanos: withBump(user.dailySpendLimitNanos, user.bumpDayStamp === today ? user.dailyBumpNanos : 0),
-            lifetimeSpendLimitNanos: withBump(user.lifetimeSpendLimitNanos, user.lifetimeBumpNanos),
-            dailyTokenLimit: user.dailyTokenLimit,
-            lifetimeTokenLimit: user.lifetimeTokenLimit,
-        });
-        if (userEval.hard)
-            return reject(userEval.hard.code, userEval.hard.reason);
-        warnings.push(...userEval.warnings);
-        let action = null;
-        if (args.actionName !== undefined) {
-            action = await getOrCreateAction(ctx, args.actionName);
-            if (action.disabled) {
-                return reject("action_disabled", `Action "${action.name}" is disabled`);
+        // Concurrency cap: reject when a bucket already has maxConcurrent requests
+        // in flight (pendingCount). A transient limit like rate-limiting, so it
+        // isn't persisted (the caller retries once something settles).
+        for (const b of buckets) {
+            if (b.maxConcurrent !== undefined && (b.pendingCount ?? 0) >= b.maxConcurrent) {
+                return reject(`${b.dimension}_max_concurrent`, `Too many concurrent requests for ${b.dimension} "${b.value}" (max ${b.maxConcurrent})`, false);
+            }
+        }
+        // Enforce the rolling 60-second rate limit on every configured dimension.
+        // User/action requests use their first-class indexes; custom dimensions use
+        // the requestTags reverse index. All reads are bounded by the configured
+        // limit (plus a small allowance for persisted blocked attempts).
+        const rateCutoff = Date.now() - 60_000;
+        for (const b of buckets) {
+            const limit = b.requestsPerMinute;
+            if (limit === undefined)
+                continue;
+            let recentCount;
+            if (b.dimension === USER_DIM) {
+                const recent = await ctx.db
+                    .query("requests")
+                    .withIndex("userId", (q) => q.eq("userId", b.value).gt("_creationTime", rateCutoff))
+                    .take(limit + 50);
+                recentCount = recent.filter((r) => r.status !== "blocked").length;
+            }
+            else if (b.dimension === ACTION_DIM) {
+                const recent = await ctx.db
+                    .query("requests")
+                    .withIndex("actionName", (q) => q.eq("actionName", b.value).gt("_creationTime", rateCutoff))
+                    .take(limit + 50);
+                recentCount = recent.filter((r) => r.status !== "blocked").length;
             }
-            const aSameDay = action.dayStamp === today;
-            const actionEval = evaluateCaps({
-                label: "action",
-                name: action.name,
-                enforcement: action.enforcement ?? "hard",
+            else {
+                recentCount = (await ctx.db
+                    .query("requestTags")
+                    .withIndex("dim_value", (q) => q
+                    .eq("dimension", b.dimension)
+                    .eq("value", b.value)
+                    .gt("_creationTime", rateCutoff))
+                    .take(limit)).length;
+            }
+            if (recentCount >= limit) {
+                const code = b.dimension === USER_DIM ? "rate_limit" : `${b.dimension}_rate_limit`;
+                return reject(code, `Rate limit exceeded for ${b.dimension} "${b.value}" (${limit}/min)`, false);
+            }
+        }
+        // Committed + already-reserved in-flight usage + this estimate must fit
+        // under EACH bucket's cap. Decided and reserved in one transaction, so
+        // Convex's serializable isolation makes concurrent admission atomic across
+        // every capped bucket. Final usage can exceed the estimate; settlement then
+        // records the actual amount. Soft enforcement turns a violation into a
+        // warning instead of a block.
+        for (const b of buckets) {
+            const sameDay = b.dayStamp === today;
+            const sameMonth = b.monthStamp === month;
+            const ev = evaluateCaps({
+                label: b.dimension,
+                name: b.value,
+                enforcement: b.enforcement ?? "hard",
+                warnAtPct: b.warnAtPct ?? defaultWarnAtPct,
                 estCost: est.cost,
                 estTokens: est.tokens,
-                spendToday: aSameDay ? action.spendTodayNanos : 0,
-                reservedSpendToday: aSameDay ? action.reservedTodayNanos ?? 0 : 0,
-                totalSpend: action.totalSpendNanos,
-                reservedSpendTotal: action.reservedTotalNanos ?? 0,
-                tokensToday: aSameDay ? action.tokensToday ?? 0 : 0,
-                reservedTokensToday: aSameDay ? action.reservedTodayTokens ?? 0 : 0,
-                totalTokens: action.totalTokens,
-                reservedTokensTotal: action.reservedTotalTokens ?? 0,
-                dailySpendLimitNanos: withBump(action.dailySpendLimitNanos, action.bumpDayStamp === today ? action.dailyBumpNanos : 0),
-                lifetimeSpendLimitNanos: withBump(action.lifetimeSpendLimitNanos, action.lifetimeBumpNanos),
-                dailyTokenLimit: action.dailyTokenLimit,
-                lifetimeTokenLimit: action.lifetimeTokenLimit,
+                spendToday: sameDay ? b.spendTodayNanos : 0,
+                reservedSpendToday: sameDay ? b.reservedTodayNanos ?? 0 : 0,
+                spendThisMonth: sameMonth ? b.spendThisMonthNanos ?? 0 : 0,
+                reservedSpendMonth: sameMonth ? b.reservedMonthNanos ?? 0 : 0,
+                totalSpend: b.totalSpendNanos,
+                reservedSpendTotal: b.reservedTotalNanos ?? 0,
+                tokensToday: sameDay ? b.tokensToday ?? 0 : 0,
+                reservedTokensToday: sameDay ? b.reservedTodayTokens ?? 0 : 0,
+                tokensThisMonth: sameMonth ? b.tokensThisMonth ?? 0 : 0,
+                reservedTokensMonth: sameMonth ? b.reservedMonthTokens ?? 0 : 0,
+                totalTokens: b.totalTokens,
+                reservedTokensTotal: b.reservedTotalTokens ?? 0,
+                dailySpendLimitNanos: withBump(b.dailySpendLimitNanos, b.bumpDayStamp === today ? b.dailyBumpNanos : 0),
+                monthlySpendLimitNanos: withBump(b.monthlySpendLimitNanos, b.bumpMonthStamp === month ? b.monthlyBumpNanos : 0),
+                lifetimeSpendLimitNanos: withBump(b.lifetimeSpendLimitNanos, b.lifetimeBumpNanos),
+                dailyTokenLimit: b.dailyTokenLimit,
+                monthlyTokenLimit: b.monthlyTokenLimit,
+                lifetimeTokenLimit: b.lifetimeTokenLimit,
             });
-            if (actionEval.hard)
-                return reject(actionEval.hard.code, actionEval.hard.reason);
-            warnings.push(...actionEval.warnings);
+            if (ev.hard)
+                return reject(ev.hard.code, ev.hard.reason);
+            warnings.push(...ev.warnings);
+            notices.push(...ev.notices);
         }
-        // Deployment-wide ("global") spend cap — same reserve-then-settle model and
-        // the SAME evaluateCaps logic as the per-entity caps above, so the guarantee
-        // statement is uniform: a request is admitted only if
+        // Deployment-wide ("global") spend cap — the same estimated-usage admission
+        // rule as the per-bucket caps above: a request is admitted only if
         // committed + reserved + estimate <= cap. The ONE difference is the holder:
-        // per-user/per-action reserve on a single row (an exact atomic
-        // check-and-reserve), while the global holder is a sharded counter for
+        // per-bucket caps reserve on a single row (an atomic check-and-reserve),
+        // while the global holder is a sharded counter for
         // throughput — its committed total is read as an eventually-consistent sum
         // with no cross-request reservation, so a hard global cap can overshoot by a
-        // bounded amount under burst. That's the deliberate exactness/throughput
+        // bounded amount under burst. That's the deliberate consistency/throughput
         // trade for a deployment-wide killswitch; it's the only approximate scope.
-        if (policy &&
-            (policy.globalDailySpendLimitNanos !== undefined ||
-                policy.globalLifetimeSpendLimitNanos !== undefined)) {
+        if (settings &&
+            (settings.globalDailySpendLimitNanos !== undefined ||
+                settings.globalLifetimeSpendLimitNanos !== undefined)) {
             const globalEval = evaluateCaps({
                 label: "global",
                 name: "deployment",
-                enforcement: policy.globalEnforcement ?? "hard",
+                enforcement: settings.globalEnforcement ?? "hard",
+                warnAtPct: defaultWarnAtPct,
                 estCost: est.cost,
                 estTokens: est.tokens,
                 spendToday: await globalSpend.count(ctx, globalDayKey(today)),
                 reservedSpendToday: 0, // sharded holder: no cross-request reservation
+                spendThisMonth: 0, // global tracks daily + lifetime only
+                reservedSpendMonth: 0,
                 totalSpend: await globalSpend.count(ctx, GLOBAL_TOTAL),
                 reservedSpendTotal: 0,
                 tokensToday: 0,
                 reservedTokensToday: 0,
+                tokensThisMonth: 0,
+                reservedTokensMonth: 0,
                 totalTokens: 0,
                 reservedTokensTotal: 0,
-                dailySpendLimitNanos: withBump(policy.globalDailySpendLimitNanos, policy.globalBumpDayStamp === today ? policy.globalDailyBumpNanos : 0),
-                lifetimeSpendLimitNanos: withBump(policy.globalLifetimeSpendLimitNanos, policy.globalLifetimeBumpNanos),
+                dailySpendLimitNanos: withBump(settings.globalDailySpendLimitNanos, settings.globalBumpDayStamp === today ? settings.globalDailyBumpNanos : 0),
+                lifetimeSpendLimitNanos: withBump(settings.globalLifetimeSpendLimitNanos, settings.globalLifetimeBumpNanos),
             });
             if (globalEval.hard)
                 return reject(globalEval.hard.code, globalEval.hard.reason);
             warnings.push(...globalEval.warnings);
+            notices.push(...globalEval.notices);
         }
-        // Passed — reserve, but ONLY on entities that actually have a cap. Writing
-        // an uncapped entity's row here would serialize every request that shares it
-        // (e.g. all callers of one action); with no cap there's no reserved amount
-        // to consult. Totals are still accrued later, asynchronously, in foldOne.
-        if (hasAnyCap(user)) {
-            await ctx.db.patch(user._id, {
+        // Passed — reserve, but ONLY on buckets that actually have a cap. Writing an
+        // uncapped bucket's row here would serialize every request that shares it
+        // (e.g. all callers of one action, or every request in one env); with no cap
+        // there's no reserved amount to consult. Totals are still accrued later,
+        // asynchronously, in foldOne — for every bucket, capped or not.
+        for (const b of buckets) {
+            if (!needsReserve(b))
+                continue;
+            const sameDay = b.dayStamp === today;
+            const sameMonth = b.monthStamp === month;
+            await ctx.db.patch(b._id, {
                 dayStamp: today,
-                spendTodayNanos: userSameDay ? user.spendTodayNanos : 0,
-                tokensToday: userSameDay ? user.tokensToday ?? 0 : 0,
-                reservedTodayNanos: (userSameDay ? user.reservedTodayNanos ?? 0 : 0) + est.cost,
-                reservedTotalNanos: (user.reservedTotalNanos ?? 0) + est.cost,
-                reservedTodayTokens: (userSameDay ? user.reservedTodayTokens ?? 0 : 0) + est.tokens,
-                reservedTotalTokens: (user.reservedTotalTokens ?? 0) + est.tokens,
-                pendingCount: (user.pendingCount ?? 0) + 1,
-            });
-        }
-        if (action && hasAnyCap(action)) {
-            const aSameDay = action.dayStamp === today;
-            await ctx.db.patch(action._id, {
-                dayStamp: today,
-                spendTodayNanos: aSameDay ? action.spendTodayNanos : 0,
-                tokensToday: aSameDay ? action.tokensToday ?? 0 : 0,
-                reservedTodayNanos: (aSameDay ? action.reservedTodayNanos ?? 0 : 0) + est.cost,
-                reservedTotalNanos: (action.reservedTotalNanos ?? 0) + est.cost,
-                reservedTodayTokens: (aSameDay ? action.reservedTodayTokens ?? 0 : 0) + est.tokens,
-                reservedTotalTokens: (action.reservedTotalTokens ?? 0) + est.tokens,
-                pendingCount: (action.pendingCount ?? 0) + 1,
+                monthStamp: month,
+                spendTodayNanos: sameDay ? b.spendTodayNanos : 0,
+                tokensToday: sameDay ? b.tokensToday ?? 0 : 0,
+                spendThisMonthNanos: sameMonth ? b.spendThisMonthNanos ?? 0 : 0,
+                tokensThisMonth: sameMonth ? b.tokensThisMonth ?? 0 : 0,
+                reservedTodayNanos: (sameDay ? b.reservedTodayNanos ?? 0 : 0) + est.cost,
+                reservedMonthNanos: (sameMonth ? b.reservedMonthNanos ?? 0 : 0) + est.cost,
+                reservedTotalNanos: (b.reservedTotalNanos ?? 0) + est.cost,
+                reservedTodayTokens: (sameDay ? b.reservedTodayTokens ?? 0 : 0) + est.tokens,
+                reservedMonthTokens: (sameMonth ? b.reservedMonthTokens ?? 0 : 0) + est.tokens,
+                reservedTotalTokens: (b.reservedTotalTokens ?? 0) + est.tokens,
+                pendingCount: (b.pendingCount ?? 0) + 1,
             });
         }
         const requestId = await ctx.db.insert("requests", {
-            ...args,
+            userId: args.userId,
+            actionName: args.actionName,
+            ...(extraTags.length ? { tags: extraTags } : {}),
+            model: args.model,
+            messages: args.messages,
+            rerunOf: args.rerunOf,
             status: "pending",
             estimatedNanos: est.cost,
             estimatedTokens: est.tokens,
             ...(priceInfo.known ? {} : { unpricedModel: true }),
             ...(warnings.length > 0 ? { overBudget: true } : {}),
         });
-        return { allowed: true, requestId, warnings };
+        // Reverse index so the request log can be filtered by any extra tag
+        // dimension (user/action are already indexed on the requests table).
+        for (const t of extraTags) {
+            await ctx.db.insert("requestTags", {
+                dimension: t.dimension,
+                value: t.value,
+                requestId,
+            });
+        }
+        return { allowed: true, requestId, warnings, notices };
     },
 });
 export const finishRequest = mutation({
@@ -362,6 +547,14 @@ export const finishRequest = mutation({
         promptTokens: v.optional(v.number()),
         completionTokens: v.optional(v.number()),
         cachedTokens: v.optional(v.number()),
+        // Provider server-tool invocations that bill a per-call fee (e.g.
+        // { web_search: 3 }). Added to the token cost when no authoritative cost is
+        // supplied; recorded either way.
+        serverToolUses: v.optional(v.record(v.string(), v.number())),
+        // Authoritative cost from the gateway/provider, if reported. When present
+        // it's recorded verbatim (already includes any tool fees); when absent we
+        // price from tokens (cache-aware) plus server-tool fees.
+        costNanos: v.optional(v.number()),
         latencyMs: v.optional(v.number()),
     },
     returns: v.object({ costNanos: v.number() }),
@@ -384,7 +577,18 @@ export const finishRequest = mutation({
         const promptTokens = Math.max(0, args.promptTokens ?? 0);
         const completionTokens = Math.max(0, args.completionTokens ?? 0);
         const cachedTokens = Math.min(promptTokens, Math.max(0, args.cachedTokens ?? 0));
-        const costNanos = Math.max(0, costOf(promptTokens, completionTokens, await getPrice(ctx, request.model)));
+        // Prefer an authoritative gateway cost when supplied (it already includes
+        // tool fees); otherwise price from tokens — discounting the cached
+        // (prompt-cache-read) slice — plus any server-tool per-call fees.
+        let costNanos;
+        if (args.costNanos !== undefined && args.costNanos >= 0) {
+            costNanos = Math.round(args.costNanos);
+        }
+        else {
+            const settings = await getSettings(ctx);
+            costNanos =
+                settleCost(promptTokens, cachedTokens, completionTokens, await getPrice(ctx, request.model)) + serverToolCost(args.serverToolUses, settings?.serverToolPrices);
+        }
         // Durable write to the request's OWN row only — uncontended, so it always
         // lands. `settled: false` hands it to the fold step; the row is never left
         // orphaned in "pending" even if the totals update below fails and retries.
@@ -395,11 +599,12 @@ export const finishRequest = mutation({
             promptTokens,
             completionTokens,
             ...(cachedTokens > 0 ? { cachedTokens } : {}),
+            ...(args.serverToolUses ? { serverToolUses: args.serverToolUses } : {}),
             costNanos,
             latencyMs: args.latencyMs,
             settled: false,
         });
-        // Fold into the (hot) user/action counters in a separate mutation. If it
+        // Fold into the (hot) per-bucket counters in a separate mutation. If it
         // exhausts retries under contention, the cron reconciler picks it up.
         await ctx.scheduler.runAfter(0, internal.lib.foldTotals, {
             requestId: args.requestId,
@@ -407,9 +612,9 @@ export const finishRequest = mutation({
         return { costNanos };
     },
 });
-// Fold one finished request into the user/action running totals, releasing its
-// reservation. Idempotent: guarded by `settled` so the scheduler and the cron
-// reconciler can never double-count.
+// Fold one finished request into every attributed bucket's running totals,
+// releasing its reservation. Idempotent: guarded by `settled` so the scheduler
+// and the cron reconciler can never double-count.
 async function foldOne(ctx, req) {
     if (!req || req.settled !== false)
         return;
@@ -418,46 +623,42 @@ async function foldOne(ctx, req) {
     const tokens = (req.promptTokens ?? 0) + (req.completionTokens ?? 0);
     const estTokens = req.estimatedTokens ?? 0;
     const today = dayStamp();
-    const user = await getOrCreateUser(ctx, req.userId);
-    const uSameDay = user.dayStamp === today;
-    await ctx.db.patch(user._id, {
-        totalSpendNanos: user.totalSpendNanos + actual,
-        totalRequests: user.totalRequests + 1,
-        totalTokens: user.totalTokens + tokens,
-        dayStamp: today,
-        spendTodayNanos: (uSameDay ? user.spendTodayNanos : 0) + actual,
-        tokensToday: (uSameDay ? user.tokensToday ?? 0 : 0) + tokens,
-        reservedTodayNanos: Math.max(0, (uSameDay ? user.reservedTodayNanos ?? 0 : 0) - estCost),
-        reservedTotalNanos: Math.max(0, (user.reservedTotalNanos ?? 0) - estCost),
-        reservedTodayTokens: Math.max(0, (uSameDay ? user.reservedTodayTokens ?? 0 : 0) - estTokens),
-        reservedTotalTokens: Math.max(0, (user.reservedTotalTokens ?? 0) - estTokens),
-        pendingCount: Math.max(0, (user.pendingCount ?? 0) - 1),
-    });
-    if (req.actionName !== undefined) {
-        const action = await getOrCreateAction(ctx, req.actionName);
-        const aSameDay = action.dayStamp === today;
-        await ctx.db.patch(action._id, {
-            totalSpendNanos: action.totalSpendNanos + actual,
-            totalRequests: action.totalRequests + 1,
-            totalTokens: action.totalTokens + tokens,
+    const month = monthStamp();
+    // Accrue into every attributed bucket (user, action, and each tag) — capped
+    // or not. Buckets that never held a reservation have their reserved fields
+    // clamped at 0 by Math.max, so subtracting an estimate is a harmless no-op.
+    // Also write the durable per-(bucket, day/month) usage rows that survive
+    // request retention, so spend history outlives the raw request log.
+    for (const t of requestBuckets(req.userId, req.actionName, req.tags)) {
+        const b = await getOrCreateBucket(ctx, t.dimension, t.value);
+        const sameDay = b.dayStamp === today;
+        const sameMonth = b.monthStamp === month;
+        await ctx.db.patch(b._id, {
+            totalSpendNanos: b.totalSpendNanos + actual,
+            totalRequests: b.totalRequests + 1,
+            totalTokens: b.totalTokens + tokens,
             dayStamp: today,
-            spendTodayNanos: (aSameDay ? action.spendTodayNanos : 0) + actual,
-            tokensToday: (aSameDay ? action.tokensToday ?? 0 : 0) + tokens,
-            reservedTodayNanos: Math.max(0, (aSameDay ? action.reservedTodayNanos ?? 0 : 0) - estCost),
-            reservedTotalNanos: Math.max(0, (action.reservedTotalNanos ?? 0) - estCost),
-            reservedTodayTokens: Math.max(0, (aSameDay ? action.reservedTodayTokens ?? 0 : 0) - estTokens),
-            reservedTotalTokens: Math.max(0, (action.reservedTotalTokens ?? 0) - estTokens),
-            pendingCount: Math.max(0, (action.pendingCount ?? 0) - 1),
+            monthStamp: month,
+            spendTodayNanos: (sameDay ? b.spendTodayNanos : 0) + actual,
+            tokensToday: (sameDay ? b.tokensToday ?? 0 : 0) + tokens,
+            spendThisMonthNanos: (sameMonth ? b.spendThisMonthNanos ?? 0 : 0) + actual,
+            tokensThisMonth: (sameMonth ? b.tokensThisMonth ?? 0 : 0) + tokens,
+            reservedTodayNanos: Math.max(0, (sameDay ? b.reservedTodayNanos ?? 0 : 0) - estCost),
+            reservedMonthNanos: Math.max(0, (sameMonth ? b.reservedMonthNanos ?? 0 : 0) - estCost),
+            reservedTotalNanos: Math.max(0, (b.reservedTotalNanos ?? 0) - estCost),
+            reservedTodayTokens: Math.max(0, (sameDay ? b.reservedTodayTokens ?? 0 : 0) - estTokens),
+            reservedMonthTokens: Math.max(0, (sameMonth ? b.reservedMonthTokens ?? 0 : 0) - estTokens),
+            reservedTotalTokens: Math.max(0, (b.reservedTotalTokens ?? 0) - estTokens),
+            pendingCount: Math.max(0, (b.pendingCount ?? 0) - 1),
         });
+        await addUsage(ctx, t.dimension, t.value, "day", today, actual, tokens, 1);
+        await addUsage(ctx, t.dimension, t.value, "month", month, actual, tokens, 1);
     }
     // Deployment-wide totals via the sharded counter (only when a global cap is
     // configured — otherwise skip the writes entirely). Distributed across shards,
     // so this does not serialize on a single row.
     if (actual > 0) {
-        const settings = await ctx.db
-            .query("settings")
-            .withIndex("key", (q) => q.eq("key", "singleton"))
-            .unique();
+        const settings = await getSettings(ctx);
         if (settings &&
             (settings.globalDailySpendLimitNanos !== undefined ||
                 settings.globalLifetimeSpendLimitNanos !== undefined)) {
@@ -520,6 +721,7 @@ export const reconcile = internalMutation({
                 // Only rows that are done and accounted: folded (settled === true) or a
                 // blocked attempt (never needs folding). Never a pending/unfolded row.
                 if (req.settled === true || req.status === "blocked") {
+                    await deleteRequestTags(ctx, req._id);
                     await ctx.db.delete(req._id);
                     purged++;
                 }
@@ -565,10 +767,19 @@ export const getRequest = query({
     handler: async (ctx, args) => ctx.db.get(args.requestId),
 });
 export const listRequests = query({
-    args: { userId: v.optional(v.string()), limit: v.optional(v.number()) },
+    args: {
+        userId: v.optional(v.string()),
+        // Filter by any attribution dimension (user/action indexed on the request;
+        // custom tag dimensions resolved via the requestTags reverse index).
+        dimension: v.optional(v.string()),
+        value: v.optional(v.string()),
+        limit: v.optional(v.number()),
+    },
     handler: async (ctx, args) => {
         const limit = args.limit ?? 50;
-        if (args.userId !== undefined) {
+        const dim = args.dimension;
+        const val = args.value ?? args.userId;
+        if (dim === undefined && args.userId !== undefined) {
             const userId = args.userId;
             return await ctx.db
                 .query("requests")
@@ -576,155 +787,234 @@ export const listRequests = query({
                 .order("desc")
                 .take(limit);
         }
+        if (dim !== undefined && val !== undefined) {
+            if (dim === USER_DIM) {
+                return await ctx.db
+                    .query("requests")
+                    .withIndex("userId", (q) => q.eq("userId", val))
+                    .order("desc")
+                    .take(limit);
+            }
+            if (dim === ACTION_DIM) {
+                return await ctx.db
+                    .query("requests")
+                    .withIndex("actionName", (q) => q.eq("actionName", val))
+                    .order("desc")
+                    .take(limit);
+            }
+            // Custom tag dimension: walk the reverse index, then fetch each request.
+            const tagRows = await ctx.db
+                .query("requestTags")
+                .withIndex("dim_value", (q) => q.eq("dimension", dim).eq("value", val))
+                .order("desc")
+                .take(limit);
+            const rows = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
+            return rows.filter((r) => r !== null);
+        }
         return await ctx.db.query("requests").order("desc").take(limit);
     },
 });
 const ADMIN_LIST_CAP = 2000;
-export const listUsers = query({
-    args: {},
-    handler: async (ctx) => {
+// List budget buckets, optionally filtered to one dimension ("user", "action",
+// or any custom tag dimension). Today's spend is zeroed for stale day windows.
+export const listBuckets = query({
+    args: { dimension: v.optional(v.string()) },
+    handler: async (ctx, args) => {
         // Bounded to avoid an unbounded full-table scan on this reactive query.
-        // Paginate (ctx.db.query("users").paginate(...)) for larger deployments.
-        const users = await ctx.db.query("users").take(ADMIN_LIST_CAP);
+        // Paginate (ctx.db.query("buckets").paginate(...)) for larger deployments.
+        const rows = args.dimension !== undefined
+            ? await ctx.db
+                .query("buckets")
+                .withIndex("dimension", (q) => q.eq("dimension", args.dimension))
+                .take(ADMIN_LIST_CAP)
+            : await ctx.db.query("buckets").take(ADMIN_LIST_CAP);
         const today = dayStamp();
-        return users.map((u) => ({
-            ...u,
-            spendTodayNanos: u.dayStamp === today ? u.spendTodayNanos : 0,
+        const month = monthStamp();
+        return rows.map((b) => ({
+            ...b,
+            spendTodayNanos: b.dayStamp === today ? b.spendTodayNanos : 0,
+            spendThisMonthNanos: b.monthStamp === month ? b.spendThisMonthNanos ?? 0 : 0,
         }));
     },
 });
-export const setLimits = mutation({
+export const getBucket = query({
+    args: { dimension: v.string(), value: v.string() },
+    handler: async (ctx, args) => {
+        const b = await getBucketDoc(ctx, args.dimension, args.value);
+        if (!b)
+            return null;
+        const today = dayStamp();
+        const month = monthStamp();
+        return {
+            ...b,
+            spendTodayNanos: b.dayStamp === today ? b.spendTodayNanos : 0,
+            spendThisMonthNanos: b.monthStamp === month ? b.spendThisMonthNanos ?? 0 : 0,
+        };
+    },
+});
+// Set a bucket's limits/controls. `user` and `action` are just dimensions here;
+// the client's ai.users / ai.actions namespaces are thin wrappers over this.
+export const setBucketLimits = mutation({
     args: {
-        userId: v.string(),
+        dimension: v.string(),
+        value: v.string(),
         requestsPerMinute: v.optional(v.number()),
+        maxConcurrent: v.optional(v.number()),
         dailySpendLimitNanos: v.optional(v.number()),
+        monthlySpendLimitNanos: v.optional(v.number()),
         lifetimeSpendLimitNanos: v.optional(v.number()),
         dailyTokenLimit: v.optional(v.number()),
+        monthlyTokenLimit: v.optional(v.number()),
         lifetimeTokenLimit: v.optional(v.number()),
+        warnAtPct: v.optional(v.number()),
         enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
         blocked: v.optional(v.boolean()),
     },
     returns: v.null(),
     handler: async (ctx, args) => {
-        const user = await getOrCreateUser(ctx, args.userId);
-        const { userId: _userId, ...limits } = args;
-        await ctx.db.patch(user._id, limits);
+        const bucket = await getOrCreateBucket(ctx, args.dimension, args.value);
+        const { dimension: _d, value: _v, ...limits } = args;
+        await ctx.db.patch(bucket._id, limits);
         return null;
     },
 });
-// Delete a user and all their request rows (e.g. account deletion / GDPR).
-// Deletes in bounded batches and self-reschedules so it never exceeds the
-// per-transaction document limit — a user with millions of rows still deletes.
-const DELETE_BATCH = 500;
-export const deleteUser = mutation({
-    args: { userId: v.string() },
-    returns: v.object({ deletedThisBatch: v.number(), done: v.boolean() }),
-    handler: async (ctx, { userId }) => {
-        const rows = await ctx.db
-            .query("requests")
-            .withIndex("userId", (q) => q.eq("userId", userId))
-            .take(DELETE_BATCH);
-        for (const r of rows)
-            await ctx.db.delete(r._id);
-        if (rows.length === DELETE_BATCH) {
-            // More to go — continue in a fresh transaction.
-            await ctx.scheduler.runAfter(0, api.lib.deleteUser, { userId });
-            return { deletedThisBatch: rows.length, done: false };
-        }
-        // Last batch: remove the user row itself.
-        const user = await ctx.db
-            .query("users")
-            .withIndex("userId", (q) => q.eq("userId", userId))
-            .unique();
-        if (user)
-            await ctx.db.delete(user._id);
-        return { deletedThisBatch: rows.length + (user ? 1 : 0), done: true };
-    },
-});
-export const listActions = query({
-    args: {},
-    handler: async (ctx) => {
-        const actions = await ctx.db.query("actions").take(ADMIN_LIST_CAP);
-        const today = dayStamp();
-        return actions.map((a) => ({
-            ...a,
-            spendTodayNanos: a.dayStamp === today ? a.spendTodayNanos : 0,
-        }));
-    },
-});
 const vBumpArgs = {
     dailyNanos: v.optional(v.number()),
+    monthlyNanos: v.optional(v.number()),
     lifetimeNanos: v.optional(v.number()),
 };
-// One-time "approve another $X" bumps, added on top of the standing cap without
-// changing it. Daily bumps apply to today only; lifetime bumps are permanent.
-export const bumpUser = mutation({
-    args: { userId: v.string(), ...vBumpArgs },
+// One-time "approve another $X" bumps, added on top of a bucket's standing cap
+// without changing it. Daily/monthly bumps apply to the current window only;
+// lifetime bumps persist.
+export const bumpBucket = mutation({
+    args: { dimension: v.string(), value: v.string(), ...vBumpArgs },
     returns: v.null(),
-    handler: async (ctx, { userId, dailyNanos, lifetimeNanos }) => {
-        const user = await getOrCreateUser(ctx, userId);
+    handler: async (ctx, { dimension, value, dailyNanos, monthlyNanos, lifetimeNanos }) => {
+        const bucket = await getOrCreateBucket(ctx, dimension, value);
         const today = dayStamp();
-        const curDaily = user.bumpDayStamp === today ? user.dailyBumpNanos ?? 0 : 0;
-        await ctx.db.patch(user._id, {
+        const month = monthStamp();
+        const curDaily = bucket.bumpDayStamp === today ? bucket.dailyBumpNanos ?? 0 : 0;
+        const curMonthly = bucket.bumpMonthStamp === month ? bucket.monthlyBumpNanos ?? 0 : 0;
+        await ctx.db.patch(bucket._id, {
             bumpDayStamp: today,
             dailyBumpNanos: curDaily + (dailyNanos ?? 0),
-            lifetimeBumpNanos: (user.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
+            bumpMonthStamp: month,
+            monthlyBumpNanos: curMonthly + (monthlyNanos ?? 0),
+            lifetimeBumpNanos: (bucket.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
         });
         return null;
     },
 });
-export const bumpAction = mutation({
-    args: { name: v.string(), ...vBumpArgs },
+// Manually credit or debit a bucket (comp a user, correct an overcharge).
+// Negative deltaNanos = credit/refund, positive = extra charge. Adjusts the
+// live day/month/lifetime windows AND the durable usage history, and records an
+// audit row. Does not touch the global sharded total or reservations.
+export const adjustBucket = mutation({
+    args: {
+        dimension: v.string(),
+        value: v.string(),
+        deltaNanos: v.number(),
+        tokens: v.optional(v.number()),
+        reason: v.optional(v.string()),
+    },
     returns: v.null(),
-    handler: async (ctx, { name, dailyNanos, lifetimeNanos }) => {
-        const action = await getOrCreateAction(ctx, name);
+    handler: async (ctx, { dimension, value, deltaNanos, tokens, reason }) => {
+        const b = await getOrCreateBucket(ctx, dimension, value);
         const today = dayStamp();
-        const curDaily = action.bumpDayStamp === today ? action.dailyBumpNanos ?? 0 : 0;
-        await ctx.db.patch(action._id, {
-            bumpDayStamp: today,
-            dailyBumpNanos: curDaily + (dailyNanos ?? 0),
-            lifetimeBumpNanos: (action.lifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
+        const month = monthStamp();
+        const dSame = b.dayStamp === today;
+        const mSame = b.monthStamp === month;
+        const dt = tokens ?? 0;
+        await ctx.db.patch(b._id, {
+            totalSpendNanos: Math.max(0, b.totalSpendNanos + deltaNanos),
+            totalTokens: Math.max(0, b.totalTokens + dt),
+            dayStamp: today,
+            monthStamp: month,
+            spendTodayNanos: Math.max(0, (dSame ? b.spendTodayNanos : 0) + deltaNanos),
+            tokensToday: Math.max(0, (dSame ? b.tokensToday ?? 0 : 0) + dt),
+            spendThisMonthNanos: Math.max(0, (mSame ? b.spendThisMonthNanos ?? 0 : 0) + deltaNanos),
+            tokensThisMonth: Math.max(0, (mSame ? b.tokensThisMonth ?? 0 : 0) + dt),
         });
+        await ctx.db.insert("adjustments", { dimension, value, deltaNanos, tokens: dt, reason });
+        await addUsage(ctx, dimension, value, "day", today, deltaNanos, dt, 0);
+        await addUsage(ctx, dimension, value, "month", month, deltaNanos, dt, 0);
         return null;
     },
 });
-export const bumpGlobal = mutation({
-    args: vBumpArgs,
-    returns: v.null(),
-    handler: async (ctx, { dailyNanos, lifetimeNanos }) => {
-        const today = dayStamp();
-        const s = await getSettings(ctx);
-        const curDaily = s?.globalBumpDayStamp === today ? s?.globalDailyBumpNanos ?? 0 : 0;
-        const patch = {
-            globalBumpDayStamp: today,
-            globalDailyBumpNanos: curDaily + (dailyNanos ?? 0),
-            globalLifetimeBumpNanos: (s?.globalLifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
-        };
-        if (s)
-            await ctx.db.patch(s._id, patch);
-        else
-            await ctx.db.insert("settings", { key: "singleton", ...patch });
-        return null;
-    },
+export const listAdjustments = query({
+    args: { dimension: v.string(), value: v.string(), limit: v.optional(v.number()) },
+    handler: async (ctx, { dimension, value, limit }) => ctx.db
+        .query("adjustments")
+        .withIndex("dim_value", (q) => q.eq("dimension", dimension).eq("value", value))
+        .order("desc")
+        .take(limit ?? 50),
 });
-export const setActionLimits = mutation({
+// Durable spend history for a bucket: per-day or per-month rows, newest first.
+// Survives request retention.
+export const usageHistory = query({
     args: {
-        name: v.string(),
-        dailySpendLimitNanos: v.optional(v.number()),
-        lifetimeSpendLimitNanos: v.optional(v.number()),
-        dailyTokenLimit: v.optional(v.number()),
-        lifetimeTokenLimit: v.optional(v.number()),
-        enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
-        disabled: v.optional(v.boolean()),
+        dimension: v.string(),
+        value: v.string(),
+        period: v.union(v.literal("day"), v.literal("month")),
+        limit: v.optional(v.number()),
     },
+    handler: async (ctx, { dimension, value, period, limit }) => ctx.db
+        .query("usage")
+        .withIndex("bucket_period_stamp", (q) => q.eq("dimension", dimension).eq("value", value).eq("period", period))
+        .order("desc")
+        .take(limit ?? 90),
+});
+// Deployment-wide default threshold for approaching-limit alerts (fraction of a
+// cap, e.g. 0.8). Buckets can override with their own warnAtPct.
+export const setAlertDefaults = mutation({
+    args: { warnAtPct: v.optional(v.number()) },
     returns: v.null(),
-    handler: async (ctx, args) => {
-        const action = await getOrCreateAction(ctx, args.name);
-        const { name: _name, ...limits } = args;
-        await ctx.db.patch(action._id, limits);
+    handler: async (ctx, { warnAtPct }) => {
+        const existing = await getSettings(ctx);
+        if (existing)
+            await ctx.db.patch(existing._id, { defaultWarnAtPct: warnAtPct });
+        else
+            await ctx.db.insert("settings", { key: "singleton", defaultWarnAtPct: warnAtPct });
         return null;
     },
 });
+// Delete a bucket and (for the `user` dimension) all of that user's request
+// rows — e.g. account deletion / GDPR. Deletes requests in bounded batches and
+// self-reschedules so it never exceeds the per-transaction document limit.
+const DELETE_BATCH = 500;
+export const deleteBucket = mutation({
+    args: { dimension: v.string(), value: v.string() },
+    returns: v.object({ deletedThisBatch: v.number(), done: v.boolean() }),
+    handler: async (ctx, { dimension, value }) => {
+        // Only the user dimension owns request rows (indexed by userId). Other
+        // dimensions just drop their budget-holder row.
+        if (dimension === USER_DIM) {
+            const rows = await ctx.db
+                .query("requests")
+                .withIndex("userId", (q) => q.eq("userId", value))
+                .take(DELETE_BATCH);
+            for (const r of rows) {
+                await deleteRequestTags(ctx, r._id);
+                await ctx.db.delete(r._id);
+            }
+            if (rows.length === DELETE_BATCH) {
+                await ctx.scheduler.runAfter(0, api.lib.deleteBucket, {
+                    dimension,
+                    value,
+                });
+                return { deletedThisBatch: rows.length, done: false };
+            }
+            const bucket = await getBucketDoc(ctx, dimension, value);
+            if (bucket)
+                await ctx.db.delete(bucket._id);
+            return { deletedThisBatch: rows.length + (bucket ? 1 : 0), done: true };
+        }
+        const bucket = await getBucketDoc(ctx, dimension, value);
+        if (bucket)
+            await ctx.db.delete(bucket._id);
+        return { deletedThisBatch: bucket ? 1 : 0, done: true };
+    },
+});
 export const getModelPolicy = query({
     args: {},
     returns: v.object({
@@ -747,6 +1037,9 @@ export const getGlobalStatus = query({
         enforcement: v.union(v.literal("hard"), v.literal("soft")),
         spentTodayNanos: v.number(),
         spentTotalNanos: v.number(),
+        // deployment-wide config (surfaced for the admin dashboard)
+        retentionMs: v.union(v.number(), v.null()),
+        defaultWarnAtPct: v.union(v.number(), v.null()),
     }),
     handler: async (ctx) => {
         const s = await ctx.db
@@ -759,6 +1052,8 @@ export const getGlobalStatus = query({
             enforcement: s?.globalEnforcement ?? "hard",
             spentTodayNanos: await globalSpend.count(ctx, globalDayKey(dayStamp())),
             spentTotalNanos: await globalSpend.count(ctx, GLOBAL_TOTAL),
+            retentionMs: s?.retentionMs ?? null,
+            defaultWarnAtPct: s?.defaultWarnAtPct ?? null,
         };
     },
 });
@@ -786,6 +1081,25 @@ export const setGlobalLimits = mutation({
         return null;
     },
 });
+export const bumpGlobal = mutation({
+    args: vBumpArgs,
+    returns: v.null(),
+    handler: async (ctx, { dailyNanos, lifetimeNanos }) => {
+        const today = dayStamp();
+        const s = await getSettings(ctx);
+        const curDaily = s?.globalBumpDayStamp === today ? s?.globalDailyBumpNanos ?? 0 : 0;
+        const patch = {
+            globalBumpDayStamp: today,
+            globalDailyBumpNanos: curDaily + (dailyNanos ?? 0),
+            globalLifetimeBumpNanos: (s?.globalLifetimeBumpNanos ?? 0) + (lifetimeNanos ?? 0),
+        };
+        if (s)
+            await ctx.db.patch(s._id, patch);
+        else
+            await ctx.db.insert("settings", { key: "singleton", ...patch });
+        return null;
+    },
+});
 export const setModelPolicy = mutation({
     args: {
         mode: v.union(v.literal("open"), v.literal("allowlist"), v.literal("denylist")),
@@ -815,12 +1129,16 @@ export const setPrice = mutation({
         model: v.string(),
         inputNanosPerMTok: v.number(),
         outputNanosPerMTok: v.number(),
+        // optional cache-read rate; if omitted, a default discount off input applies
+        cachedNanosPerMTok: v.optional(v.number()),
     },
     returns: v.null(),
     handler: async (ctx, args) => {
         // Negative prices would make costOf return a negative cost, which folds
         // into totals as a spend *refund* — pushing a user back under their cap.
-        if (args.inputNanosPerMTok < 0 || args.outputNanosPerMTok < 0) {
+        if (args.inputNanosPerMTok < 0 ||
+            args.outputNanosPerMTok < 0 ||
+            (args.cachedNanosPerMTok ?? 0) < 0) {
             throw new Error("Prices must be non-negative");
         }
         const existing = await ctx.db
@@ -848,9 +1166,34 @@ export const listPrices = query({
             merged[o.model] = {
                 input: o.inputNanosPerMTok,
                 output: o.outputNanosPerMTok,
+                cached: o.cachedNanosPerMTok,
                 overridden: true,
             };
         }
         return merged;
     },
 });
+// Per-call fees for provider server tools (web search, etc.), defaults merged
+// with any deployment overrides.
+export const listServerToolPrices = query({
+    args: {},
+    handler: async (ctx) => {
+        const s = await getSettings(ctx);
+        return { ...DEFAULT_SERVER_TOOL_PRICES, ...(s?.serverToolPrices ?? {}) };
+    },
+});
+export const setServerToolPrice = mutation({
+    args: { tool: v.string(), nanosPerCall: v.number() },
+    returns: v.null(),
+    handler: async (ctx, { tool, nanosPerCall }) => {
+        if (nanosPerCall < 0)
+            throw new Error("Prices must be non-negative");
+        const s = await getSettings(ctx);
+        const serverToolPrices = { ...(s?.serverToolPrices ?? {}), [tool]: nanosPerCall };
+        if (s)
+            await ctx.db.patch(s._id, { serverToolPrices });
+        else
+            await ctx.db.insert("settings", { key: "singleton", serverToolPrices });
+        return null;
+    },
+});