@convex-dev/ai-budget 0.0.2-alpha.16 → 0.0.2-alpha.18

package/README.md

@@ -390,6 +390,13 @@ reserve-then-settle admission check runs per bucket. Uncapped buckets never seri
 adding tags you don't cap is free at admission; their totals still accrue for
 reporting.
 
+> **Choosing what to cap.** A cap makes its bucket's admissions atomic by
+> reserving on that bucket's single row, so capping a dimension that *all* traffic
+> shares (one `env`/`action` every request carries) serializes those admissions on
+> one document under load. Prefer capping **naturally-sharded** dimensions — per
+> `user`, per `customer` — and use the [global cap](#global-cap) for a
+> deployment-wide ceiling. Uncapped high-traffic tags are always free.
+
 ### Tags — budgeting by any dimension
 
 ```ts
@@ -488,9 +495,12 @@ ai.global.status(ctx)   // { limits, spentTodayNanos, spentTotalNanos, … }
 ai.global.bump(ctx, { dailyNanos?, lifetimeNanos? })
 ```
 
-A killswitch across everything. Backed by a sharded counter for throughput, so
-it's enforced **approximately** (bounded concurrency overshoot under burst).
-Per-bucket admission is atomic against each request's estimate.
+A best-effort killswitch across everything. Backed by a sharded counter for
+throughput, so it's enforced **approximately** — under burst it can overshoot the
+cap by a bounded amount, and it excludes in-flight (not-yet-settled) spend. It is
+**not** a to-the-dollar ceiling: for an exact limit use a per-bucket cap (those
+reserve atomically); reach for the global cap when you want a deployment-wide
+"stop everything" switch.
 
 ### Model policy
 
@@ -552,6 +562,12 @@ ai.global.setRetention(ctx, { retentionMs })   // default 1h; 0 disables
 Full request rows (prompts + responses) are swept after the window to bound the
 audit table. **Spend history survives** — it lives in separate durable rollups.
 
+> **Prompt/response content is stored** on the request row until the window sweeps
+> it (default 1h), so it's visible in the request log and to admin reads until
+> then. If your prompts carry PII you don't want retained, shorten `retentionMs`
+> (or set it low enough that content lives only as long as you need replay). Spend
+> rollups never contain content, so trimming retention doesn't cost you history.
+
 ---
 
 ## Admin dashboard
@@ -568,19 +584,29 @@ import { AIBudget } from "@convex-dev/ai-budget";
 const ai = new AIBudget(components.aiBudget);
 const http = httpRouter();
 
-ai.registerRoutes(http, {
-  // Gate it — the endpoint is public. Recommended: check the caller is an admin.
-  authorize: async (ctx) => (await ctx.auth.getUserIdentity())?.role === "admin",
-});
+// Recommended: token mode. Set AI_BUDGET_DASHBOARD_TOKEN in the deployment env;
+// open the dashboard with ?token=<it> once and the page keeps using it.
+ai.registerRoutes(http);
 
 export default http;
 ```
 
 It lives at `https://<deployment>.convex.site/aibudget` (override with `path`).
-**It is a public internet endpoint, so you must gate it**: pass `authorize`
-(return `true` to allow) or set `AI_BUDGET_DASHBOARD_TOKEN` (sent as
-`Authorization: Bearer …`). With neither, every route returns 401. Everything the
-page shows is backed by the component's own functions — nothing else to wire up.
+**It is a public internet endpoint, so you must gate it**, one of two ways:
+
+- **Token (recommended, works end-to-end):** set the `AI_BUDGET_DASHBOARD_TOKEN`
+  env var. Open `…/aibudget?token=<token>` once; the page strips it from the URL
+  and sends it as a bearer on every API call. `?token=` is accepted only on the
+  page navigation, and the compare is constant-time.
+- **`authorize(ctx, request)` callback:** must authenticate from something the
+  **browser sends on a top-level navigation** — a cookie or a header *you*
+  control — because a page load carries no `Authorization` bearer. `ctx.auth`
+  (deployment JWT) is `null` for the HTML page, so `authorize: (ctx) => (await
+  ctx.auth.getUserIdentity())?.role === "admin"` will 401 the page. Use it only
+  when you have your own session cookie to check.
+
+With neither configured, every route returns 401. Everything the page shows is
+backed by the component's own functions — nothing else to wire up.
 
 ---
 
@@ -625,6 +651,19 @@ more tokens or costs more than estimated, settlement records the real amount and
 the final total can exceed a hard cap by that request's estimation delta. The next
 admission sees the settled total and blocks until there is headroom again.
 
+**Throughput characteristics (know these before you turn on a global cap).**
+Reconciliation runs as small, independently-rescheduling phases, so folding,
+expiry, and retention can't stall each other and each drains its own backlog.
+Two shared-write hot spots remain, by design:
+- The **global cap** reads a sharded counter *inside* admission, which contends
+  with every fold that writes it — so a configured global cap adds contention on
+  the hot path. It's a deployment-wide killswitch, not a high-throughput per-call
+  limit; prefer per-bucket caps for the common case.
+- **Settlement writes every attributed bucket**, including the per-request
+  `action` bucket that every call shares. That single row is written by every
+  settle, so an extremely high single-action settle rate can lag totals. Spread
+  load across naturally-sharded dimensions (per user/customer) where you can.
+
 The `error.md` file documents the adversarial audits this design survived, with
 live repros.
 

package/dist/client/dashboard.js

@@ -97,8 +97,18 @@ async function renderBuckets() {
   const dims = ["user", "action"];
   const state = { dimension: window.__dim ?? "" };
   const rows = await get("/buckets", state.dimension ? { dimension: state.dimension } : {});
-  const grand = rows.reduce((s, b) => s + b.totalSpendNanos, 0);
-  document.getElementById("total").textContent = rows.length + " buckets · " + usd(grand) + " total";
+  // Within ONE dimension each request bills exactly one bucket, so summing is a
+  // true total. Across dimensions a request bills several buckets (user + action
+  // + tags), so summing multi-counts — use the deployment-wide spend instead.
+  let totalText;
+  if (state.dimension) {
+    const grand = rows.reduce((s, b) => s + b.totalSpendNanos, 0);
+    totalText = rows.length + " buckets · " + usd(grand) + " total (" + state.dimension + ")";
+  } else {
+    const g = await get("/global", {});
+    totalText = rows.length + " buckets · " + usd(g.spentTotalNanos ?? 0) + " spent (deployment)";
+  }
+  document.getElementById("total").textContent = totalText;
   const dimSet = [...new Set(rows.map((b) => b.dimension).concat(dims))];
   const sel = el("select", { value: state.dimension, style: "width:140px",
     onchange: (e) => { window.__dim = e.target.value; render(); } },

package/dist/client/index.js

@@ -39,8 +39,17 @@ function toTokenCount(x) {
 }
 function extractUsage(usage) {
     return {
-        promptTokens: toTokenCount(usage?.inputTokens ?? usage?.promptTokens),
-        completionTokens: toTokenCount(usage?.outputTokens ?? usage?.completionTokens),
+        // Cover AI SDK camelCase (v5/v7), raw OpenAI-compatible snake_case, AND raw
+        // Anthropic (`input_tokens`/`output_tokens`), so a `meter` caller passing any
+        // provider's raw `usage` object still gets counts.
+        promptTokens: toTokenCount(usage?.inputTokens ??
+            usage?.promptTokens ??
+            usage?.prompt_tokens ??
+            usage?.input_tokens),
+        completionTokens: toTokenCount(usage?.outputTokens ??
+            usage?.completionTokens ??
+            usage?.completion_tokens ??
+            usage?.output_tokens),
         // cached prompt tokens. The Convex gateway reports these at
         // `usage.inputTokenDetails.cacheReadTokens`; the other paths cover AI SDK v5
         // (`cachedInputTokens`) and raw OpenAI-compatible shapes.
@@ -48,7 +57,8 @@ function extractUsage(usage) {
             usage?.cachedInputTokens ??
             usage?.promptTokensDetails?.cachedTokens ??
             usage?.prompt_tokens_details?.cached_tokens ??
-            usage?.cached_tokens),
+            usage?.cached_tokens ??
+            usage?.cache_read_input_tokens),
     };
 }
 // The AI Gateway reports the authoritative dollar cost of each request.
@@ -70,6 +80,10 @@ function extractGatewayCostNanos(result) {
     return undefined;
 }
 const NANOS_PER_DOLLAR = 1e9;
+// Cap stored prompt/response content so one row can't approach Convex's 1 MiB
+// document limit (which would fail the call) or bloat the reconciler's scans.
+const MAX_STORED_CONTENT = 32 * 1024;
+const capContent = (s) => s.length > MAX_STORED_CONTENT ? s.slice(0, MAX_STORED_CONTENT) + "…[truncated]" : s;
 // Flatten an AI SDK prompt (roles + content parts) into simple storable messages.
 function simplifyPrompt(prompt) {
     if (!Array.isArray(prompt))
@@ -81,13 +95,22 @@ function simplifyPrompt(prompt) {
         }
         else if (Array.isArray(m.content)) {
             content = m.content
-                .map((part) => part?.type === "text" ? part.text : JSON.stringify(part))
+                .map((part) => {
+                if (part?.type === "text")
+                    return part.text ?? "";
+                // NEVER inline a base64 image/file part — a data URL or Uint8Array here
+                // becomes megabytes, pushing the stored row past the 1 MiB doc limit
+                // (failing the call) and turning a 200 KB image into a ~50k-token
+                // estimate. Store a compact placeholder.
+                const bytes = part?.data?.length ?? part?.image?.length ?? part?.data?.byteLength;
+                return `[${part?.type ?? "part"}${typeof bytes === "number" ? ` ${bytes}b` : ""}]`;
+            })
                 .join("");
         }
         else {
             content = JSON.stringify(m.content);
         }
-        return { role: String(m.role), content };
+        return { role: String(m.role), content: capContent(content) };
     });
 }
 function extractText(result) {
@@ -240,35 +263,42 @@ export class AIBudget {
         }
         const { requestId, warnings, notices } = started;
         const start = Date.now();
+        // Run the provider call. ONLY a failure of the call itself settles as an
+        // error (no charge expected).
+        let out;
         try {
-            const out = await run();
-            const { costNanos } = await this.settle(ctx, {
-                requestId,
-                responseText: out.text,
-                usage: out.usage,
-                promptTokens: out.promptTokens,
-                completionTokens: out.completionTokens,
-                cachedTokens: out.cachedTokens,
-                serverToolUses: out.serverToolUses,
-                costNanos: out.costNanos,
-                latencyMs: Date.now() - start,
-            });
-            // Re-derive the recorded usage for the return value.
-            const usage = out.promptTokens !== undefined ||
-                out.completionTokens !== undefined ||
-                out.cachedTokens !== undefined
-                ? {
-                    promptTokens: out.promptTokens ?? 0,
-                    completionTokens: out.completionTokens ?? 0,
-                    cachedTokens: out.cachedTokens ?? 0,
-                }
-                : extractUsage(out.usage);
-            return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
+            out = await run();
         }
         catch (e) {
             await this.settle(ctx, { requestId, error: String(e), latencyMs: Date.now() - start });
             throw e;
         }
+        // The call SUCCEEDED (the provider may have charged). Settle the real usage.
+        // If settlement itself fails here, do NOT fall into an error-settle that
+        // records zero — that would erase a real charge. Rethrow and leave the
+        // reservation for the reconciler; billing stays "unknown", never a false zero.
+        const { costNanos } = await this.settle(ctx, {
+            requestId,
+            responseText: out.text,
+            usage: out.usage,
+            promptTokens: out.promptTokens,
+            completionTokens: out.completionTokens,
+            cachedTokens: out.cachedTokens,
+            serverToolUses: out.serverToolUses,
+            costNanos: out.costNanos,
+            latencyMs: Date.now() - start,
+        });
+        // Re-derive the recorded usage for the return value.
+        const usage = out.promptTokens !== undefined ||
+            out.completionTokens !== undefined ||
+            out.cachedTokens !== undefined
+            ? {
+                promptTokens: out.promptTokens ?? 0,
+                completionTokens: out.completionTokens ?? 0,
+                cachedTokens: out.cachedTokens ?? 0,
+            }
+            : extractUsage(out.usage);
+        return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
     }
     /**
      * One-shot chat through the AI Gateway with tracking + limits — sugar over
@@ -415,15 +445,10 @@ export class AIBudget {
                 wrapGenerate: async ({ doGenerate, params }) => {
                     const requestId = await begin(params);
                     const start = Date.now();
+                    // Only a failure of the generation itself settles as an error.
+                    let result;
                     try {
-                        const result = await doGenerate();
-                        await finish(requestId, {
-                            responseText: extractText(result),
-                            ...extractUsage(result.usage),
-                            costNanos: extractGatewayCostNanos(result),
-                            latencyMs: Date.now() - start,
-                        });
-                        return result;
+                        result = await doGenerate();
                     }
                     catch (e) {
                         await finish(requestId, {
@@ -432,6 +457,15 @@ export class AIBudget {
                         });
                         throw e;
                     }
+                    // Generation succeeded (provider may have charged). Settle the real
+                    // usage; a failure here rethrows rather than recording a false zero.
+                    await finish(requestId, {
+                        responseText: extractText(result),
+                        ...extractUsage(result.usage),
+                        costNanos: extractGatewayCostNanos(result),
+                        latencyMs: Date.now() - start,
+                    });
+                    return result;
                 },
                 wrapStream: async ({ doStream, params }) => {
                     const requestId = await begin(params);
@@ -446,21 +480,22 @@ export class AIBudget {
                         // chunk, or a cancel — is safe: the first wins, the rest no-op.
                         // Without this an errored or abandoned stream would never settle and
                         // its real usage would be lost (recorded as free by the reconciler).
-                        let settled = false;
-                        const settle = (error) => {
-                            if (settled)
-                                return;
-                            settled = true;
-                            return finish(requestId, {
-                                responseText: text,
-                                error,
-                                ...extractUsage(usage),
-                                costNanos: extractGatewayCostNanos({ providerMetadata }),
-                                latencyMs: Date.now() - start,
-                            });
-                        };
+                        // Settle at most once, memoizing the PROMISE so the finish chunk, an
+                        // error chunk, and flush all await the same settlement instead of
+                        // racing, dropping it (the old `void settle()`), or flipping a
+                        // "settled" flag before the mutation actually committed. A stream
+                        // that is cancelled/never fully consumed won't deliver finish or
+                        // flush; the reconciler's reservation expiry is the backstop there.
+                        let settlement;
+                        const settle = (error) => (settlement ??= finish(requestId, {
+                            responseText: text,
+                            error,
+                            ...extractUsage(usage),
+                            costNanos: extractGatewayCostNanos({ providerMetadata }),
+                            latencyMs: Date.now() - start,
+                        }));
                         const tapped = result.stream.pipeThrough(new TransformStream({
-                            transform(chunk, controller) {
+                            async transform(chunk, controller) {
                                 if (chunk?.type === "text-delta") {
                                     text += chunk.delta ?? chunk.textDelta ?? "";
                                 }
@@ -468,13 +503,30 @@ export class AIBudget {
                                     usage = chunk.usage;
                                     providerMetadata = chunk.providerMetadata ?? providerMetadata;
                                 }
-                                if (chunk?.type === "error")
-                                    void settle(String(chunk.error));
                                 controller.enqueue(chunk);
+                                // Settle after forwarding the terminal error chunk, and AWAIT
+                                // it so a failed settle surfaces instead of being dropped.
+                                if (chunk?.type === "error")
+                                    await settle(String(chunk.error));
                             },
                             async flush() {
                                 await settle();
                             },
+                            // A cancelled/aborted stream (client disconnect, AbortSignal, or
+                            // breaking out of `for await`) does NOT run `flush`. Without this
+                            // the request would sit pending until the 30-min reservation
+                            // expiry and be recorded as $0 though the provider charged for
+                            // what streamed. Settle here with the text so far, estimating the
+                            // completion tokens when the provider gave no usage.
+                            // `cancel` is a newer Streams-spec transformer hook not yet in
+                            // the TS DOM lib types; cast the literal so it compiles (runtimes
+                            // without it simply won't call it, and the reconciler backstops).
+                            async cancel(reason) {
+                                if (usage === undefined && text) {
+                                    usage = { outputTokens: Math.ceil(text.length / 4) };
+                                }
+                                await settle(`cancelled: ${String(reason)}`);
+                            },
                         }));
                         return { ...result, stream: tapped };
                     }
@@ -658,16 +710,26 @@ export class AIBudget {
                 return { ok: false, token: "" };
             const url = new URL(request.url);
             const bearer = (request.headers.get("authorization") ?? "").replace(/^Bearer\s+/i, "");
-            // `?token=` is accepted only for the initial page navigation (a browser
-            // GET can't set headers); the page strips it from the URL on load and the
-            // JSON API is called with the bearer header. Compared in constant time.
-            const provided = bearer || url.searchParams.get("token") || "";
+            // `?token=` is accepted ONLY for the initial page navigation (a browser GET
+            // can't set headers); the page strips it from the URL on load and calls the
+            // JSON API with the bearer header. Restrict it to that GET page route so a
+            // token can't be smuggled in a query string on API/mutation calls (where it
+            // would also land in access logs). Compared in constant time.
+            const sub = url.pathname.slice(prefix.length) || "/";
+            const isPageNav = request.method === "GET" && !sub.startsWith("/api");
+            const provided = bearer || (isPageNav ? url.searchParams.get("token") ?? "" : "");
             return { ok: timingSafeEqual(provided, token), token };
         };
         const json = (data, status = 200) => new Response(JSON.stringify(data ?? null), {
             status,
-            headers: { "content-type": "application/json" },
+            // Never let a shared cache/proxy retain budget data or the token-bearing
+            // page — these responses are per-viewer and sensitive.
+            headers: { "content-type": "application/json", "cache-control": "no-store" },
         });
+        // JSON.stringify does NOT escape `<`, so a value containing `</script>`
+        // would close the inline <script> and break out. Escape `<` (and the JS line
+        // separators) before embedding in HTML.
+        const jsonForScript = (x) => JSON.stringify(x).replace(/[<\u2028\u2029]/g, (ch) => "\\u" + ch.charCodeAt(0).toString(16).padStart(4, "0"));
         const handle = async (ctx, request) => {
             const url = new URL(request.url);
             const sub = url.pathname.slice(prefix.length) || "/";
@@ -724,12 +786,16 @@ export class AIBudget {
                         return json({ error: "not found" }, 404);
                 }
             }
-            // Inject as JSON literals (function replacers so `$` in the value isn't
-            // treated as a replacement pattern). This keeps a token/prefix containing
-            // quotes, backslashes, or `</script>` from breaking out of the JS string.
-            const html = DASHBOARD_HTML.replace(/__API_BASE__/g, () => JSON.stringify(`${prefix}/api`)).replace(/__TOKEN__/g, () => JSON.stringify(token));
+            // Inject as script-safe JSON literals (function replacers so `$` in the
+            // value isn't treated as a replacement pattern; `jsonForScript` escapes
+            // `<` so a token containing `</script>` can't break out of the inline JS).
+            const html = DASHBOARD_HTML.replace(/__API_BASE__/g, () => jsonForScript(`${prefix}/api`)).replace(/__TOKEN__/g, () => jsonForScript(token));
+            // The page embeds the bearer token — never let a shared cache retain it.
             return new Response(html, {
-                headers: { "content-type": "text/html; charset=utf-8" },
+                headers: {
+                    "content-type": "text/html; charset=utf-8",
+                    "cache-control": "no-store",
+                },
             });
         };
         const handler = httpActionGeneric(handle);

package/dist/component/_generated/component.d.ts

@@ -117,9 +117,9 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             warnAtPct?: number;
         }, null, Name>;
         setGlobalLimits: FunctionReference<"mutation", "internal", {
-            dailySpendLimitNanos?: number;
-            enforcement?: "hard" | "soft";
-            lifetimeSpendLimitNanos?: number;
+            dailySpendLimitNanos?: number | null;
+            enforcement?: "hard" | "soft" | null;
+            lifetimeSpendLimitNanos?: number | null;
         }, null, Name>;
         setModelPolicy: FunctionReference<"mutation", "internal", {
             mode: "open" | "allowlist" | "denylist";

package/dist/component/lib.d.ts

@@ -39,9 +39,14 @@ export declare const finishRequest: import("convex/server").RegisteredMutation<"
 export declare const foldTotals: import("convex/server").RegisteredMutation<"internal", {
     requestId: import("convex/values").GenericId<"requests">;
 }, Promise<null>>;
-export declare const reconcile: import("convex/server").RegisteredMutation<"internal", {}, Promise<{
+export declare const reconcile: import("convex/server").RegisteredMutation<"internal", {}, Promise<null>>;
+export declare const foldPhase: import("convex/server").RegisteredMutation<"internal", {}, Promise<{
     folded: number;
+}>>;
+export declare const expirePhase: import("convex/server").RegisteredMutation<"internal", {}, Promise<{
     expired: number;
+}>>;
+export declare const retentionPhase: import("convex/server").RegisteredMutation<"internal", {}, Promise<{
     purged: number;
 }>>;
 export declare const setRetention: import("convex/server").RegisteredMutation<"public", {
@@ -377,9 +382,9 @@ export declare const getGlobalStatus: import("convex/server").RegisteredQuery<"p
     defaultWarnAtPct: number | null;
 }>>;
 export declare const setGlobalLimits: import("convex/server").RegisteredMutation<"public", {
-    dailySpendLimitNanos?: number | undefined;
-    lifetimeSpendLimitNanos?: number | undefined;
-    enforcement?: "hard" | "soft" | undefined;
+    dailySpendLimitNanos?: number | null | undefined;
+    lifetimeSpendLimitNanos?: number | null | undefined;
+    enforcement?: "hard" | "soft" | null | undefined;
 }, Promise<null>>;
 export declare const bumpGlobal: import("convex/server").RegisteredMutation<"public", {
     dailyNanos?: number | undefined;