@convex-dev/ai-budget 0.0.2-alpha.15 → 0.0.2-alpha.17

package/dist/component/lib.js

@@ -12,6 +12,30 @@ 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)}`;
+// Convex's `v.number()` accepts NaN and ±Infinity. Those are poison here: a NaN
+// cap or count silently defeats every `used > cap` comparison (NaN > x is
+// false), so an unvalidated NaN would make admission fail OPEN and admit
+// unlimited spend; +Infinity in totals is just as corrupting. Validate every
+// externally-supplied accounting amount at the mutation boundary.
+function assertAmount(n, name, { signed = false } = {}) {
+    if (n === undefined)
+        return;
+    if (!Number.isFinite(n) || !Number.isSafeInteger(n)) {
+        throw new Error(`${name} must be a finite safe integer (got ${n})`);
+    }
+    if (!signed && n < 0)
+        throw new Error(`${name} must be nonnegative (got ${n})`);
+}
+function assertFraction(n, name) {
+    if (n === undefined)
+        return;
+    if (!Number.isFinite(n) || n < 0 || n > 1) {
+        throw new Error(`${name} must be a number in [0, 1] (got ${n})`);
+    }
+}
+// Coerce a caller/provider-supplied count to a finite nonnegative integer,
+// mapping NaN/Infinity/garbage to 0 rather than poisoning downstream totals.
+const safeCount = (n) => Number.isFinite(n) ? Math.max(0, Math.floor(n)) : 0;
 // 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 —
@@ -525,8 +549,14 @@ export const startRequest = mutation({
             notices.push(...globalEval.notices);
         }
         // Consume only after every admission check succeeds, in the same
-        // transaction as the reservations and request insert. An unexpected
-        // rejection throws to roll back all previously consumed dimensions.
+        // transaction as the reservations and request insert. `throws: true` is
+        // deliberate, not a rough edge: we already `.check`ed every bucket above in
+        // this same serializable transaction, so a `.limit` here cannot fail on a
+        // bucket that passed check — and if it somehow did (or a later bucket did),
+        // throwing rolls back the WHOLE transaction, including the rate we already
+        // consumed on earlier buckets. Converting this to a graceful `{allowed:false}`
+        // return would COMMIT the partial consumption and leak rate capacity, so keep
+        // the throw.
         for (const b of buckets) {
             if (b.requestsPerMinute !== undefined) {
                 await requestRateLimiter.limit(ctx, "requests", {
@@ -612,23 +642,28 @@ export const finishRequest = mutation({
     returns: v.object({ costNanos: v.number() }),
     handler: async (ctx, args) => {
         const request = await ctx.db.get(args.requestId);
+        // The request may be gone — retention purged it, or the owning bucket was
+        // deleted. A late/duplicate webhook must be an idempotent no-op, not a 500
+        // (the caller can't do anything useful with the error, and it triggers retries).
         if (!request)
-            throw new Error("Unknown request");
+            return { costNanos: 0 };
         // Expiry releases capacity, but is not evidence that the provider charged
         // nothing. Accept one final result even after expiry; duplicates stay no-ops.
         if (request.status !== "pending" && !request.reservationExpired) {
             return { costNanos: request.costNanos ?? 0 };
         }
-        // Clamp caller-supplied token counts: negatives would produce negative cost
-        // and could refund a user below their cap.
-        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));
+        // Coerce caller/provider-supplied token counts to finite nonnegative
+        // integers: negatives would refund below a cap, and NaN/Infinity (which
+        // v.number() allows) would poison every downstream total and cap check.
+        const promptTokens = safeCount(args.promptTokens);
+        const completionTokens = safeCount(args.completionTokens);
+        const cachedTokens = Math.min(promptTokens, safeCount(args.cachedTokens));
         // 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.
+        // (prompt-cache-read) slice — plus any server-tool per-call fees. Require it
+        // finite: +Infinity passes a bare `>= 0` and would corrupt the totals.
         let costNanos;
-        if (args.costNanos !== undefined && args.costNanos >= 0) {
+        if (args.costNanos !== undefined && Number.isFinite(args.costNanos) && args.costNanos >= 0) {
             costNanos = Math.round(args.costNanos);
         }
         else {
@@ -937,10 +972,15 @@ export const setBucketLimits = mutation({
     },
     returns: v.null(),
     handler: async (ctx, args) => {
-        if (args.requestsPerMinute !== undefined &&
-            (!Number.isSafeInteger(args.requestsPerMinute) || args.requestsPerMinute < 0)) {
-            throw new Error("requestsPerMinute must be a nonnegative safe integer");
-        }
+        assertAmount(args.requestsPerMinute, "requestsPerMinute");
+        assertAmount(args.maxConcurrent, "maxConcurrent");
+        assertAmount(args.dailySpendLimitNanos, "dailySpendLimitNanos");
+        assertAmount(args.monthlySpendLimitNanos, "monthlySpendLimitNanos");
+        assertAmount(args.lifetimeSpendLimitNanos, "lifetimeSpendLimitNanos");
+        assertAmount(args.dailyTokenLimit, "dailyTokenLimit");
+        assertAmount(args.monthlyTokenLimit, "monthlyTokenLimit");
+        assertAmount(args.lifetimeTokenLimit, "lifetimeTokenLimit");
+        assertFraction(args.warnAtPct, "warnAtPct");
         const bucket = await getOrCreateBucket(ctx, args.dimension, args.value);
         const { dimension: _d, value: _v, ...limits } = args;
         await ctx.db.patch(bucket._id, limits);
@@ -960,6 +1000,9 @@ export const bumpBucket = mutation({
     args: { dimension: v.string(), value: v.string(), ...vBumpArgs },
     returns: v.null(),
     handler: async (ctx, { dimension, value, dailyNanos, monthlyNanos, lifetimeNanos }) => {
+        assertAmount(dailyNanos, "dailyNanos");
+        assertAmount(monthlyNanos, "monthlyNanos");
+        assertAmount(lifetimeNanos, "lifetimeNanos");
         const bucket = await getOrCreateBucket(ctx, dimension, value);
         const today = dayStamp();
         const month = monthStamp();
@@ -989,6 +1032,8 @@ export const adjustBucket = mutation({
     },
     returns: v.null(),
     handler: async (ctx, { dimension, value, deltaNanos, tokens, reason }) => {
+        assertAmount(deltaNanos, "deltaNanos", { signed: true });
+        assertAmount(tokens, "tokens", { signed: true });
         const b = await getOrCreateBucket(ctx, dimension, value);
         const today = dayStamp();
         const month = monthStamp();
@@ -1004,6 +1049,12 @@ export const adjustBucket = mutation({
             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),
+            // Advancing the window here must also clear the OLD window's reserved
+            // holds, or an in-flight request from the previous day/month would be
+            // treated as reserving against the new window and its later release would
+            // no longer match — stranding those reserved nanos/tokens.
+            ...(dSame ? {} : { reservedTodayNanos: 0, reservedTodayTokens: 0 }),
+            ...(mSame ? {} : { reservedMonthNanos: 0, reservedMonthTokens: 0 }),
         });
         await ctx.db.insert("adjustments", { dimension, value, deltaNanos, tokens: dt, reason });
         await addUsage(ctx, dimension, value, "day", today, deltaNanos, dt, 0);
@@ -1064,6 +1115,17 @@ export const deleteBucket = mutation({
                 .withIndex("userId", (q) => q.eq("userId", value))
                 .take(DELETE_BATCH);
             for (const r of rows) {
+                // Before dropping the row, free or settle any hold it placed on OTHER
+                // (shared) buckets — an action/customer bucket this user's request
+                // reserved against. Otherwise deleting the only row that could release
+                // that hold strands the shared bucket's reservation + pendingCount
+                // forever, and a late finish would throw. A finished-but-unfolded row
+                // is folded (the charge lands on the shared buckets); a still-pending
+                // one just has its reservation released.
+                if (r.settled === false)
+                    await foldOne(ctx, r);
+                else if (r.status === "pending")
+                    await releaseReservation(ctx, r);
                 await deleteRequestTags(ctx, r._id);
                 await ctx.db.delete(r._id);
             }
@@ -1139,18 +1201,29 @@ export const getGlobalStatus = query({
 });
 export const setGlobalLimits = mutation({
     args: {
-        dailySpendLimitNanos: v.optional(v.number()),
-        lifetimeSpendLimitNanos: v.optional(v.number()),
-        enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"))),
+        // Absent = leave unchanged; explicit null = clear that limit. (A bare
+        // v.optional(number) that always rebuilt the full patch would let a
+        // one-field edit — exactly what the dashboard sends — silently wipe the
+        // other global controls by patching them to undefined.)
+        dailySpendLimitNanos: v.optional(v.union(v.number(), v.null())),
+        lifetimeSpendLimitNanos: v.optional(v.union(v.number(), v.null())),
+        enforcement: v.optional(v.union(v.literal("hard"), v.literal("soft"), v.null())),
     },
     returns: v.null(),
     handler: async (ctx, args) => {
-        // The settings fields are "global"-prefixed; map the friendly arg names.
-        const patch = {
-            globalDailySpendLimitNanos: args.dailySpendLimitNanos,
-            globalLifetimeSpendLimitNanos: args.lifetimeSpendLimitNanos,
-            globalEnforcement: args.enforcement,
-        };
+        if (typeof args.dailySpendLimitNanos === "number")
+            assertAmount(args.dailySpendLimitNanos, "dailySpendLimitNanos");
+        if (typeof args.lifetimeSpendLimitNanos === "number")
+            assertAmount(args.lifetimeSpendLimitNanos, "lifetimeSpendLimitNanos");
+        // The settings fields are "global"-prefixed; map the friendly arg names,
+        // touching ONLY the keys the caller actually passed. null -> clear.
+        const patch = {};
+        if ("dailySpendLimitNanos" in args)
+            patch.globalDailySpendLimitNanos = args.dailySpendLimitNanos ?? undefined;
+        if ("lifetimeSpendLimitNanos" in args)
+            patch.globalLifetimeSpendLimitNanos = args.lifetimeSpendLimitNanos ?? undefined;
+        if ("enforcement" in args)
+            patch.globalEnforcement = args.enforcement ?? undefined;
         const existing = await getSettings(ctx);
         if (existing) {
             await ctx.db.patch(existing._id, patch);
@@ -1216,11 +1289,11 @@ export const setPrice = mutation({
     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 ||
-            (args.cachedNanosPerMTok ?? 0) < 0) {
-            throw new Error("Prices must be non-negative");
-        }
+        // NaN/Infinity (allowed by v.number()) are just as corrupting — a NaN rate
+        // poisons every settled cost for the model — so require finite integers.
+        assertAmount(args.inputNanosPerMTok, "inputNanosPerMTok");
+        assertAmount(args.outputNanosPerMTok, "outputNanosPerMTok");
+        assertAmount(args.cachedNanosPerMTok, "cachedNanosPerMTok");
         const existing = await ctx.db
             .query("prices")
             .withIndex("model", (q) => q.eq("model", args.model))

package/package.json

@@ -7,7 +7,7 @@
     "email": "support@convex.dev",
     "url": "https://github.com/get-convex/ai-budget/issues"
   },
-  "version": "0.0.2-alpha.15",
+  "version": "0.0.2-alpha.17",
   "license": "Apache-2.0",
   "type": "module",
   "keywords": [

package/src/client/dashboard.ts

@@ -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/src/client/index.ts

@@ -71,6 +71,8 @@ export type BudgetEventInfo = {
 export type SoftLimitInfo = BudgetEventInfo & { warnings: string[] };
 export type AIBudgetOptions = {
   defaultModel?: string;
+  /** Default model for `decisions()` (the Decisions/"Jev" endpoint). */
+  defaultEvalModel?: string;
   /**
    * A *soft* limit was exceeded (request still allowed). Lets you surface budget
    * warnings even on the languageModel/Agent path where they can't be returned.
@@ -143,6 +145,17 @@ export type ChatResult = {
   notices: string[];
 };
 
+/** The tracked result of a `decisions()` call: budgeting metadata plus the
+ * structured answers from the Decisions ("Jev") endpoint. */
+export type DecisionResult = Omit<ChatResult, "text"> & {
+  /** Structured answers keyed by your question names (shape depends on each
+   * question type: `choice`, `score`, or `boolean`). */
+  answers: Record<string, any>;
+  /** The raw gateway response, including provider-specific fields (e.g.
+   * `confidence`) under `response.body`. */
+  response?: any;
+};
+
 // ---------- helpers ----------
 
 // Token counts across AI SDK versions come as plain numbers or, in v7, as a
@@ -159,8 +172,14 @@ function extractUsage(usage: any): {
   cachedTokens: number;
 } {
   return {
-    promptTokens: toTokenCount(usage?.inputTokens ?? usage?.promptTokens),
-    completionTokens: toTokenCount(usage?.outputTokens ?? usage?.completionTokens),
+    // Cover AI SDK camelCase (v5/v7) AND raw OpenAI-compatible snake_case, so a
+    // `meter` caller passing a raw provider `usage` object still gets counts.
+    promptTokens: toTokenCount(
+      usage?.inputTokens ?? usage?.promptTokens ?? usage?.prompt_tokens
+    ),
+    completionTokens: toTokenCount(
+      usage?.outputTokens ?? usage?.completionTokens ?? usage?.completion_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.
@@ -262,6 +281,7 @@ export type BumpArgs = {
 
 export class AIBudget {
   public defaultModel: string;
+  public defaultEvalModel: string;
   private onSoftLimit?: AIBudgetOptions["onSoftLimit"];
   private onThreshold?: AIBudgetOptions["onThreshold"];
   private onLimitReached?: AIBudgetOptions["onLimitReached"];
@@ -270,6 +290,7 @@ export class AIBudget {
     options?: AIBudgetOptions
   ) {
     this.defaultModel = options?.defaultModel ?? "openai/gpt-4o-mini";
+    this.defaultEvalModel = options?.defaultEvalModel ?? "typesafe/jev-1.13";
     this.onSoftLimit = options?.onSoftLimit;
     this.onThreshold = options?.onThreshold;
     this.onLimitReached = options?.onLimitReached;
@@ -448,35 +469,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: Awaited<ReturnType<typeof run>>;
     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 };
   }
 
   /**
@@ -534,6 +562,103 @@ export class AIBudget {
     );
   }
 
+  /**
+   * Budget a structured decision through the AI Gateway's Decisions ("Jev")
+   * endpoint — sugar over `meter`. Evaluates typed `questions` (choice / score /
+   * boolean) about the `state` you provide, with the same reserve→settle
+   * limits, audit log, cost tracking, and per-tag attribution as `chat`. Call
+   * from an action. `userId` defaults to the authenticated caller.
+   *
+   * Requires `@convex-dev/ai-sdk-provider` >= 0.2.1 and an `ai` version that
+   * exposes `experimental_evaluate` (AI SDK 7's evaluation interface); both are
+   * imported lazily, so consumers who never call `decisions()` are unaffected.
+   *
+   *   const { answers } = await ai.decisions(ctx, {
+   *     state: { ticket: "Customer cannot sign in" },
+   *     questions: {
+   *       priority: { type: "choice", instructions: "...", criteria: { urgent: "...", normal: "..." } },
+   *       needsReview: { type: "boolean", instructions: "..." },
+   *     },
+   *   });
+   *   answers.priority.choice; // "urgent" | "normal"
+   */
+  async decisions(
+    ctx: RunMutationCtx,
+    args: {
+      /** The evaluation model. Defaults to `defaultEvalModel` ("typesafe/jev-1.13"). */
+      model?: string;
+      /** Context the questions are evaluated against (a string or an object). */
+      state: unknown;
+      /** Typed questions (choice / score / boolean) keyed by name. */
+      questions: Record<string, unknown>;
+      /** Whom to bill. Defaults to the authenticated user (ctx.auth). */
+      userId?: string;
+      /** Attribute spend to this action name. Defaults to the calling action. */
+      action?: string;
+      /** Extra attribution dimensions to bill/limit (team, customer, env, …). */
+      tags?: Tag[];
+      /** Reserve this exact amount (nanodollars) up front — the decision cost
+       * isn't known before the call, so a hard cap is only exact with this. */
+      estimatedCostNanos?: number;
+      rerunOf?: string;
+      /** Cancel the underlying request. */
+      abortSignal?: AbortSignal;
+    }
+  ): Promise<DecisionResult> {
+    const model = args.model ?? this.defaultEvalModel;
+    // `evaluate` is an experimental, version-gated export; import it lazily and
+    // untyped so consumers on an older `ai` (who never call this) aren't broken.
+    const evaluate = ((await import("ai")) as any).experimental_evaluate;
+    if (typeof evaluate !== "function") {
+      throw new Error(
+        "ai-budget: decisions() needs `experimental_evaluate` from the `ai` " +
+          "package (AI SDK 7's evaluation interface). Upgrade `ai` to a " +
+          "version that exports it."
+      );
+    }
+    // Likewise, `evaluationModel` exists on @convex-dev/ai-sdk-provider >= 0.2.1.
+    const evaluationModel = (convexGateway as any).evaluationModel;
+    if (typeof evaluationModel !== "function") {
+      throw new Error(
+        "ai-budget: decisions() needs `convexGateway.evaluationModel` from " +
+          "@convex-dev/ai-sdk-provider >= 0.2.1. Upgrade the provider."
+      );
+    }
+    let decision: any;
+    const result = await this.meter(
+      ctx,
+      {
+        model,
+        // Store the structured request for audit/replay.
+        messages: [
+          {
+            role: "user",
+            content: JSON.stringify({ state: args.state, questions: args.questions }),
+          },
+        ],
+        userId: args.userId,
+        action: args.action,
+        tags: args.tags,
+        estimatedCostNanos: args.estimatedCostNanos,
+        rerunOf: args.rerunOf,
+      },
+      async () => {
+        decision = await evaluate({
+          model: evaluationModel(model),
+          state: args.state,
+          questions: args.questions,
+          ...(args.abortSignal ? { abortSignal: args.abortSignal } : {}),
+        });
+        return {
+          usage: decision?.usage,
+          costNanos: extractGatewayCostNanos(decision),
+        };
+      }
+    );
+    const { text: _text, ...tracking } = result;
+    return { ...tracking, answers: decision?.answers ?? {}, response: decision?.response };
+  }
+
   /**
    * An AI SDK LanguageModel that enforces limits and records usage/cost for
    * `userId` on every call. Drop it into `generateText`, `streamText`, or the
@@ -605,15 +730,10 @@ export class AIBudget {
         wrapGenerate: async ({ doGenerate, params }: any) => {
           const requestId = await begin(params);
           const start = Date.now();
+          // Only a failure of the generation itself settles as an error.
+          let result: any;
           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, {
               error: String(e),
@@ -621,6 +741,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 }: any) => {
           const requestId = await begin(params);
@@ -635,21 +764,24 @@ 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?: string) => {
-              if (settled) return;
-              settled = true;
-              return finish(requestId, {
+            // 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: Promise<{ costNanos: number }> | undefined;
+            const settle = (error?: string) =>
+              (settlement ??= finish(requestId, {
                 responseText: text,
                 error,
                 ...extractUsage(usage),
                 costNanos: extractGatewayCostNanos({ providerMetadata }),
                 latencyMs: Date.now() - start,
-              });
-            };
+              }));
             const tapped = result.stream.pipeThrough(
               new TransformStream({
-                transform(chunk: any, controller) {
+                async transform(chunk: any, controller) {
                   if (chunk?.type === "text-delta") {
                     text += chunk.delta ?? chunk.textDelta ?? "";
                   }
@@ -657,8 +789,10 @@ 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();
@@ -925,17 +1059,29 @@ export class AIBudget {
       if (!token) 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: unknown, 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: unknown) =>
+      JSON.stringify(x).replace(/[<\u2028\u2029]/g, (ch) =>
+        "\\u" + ch.charCodeAt(0).toString(16).padStart(4, "0"));
 
     const handle = async (ctx: any, request: Request): Promise<Response> => {
       const url = new URL(request.url);
@@ -999,15 +1145,19 @@ export class AIBudget {
         }
       }
 
-      // 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.
+      // 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,
-        () => JSON.stringify(`${prefix}/api`)
-      ).replace(/__TOKEN__/g, () => JSON.stringify(token));
+        () => 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",
+        },
       });
     };