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

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.17",
+  "version": "0.0.2-alpha.19",
   "license": "Apache-2.0",
   "type": "module",
   "keywords": [

package/src/client/index.ts

@@ -172,13 +172,20 @@ function extractUsage(usage: any): {
   cachedTokens: number;
 } {
   return {
-    // 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.
+    // 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?.inputTokens ??
+        usage?.promptTokens ??
+        usage?.prompt_tokens ??
+        usage?.input_tokens
     ),
     completionTokens: toTokenCount(
-      usage?.outputTokens ?? usage?.completionTokens ?? usage?.completion_tokens
+      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
@@ -188,7 +195,8 @@ function extractUsage(usage: any): {
         usage?.cachedInputTokens ??
         usage?.promptTokensDetails?.cachedTokens ??
         usage?.prompt_tokens_details?.cached_tokens ??
-        usage?.cached_tokens
+        usage?.cached_tokens ??
+        usage?.cache_read_input_tokens
     ),
   };
 }
@@ -214,6 +222,12 @@ function extractGatewayCostNanos(result: any): number | 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: string) =>
+  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: any): Message[] {
   if (!Array.isArray(prompt)) return [];
@@ -223,14 +237,20 @@ function simplifyPrompt(prompt: any): Message[] {
       content = m.content;
     } else if (Array.isArray(m.content)) {
       content = m.content
-        .map((part: any) =>
-          part?.type === "text" ? part.text : JSON.stringify(part)
-        )
+        .map((part: any) => {
+          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) };
   });
 }
 
@@ -781,7 +801,7 @@ export class AIBudget {
               }));
             const tapped = result.stream.pipeThrough(
               new TransformStream({
-                async transform(chunk: any, controller) {
+                async transform(chunk: any, controller: any) {
                   if (chunk?.type === "text-delta") {
                     text += chunk.delta ?? chunk.textDelta ?? "";
                   }
@@ -797,7 +817,22 @@ export class AIBudget {
                 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: any) {
+                  if (usage === undefined && text) {
+                    usage = { outputTokens: Math.ceil(text.length / 4) };
+                  }
+                  await settle(`cancelled: ${String(reason)}`);
+                },
+              } as any)
             );
             return { ...result, stream: tapped };
           } catch (e) {
@@ -956,13 +991,17 @@ export class AIBudget {
     return {
       /** Limits + spend today/total. */
       status: (ctx: RunQueryCtx) => ctx.runQuery(c.lib.getGlobalStatus, {}),
-      /** A killswitch spend cap across all users/actions (enforced approximately). */
+      /**
+       * A killswitch spend cap across all users/actions. Enforced
+       * **approximately** (bounded overshoot, no per-request reservation) — for
+       * an exact ceiling use a per-bucket cap. Pass `null` to clear a field.
+       */
       setLimits: (
         ctx: RunMutationCtx,
         args: {
-          dailySpendLimitNanos?: number;
-          lifetimeSpendLimitNanos?: number;
-          enforcement?: "hard" | "soft";
+          dailySpendLimitNanos?: number | null;
+          lifetimeSpendLimitNanos?: number | null;
+          enforcement?: "approximate" | "soft" | null;
         }
       ) => ctx.runMutation(c.lib.setGlobalLimits, args),
       bump: (
@@ -975,6 +1014,17 @@ export class AIBudget {
       /** Request-row retention window in ms (default 1h; 0 disables). */
       setRetention: (ctx: RunMutationCtx, args: { retentionMs: number }) =>
         ctx.runMutation(c.lib.setRetention, args),
+      /**
+       * Deployment-wide data/pricing policy (only the fields you pass change):
+       * - `allowUnpricedModels: false` rejects models with no configured price
+       *   under hard enforcement (default true — charge the conservative fallback).
+       * - `storeContent: false` stops persisting prompt/response content on
+       *   request rows (default true).
+       */
+      setPolicy: (
+        ctx: RunMutationCtx,
+        args: { allowUnpricedModels?: boolean; storeContent?: boolean }
+      ) => ctx.runMutation(c.lib.setDeploymentPolicy, args),
     };
   }
 

package/src/component/_generated/component.ts

@@ -95,7 +95,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         {
           dailySpendLimitNanos: number | null;
           defaultWarnAtPct: number | null;
-          enforcement: "hard" | "soft";
+          enforcement: "approximate" | "soft";
           lifetimeSpendLimitNanos: number | null;
           retentionMs: number | null;
           spentTodayNanos: number;
@@ -181,13 +181,20 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         null,
         Name
       >;
+      setDeploymentPolicy: FunctionReference<
+        "mutation",
+        "internal",
+        { allowUnpricedModels?: boolean; storeContent?: boolean },
+        null,
+        Name
+      >;
       setGlobalLimits: FunctionReference<
         "mutation",
         "internal",
         {
-          dailySpendLimitNanos?: number;
-          enforcement?: "hard" | "soft";
-          lifetimeSpendLimitNanos?: number;
+          dailySpendLimitNanos?: number | null;
+          enforcement?: "approximate" | "soft" | null;
+          lifetimeSpendLimitNanos?: number | null;
         },
         null,
         Name

package/src/component/lib.test.ts

@@ -245,7 +245,7 @@ describe("durable usage history", () => {
 });
 
 describe("manual adjustments", () => {
-  test("a credit reduces spend and is logged", async () => {
+  test("a credit accrues separately from gross spend and grants headroom", async () => {
     const t = initTest();
     const r = await start(t, { userId: "u" });
     await settleWith(t, r.requestId, { promptTokens: 1_000_000, completionTokens: 1_000_000 });
@@ -256,11 +256,26 @@ describe("manual adjustments", () => {
       reason: "goodwill credit",
     });
     const u = await userOf(t, "u");
-    expect(u.totalSpendNanos).toBe(500_000_000); // 750M - 250M
+    // Gross spend is unchanged (credits never reduce it); the credit is tracked
+    // separately. Net = gross - credits = 500M.
+    expect(u.totalSpendNanos).toBe(750_000_000);
+    expect(u.creditsNanos).toBe(250_000_000);
     const log = await t.query(api.lib.listAdjustments, { dimension: "user", value: "u" });
     expect(log.length).toBe(1);
     expect(log[0].deltaNanos).toBe(-250_000_000);
   });
+
+  test("a credit grants headroom under a cap; a debit consumes gross spend", async () => {
+    const t = initTest();
+    await setUserLimits(t, "u", { lifetimeSpendLimitNanos: 1_000_000_000 });
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, { costNanos: 900_000_000 }); // $0.90 gross
+    // Right at the edge: a $0.20 estimate would exceed the $1 cap...
+    expect((await start(t, { userId: "u", estimatedCostNanos: 200_000_000 })).allowed).toBe(false);
+    // ...but a $0.30 credit (net spend $0.60) reopens headroom.
+    await t.mutation(api.lib.adjustBucket, { dimension: "user", value: "u", deltaNanos: -300_000_000 });
+    expect((await start(t, { userId: "u", estimatedCostNanos: 200_000_000 })).allowed).toBe(true);
+  });
 });
 
 describe("threshold alerts", () => {
@@ -554,8 +569,9 @@ describe("F-04 fail-closed pricing", () => {
     expect(r.allowed).toBe(true);
     await settle(t, r.requestId, 1_000_000, 1_000_000);
     const u = await userOf(t, "u");
-    // conservative = max over table = {$3 in, $15 out}/Mtok => $18 = 18e9 nano.
-    expect(u.totalSpendNanos).toBe(18_000_000_000);
+    // Conservative fallback is a frontier ceiling of {$20 in, $100 out}/Mtok, so
+    // 1M in + 1M out => $120 = 120e9 nano (over-count is the safe direction).
+    expect(u.totalSpendNanos).toBe(120_000_000_000);
     const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
     expect(req.unpricedModel).toBe(true);
   });
@@ -598,7 +614,7 @@ describe("accounting lifecycle regressions", () => {
       await setUserLimits(t, "u", { dailySpendLimitNanos: 1000 });
       const job = await start(t, { userId: "u", estimatedCostNanos: 100 });
       vi.advanceTimersByTime(2 * 60 * 60_000);
-      await t.mutation(internal.lib.reconcile, {});
+      await t.mutation(internal.lib.expirePhase, {});
       expect((await userOf(t, "u")).reservedTotalNanos).toBe(0);
       await start(t, { userId: "u", estimatedCostNanos: 100 });
       await t.mutation(api.lib.finishRequest, { requestId: job.requestId, costNanos: 75 });
@@ -625,7 +641,7 @@ describe("accounting lifecycle regressions", () => {
       });
       const short = await start(t, { userId: "short" });
       vi.advanceTimersByTime(31 * 60_000);
-      const result = await t.mutation(internal.lib.reconcile, {});
+      const result = await t.mutation(internal.lib.expirePhase, {});
       expect(result.expired).toBe(1);
       expect((await t.run(ctx => ctx.db.get(short.requestId))).reservationExpired).toBe(true);
     } finally { vi.useRealTimers(); }
@@ -638,6 +654,11 @@ describe("accounting lifecycle regressions", () => {
     await t.mutation(internal.lib.foldTotals, { requestId: job.requestId });
     expect((await t.query(api.lib.getGlobalStatus, {})).spentTotalNanos).toBe(100);
     await t.mutation(api.lib.setGlobalLimits, { lifetimeSpendLimitNanos: 100 });
+    // The killswitch trips out-of-band (H4): the reconciler's globalPhase
+    // compares the sharded total to the cap and flags settings; admission reads
+    // the flag. So a request admits until the flag is set, then blocks.
+    expect((await start(t, { userId: "u", estimatedCostNanos: 1 })).allowed).toBe(true);
+    await t.mutation(internal.lib.globalPhase, {});
     expect((await start(t, { userId: "u", estimatedCostNanos: 1 })).allowed).toBe(false);
   });
 
@@ -664,9 +685,11 @@ test("legacy pending rows acquire deadlines without starving newer expired work"
     });
     const job = await start(t, { userId: "new" });
     vi.advanceTimersByTime(31 * 60_000);
-    expect((await t.mutation(internal.lib.reconcile, {})).expired).toBe(1);
+    expect((await t.mutation(internal.lib.expirePhase, {})).expired).toBe(1);
     expect((await t.run(ctx => ctx.db.get(job.requestId))).reservationExpired).toBe(true);
-    await t.mutation(internal.lib.reconcile, {});
+    // The phase self-reschedules to backfill the remaining legacy rows in
+    // batches; drain those scheduled continuations.
+    await t.finishAllScheduledFunctions(vi.runAllTimers);
     expect(await t.run(ctx => ctx.db.query("requests").withIndex("status_expires", q =>
       q.eq("status", "pending").eq("expiresAt", undefined)).take(1))).toHaveLength(0);
   } finally { vi.useRealTimers(); }
@@ -685,7 +708,7 @@ test("retention progresses past unresolved jobs", async () => {
     await t.mutation(api.lib.finishRequest, { requestId: job.requestId, costNanos: 0 });
     await t.mutation(internal.lib.foldTotals, { requestId: job.requestId });
     vi.advanceTimersByTime(2 * 60 * 60_000);
-    expect((await t.mutation(internal.lib.reconcile, {})).purged).toBe(1);
+    expect((await t.mutation(internal.lib.retentionPhase, {})).purged).toBe(1);
     expect(await t.run(ctx => ctx.db.get(job.requestId))).toBeNull();
   } finally { vi.useRealTimers(); }
 });
@@ -791,3 +814,80 @@ describe("v1 hardening", () => {
     ).rejects.toThrow(/dailySpendLimitNanos/);
   });
 });
+
+describe("v1 hardening (round 2)", () => {
+  test("a non-finite estimatedCostNanos is rejected before it can poison a bucket", async () => {
+    const t = initTest();
+    for (const estimatedCostNanos of [Infinity, NaN, -1]) {
+      await expect(
+        start(t, { userId: "u", estimatedCostNanos })
+      ).rejects.toThrow(/estimatedCostNanos/);
+    }
+  });
+
+  test("a settle with no cost signal falls back to the reserved estimate, not $0", async () => {
+    const t = initTest();
+    // $2 reserved up front (e.g. a video job).
+    const r = await start(t, { userId: "u", estimatedCostNanos: 2_000_000_000 });
+    // Settle with an unpriced server tool and no tokens/authoritative cost.
+    const out = await t.mutation(api.lib.finishRequest, {
+      requestId: r.requestId,
+      serverToolUses: { video_seconds: 8 }, // no configured price
+    });
+    expect(out.costNanos).toBe(2_000_000_000); // NOT 0
+    vi.useFakeTimers();
+    await t.finishAllScheduledFunctions(vi.runAllTimers);
+    vi.useRealTimers();
+    const u = await userOf(t, "u");
+    expect(u.totalSpendNanos).toBe(2_000_000_000);
+  });
+});
+
+describe("v1 hardening (round 3): reconcile phases", () => {
+  test("expired tombstones are deleted after the late-settle horizon", async () => {
+    vi.useFakeTimers();
+    try {
+      const t = initTest();
+      const job = await start(t, { userId: "u", estimatedCostNanos: 100 });
+      vi.advanceTimersByTime(31 * 60_000);
+      // Expire it into a billing tombstone (reservationExpired + settled).
+      await t.mutation(internal.lib.expirePhase, {});
+      // Within the 7-day late-settle horizon: retention keeps the tombstone.
+      await t.mutation(internal.lib.retentionPhase, {});
+      expect(await t.run((ctx) => ctx.db.get(job.requestId))).not.toBeNull();
+      // Past the horizon: the tombstone is deleted so they can't accumulate.
+      vi.advanceTimersByTime(8 * 24 * 60 * 60_000);
+      await t.mutation(internal.lib.retentionPhase, {});
+      expect(await t.run((ctx) => ctx.db.get(job.requestId))).toBeNull();
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+});
+
+describe("v1: deployment policy (unpriced models, content storage)", () => {
+  test("blocking unpriced models rejects a model with no configured price", async () => {
+    const t = initTest();
+    await t.mutation(api.lib.setDeploymentPolicy, { allowUnpricedModels: false });
+    const r = await start(t, { userId: "u", model: "made/up-model" });
+    expect(r.allowed).toBe(false);
+    expect(r.code).toBe("model_unpriced");
+    // A priced model still admits.
+    expect((await start(t, { userId: "u", model: MODEL })).allowed).toBe(true);
+  });
+
+  test("storeContent:false keeps metadata but drops prompt/response content", async () => {
+    const t = initTest();
+    await t.mutation(api.lib.setDeploymentPolicy, { storeContent: false });
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, {
+      promptTokens: 10,
+      completionTokens: 5,
+      responseText: "secret answer",
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.messages).toEqual([]); // prompt not stored
+    expect(req.responseText ?? undefined).toBe(undefined); // response not stored
+    expect(req.promptTokens).toBe(10); // metadata/cost still recorded
+  });
+});