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

package/src/component/lib.test.ts

@@ -1,7 +1,9 @@
 import { convexTest } from "convex-test";
 import { describe, expect, test, vi } from "vitest";
+import shardedCounterTest from "@convex-dev/sharded-counter/test";
+import rateLimiterTest from "@convex-dev/rate-limiter/test";
 import schema from "./schema";
-import { api } from "./_generated/api";
+import { api, internal } from "./_generated/api";
 
 // convex-test loads the component's own modules; exclude convex.config (not a
 // function module) and the test files themselves.
@@ -11,6 +13,13 @@ const modules = import.meta.glob([
   "!./**/convex.config.ts",
 ]);
 
+function initTest() {
+  const t = convexTest(schema, modules);
+  rateLimiterTest.register(t);
+  shardedCounterTest.register(t);
+  return t;
+}
+
 const MODEL = "openai/gpt-4o-mini";
 const msg = (content: string) => [{ role: "user", content }];
 
@@ -46,7 +55,7 @@ const userOf = (t: any, userId: string) => bucketOf(t, "user", userId);
 
 describe("reserve / settle spend caps", () => {
   test("a daily cap below one request's reservation blocks up front", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     // one gpt-4o-mini request reserves ~480_000 nanodollars ($0.00048); a
     // 1_000-nano ($0.000001) cap can't fit it.
     await setUserLimits(t, "u", { dailySpendLimitNanos: 1_000 });
@@ -56,7 +65,7 @@ describe("reserve / settle spend caps", () => {
   });
 
   test("reservation is released and settled to the real cost", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { dailySpendLimitNanos: 1_000_000_000 }); // $1/day
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(true);
@@ -79,7 +88,7 @@ async function settleWith(t: any, requestId: any, fields: any) {
 
 describe("monthly budgets", () => {
   test("a tiny monthly cap blocks up front", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { monthlySpendLimitNanos: 1_000 });
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(false);
@@ -89,7 +98,7 @@ describe("monthly budgets", () => {
 
 describe("cache-aware pricing", () => {
   test("cached prompt tokens are billed at the discount, not full input", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     // 1M prompt, ALL cached, 0 completion. gpt-4o-mini input $0.15/Mtok; the
     // cache default is 10% of input → 0.1 * 150_000_000 = 15_000_000 nano.
@@ -104,7 +113,7 @@ describe("cache-aware pricing", () => {
   });
 
   test("an authoritative gateway cost overrides the token estimate", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     await settleWith(t, r.requestId, {
       promptTokens: 1_000_000,
@@ -118,7 +127,7 @@ describe("cache-aware pricing", () => {
 
 describe("server-tool pricing", () => {
   test("server-tool uses add a per-call fee on top of tokens", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     // 0 tokens; 3 web searches at the $0.01 default = 30_000_000 nano.
     await settleWith(t, r.requestId, {
@@ -132,7 +141,7 @@ describe("server-tool pricing", () => {
   });
 
   test("an override price is applied", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await t.mutation(api.lib.setServerToolPrice, {
       tool: "web_search",
       nanosPerCall: 12_000_000,
@@ -148,7 +157,7 @@ describe("server-tool pricing", () => {
   });
 
   test("an authoritative cost already includes tool fees (not double-charged)", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     await settleWith(t, r.requestId, {
       promptTokens: 1_000_000,
@@ -163,7 +172,7 @@ describe("server-tool pricing", () => {
 
 describe("cost known up front (image gen, per-call APIs)", () => {
   test("estimatedCostNanos drives the reservation for a hard cap", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { dailySpendLimitNanos: 100_000_000 }); // $0.10
     // A $0.13 image is known before the call; reserving it exceeds the cap,
     // even though the token estimate for the prompt alone would pass.
@@ -177,7 +186,7 @@ describe("cost known up front (image gen, per-call APIs)", () => {
   });
 
   test("admits when it fits, then settles to the real per-image cost", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { dailySpendLimitNanos: 500_000_000 });
     const r = await start(t, {
       userId: "u",
@@ -194,7 +203,7 @@ describe("cost known up front (image gen, per-call APIs)", () => {
 
 describe("async lifecycle (video jobs): begin now, settle later", () => {
   test("reserveTtlMs is stored, and settle records the real cost", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { dailySpendLimitNanos: 5_000_000_000 });
     // Reserve $2 for a long job that will settle minutes later.
     const r = await start(t, {
@@ -221,7 +230,7 @@ describe("async lifecycle (video jobs): begin now, settle later", () => {
 
 describe("durable usage history", () => {
   test("settled spend lands in a per-day usage row", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     await settleWith(t, r.requestId, { promptTokens: 1_000_000, completionTokens: 1_000_000 });
     const hist = await t.query(api.lib.usageHistory, {
@@ -237,7 +246,7 @@ describe("durable usage history", () => {
 
 describe("manual adjustments", () => {
   test("a credit reduces spend and is logged", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     await settleWith(t, r.requestId, { promptTokens: 1_000_000, completionTokens: 1_000_000 });
     await t.mutation(api.lib.adjustBucket, {
@@ -256,7 +265,7 @@ describe("manual adjustments", () => {
 
 describe("threshold alerts", () => {
   test("crossing warnAtPct returns a notice but still admits", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     // One "hi" estimate is ~480_150 nano. Cap 800_000, warn at 50% (400_000).
     await setUserLimits(t, "u", { dailySpendLimitNanos: 800_000, warnAtPct: 0.5 });
     const r = await start(t, { userId: "u" });
@@ -267,7 +276,7 @@ describe("threshold alerts", () => {
 
 describe("concurrency cap", () => {
   test("maxConcurrent blocks a second in-flight request", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { maxConcurrent: 1 });
     const first = await start(t, { userId: "u" });
     expect(first.allowed).toBe(true); // reserved, still pending
@@ -278,8 +287,67 @@ describe("concurrency cap", () => {
 });
 
 describe("per-bucket rate limits", () => {
+  test("refills continuously and does not consume other buckets on rejection", async () => {
+    vi.useFakeTimers();
+    try {
+      const t = initTest();
+      await setUserLimits(t, "u", { requestsPerMinute: 2 });
+      await t.mutation(api.lib.setBucketLimits, {
+        dimension: "action", value: "busy", requestsPerMinute: 1,
+      });
+      expect((await start(t, { userId: "other", actionName: "busy" })).allowed).toBe(true);
+      expect((await start(t, { userId: "u", actionName: "busy" })).code).toBe("action_rate_limit");
+      expect((await start(t, { userId: "u", actionName: "free" })).allowed).toBe(true);
+      expect((await start(t, { userId: "u", actionName: "free" })).allowed).toBe(true);
+      expect((await start(t, { userId: "u", actionName: "free" })).allowed).toBe(false);
+      vi.advanceTimersByTime(30_000);
+      expect((await start(t, { userId: "u", actionName: "free" })).allowed).toBe(true);
+      expect((await start(t, { userId: "u", actionName: "free" })).allowed).toBe(false);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  test("changing the rate preserves consumption and clamps to capacity", async () => {
+    vi.useFakeTimers();
+    try {
+      const t = initTest();
+      await setUserLimits(t, "u", { requestsPerMinute: 2 });
+      expect((await start(t, { userId: "u" })).allowed).toBe(true);
+      expect((await start(t, { userId: "u" })).allowed).toBe(true);
+      await setUserLimits(t, "u", { requestsPerMinute: 4 });
+      expect((await start(t, { userId: "u" })).allowed).toBe(false);
+      vi.advanceTimersByTime(15_000);
+      expect((await start(t, { userId: "u" })).allowed).toBe(true);
+      vi.advanceTimersByTime(60_000);
+      await setUserLimits(t, "u", { requestsPerMinute: 1 });
+      expect((await start(t, { userId: "u" })).allowed).toBe(true);
+      expect((await start(t, { userId: "u" })).allowed).toBe(false);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  test("zero blocks and invalid rates are rejected", async () => {
+    const t = initTest();
+    await setUserLimits(t, "u", { requestsPerMinute: 0 });
+    expect((await start(t, { userId: "u" })).code).toBe("rate_limit");
+    for (const requestsPerMinute of [-1, 0.5, Number.MAX_SAFE_INTEGER + 1]) {
+      await expect(setUserLimits(t, "u", { requestsPerMinute })).rejects.toThrow("nonnegative safe integer");
+    }
+  });
+
+  test("deleting and recreating a bucket starts a fresh rate balance", async () => {
+    const t = initTest();
+    await setUserLimits(t, "u", { requestsPerMinute: 1 });
+    expect((await start(t, { userId: "u" })).allowed).toBe(true);
+    await t.mutation(api.lib.deleteBucket, { dimension: "user", value: "u" });
+    await setUserLimits(t, "u", { requestsPerMinute: 1 });
+    expect((await start(t, { userId: "u" })).allowed).toBe(true);
+  });
+
   test("the existing user rate limit remains compatible", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { requestsPerMinute: 1 });
     const first = await start(t, { userId: "u" });
     expect(first.allowed).toBe(true);
@@ -289,7 +357,7 @@ describe("per-bucket rate limits", () => {
   });
 
   test("an action rate limit blocks the next request for that action", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await t.mutation(api.lib.setBucketLimits, {
       dimension: "action",
       value: "ai:summarize",
@@ -303,7 +371,7 @@ describe("per-bucket rate limits", () => {
   });
 
   test("a custom-tag rate limit blocks the next request for that value", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await t.mutation(api.lib.setBucketLimits, {
       dimension: "customer",
       value: "acme",
@@ -320,7 +388,7 @@ describe("per-bucket rate limits", () => {
 
 describe("tag-filtered request log", () => {
   test("listRequests filters by a custom tag dimension", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await start(t, { userId: "u", tags: [{ dimension: "customer", value: "acme" }] });
     await start(t, { userId: "u", tags: [{ dimension: "customer", value: "globex" }] });
     const acme = await t.query(api.lib.listRequests, {
@@ -332,7 +400,7 @@ describe("tag-filtered request log", () => {
   });
 
   test("blocked attempts appear in the tag-filtered log", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const tags = [{ dimension: "burst", value: "run-1" }];
     // Cap the tag bucket below one request's reservation so the attempt is
     // budget-blocked (a persisted rejection).
@@ -352,7 +420,7 @@ describe("tag-filtered request log", () => {
   });
 
   test("persisted blocked attempts don't consume a custom-tag rate limit", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await t.mutation(api.lib.setBucketLimits, {
       dimension: "customer",
       value: "acme",
@@ -379,7 +447,7 @@ describe("tag-filtered request log", () => {
 
 describe("tagged attribution buckets", () => {
   test("a cap on a custom tag blocks, and settlement accrues to every bucket", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     // A tiny cap on customer "acme" — the user is uncapped.
     await t.mutation(api.lib.setBucketLimits, {
       dimension: "customer",
@@ -411,7 +479,7 @@ describe("tagged attribution buckets", () => {
 
 describe("D-00 exactly-once settlement", () => {
   test("a duplicate finishRequest does not double-count", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u" });
     await settle(t, r.requestId); // first settle
     const before = await userOf(t, "u");
@@ -426,7 +494,7 @@ describe("D-00 exactly-once settlement", () => {
 
 describe("token quotas", () => {
   test("a tiny daily token cap blocks (estimate exceeds it)", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", { dailyTokenLimit: 10 });
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(false);
@@ -436,7 +504,7 @@ describe("token quotas", () => {
 
 describe("soft enforcement", () => {
   test("over a soft budget: allowed, warned, flagged overBudget", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await setUserLimits(t, "u", {
       dailySpendLimitNanos: 1, // 1 nanodollar — one estimate blows past it
       enforcement: "soft",
@@ -451,7 +519,7 @@ describe("soft enforcement", () => {
 
 describe("model policy", () => {
   test("allowlist blocks an off-list model", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await t.mutation(api.lib.setModelPolicy, {
       mode: "allowlist",
       models: ["openai/gpt-4o-mini"],
@@ -466,7 +534,7 @@ describe("model policy", () => {
 
 describe("D-02 pricing validation", () => {
   test("setPrice rejects negative rates", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     await expect(
       t.mutation(api.lib.setPrice, {
         model: "x/y",
@@ -479,7 +547,7 @@ describe("D-02 pricing validation", () => {
 
 describe("F-04 fail-closed pricing", () => {
   test("an unknown model is charged the conservative max, not zero", async () => {
-    const t = convexTest(schema, modules);
+    const t = initTest();
     const r = await start(t, { userId: "u", model: "made/up-model" });
     expect(r.allowed).toBe(true);
     await settle(t, r.requestId, 1_000_000, 1_000_000);
@@ -490,3 +558,152 @@ describe("F-04 fail-closed pricing", () => {
     expect(req.unpricedModel).toBe(true);
   });
 });
+
+describe("accounting lifecycle regressions", () => {
+  test("old-day and old-month settlements preserve new holds", async () => {
+    vi.useFakeTimers();
+    try {
+      vi.setSystemTime(new Date("2026-09-30T23:59:00Z"));
+      const t = initTest();
+      await setUserLimits(t, "u", { dailySpendLimitNanos: 1000, monthlySpendLimitNanos: 1000 });
+      const old = await start(t, { userId: "u", estimatedCostNanos: 100 });
+      vi.setSystemTime(new Date("2026-10-01T00:01:00Z"));
+      await start(t, { userId: "u", estimatedCostNanos: 100 });
+      await t.mutation(api.lib.finishRequest, { requestId: old.requestId, costNanos: 0 });
+      await t.mutation(internal.lib.foldTotals, { requestId: old.requestId });
+      const b = await userOf(t, "u");
+      expect(b.reservedTodayNanos).toBe(100);
+      expect(b.reservedMonthNanos).toBe(100);
+      expect(b.reservedTotalNanos).toBe(100);
+      expect(b.pendingCount).toBe(1);
+    } finally { vi.useRealTimers(); }
+  });
+
+  test("a request admitted before caps were enabled cannot release a later hold", async () => {
+    const t = initTest();
+    const old = await start(t, { userId: "u", estimatedCostNanos: 100 });
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 1000 });
+    await start(t, { userId: "u", estimatedCostNanos: 100 });
+    await t.mutation(api.lib.finishRequest, { requestId: old.requestId, costNanos: 0 });
+    await t.mutation(internal.lib.foldTotals, { requestId: old.requestId });
+    expect((await userOf(t, "u")).reservedTotalNanos).toBe(100);
+  });
+
+  test("expiry releases once, survives retention, and accepts one late charge", async () => {
+    vi.useFakeTimers();
+    try {
+      const t = initTest();
+      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, {});
+      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 });
+      await t.mutation(internal.lib.foldTotals, { requestId: job.requestId });
+      await t.mutation(api.lib.finishRequest, { requestId: job.requestId, costNanos: 999 });
+      await t.mutation(internal.lib.foldTotals, { requestId: job.requestId });
+      const b = await userOf(t, "u");
+      expect(b.totalSpendNanos).toBe(75);
+      expect(b.totalRequests).toBe(1);
+      expect(b.reservedTotalNanos).toBe(100);
+      expect((await t.query(api.lib.getGlobalStatus, {})).spentTotalNanos).toBe(75);
+    } finally { vi.useRealTimers(); }
+  });
+
+  test("long TTL jobs cannot hide expired jobs", async () => {
+    vi.useFakeTimers();
+    try {
+      const t = initTest();
+      await t.run(async ctx => {
+        for (let i = 0; i < 201; i++) await ctx.db.insert("requests", {
+          userId: "long", model: MODEL, messages: [], status: "pending",
+          expiresAt: Date.now() + 86400_000, heldBucketIds: [],
+        });
+      });
+      const short = await start(t, { userId: "short" });
+      vi.advanceTimersByTime(31 * 60_000);
+      const result = await t.mutation(internal.lib.reconcile, {});
+      expect(result.expired).toBe(1);
+      expect((await t.run(ctx => ctx.db.get(short.requestId))).reservationExpired).toBe(true);
+    } finally { vi.useRealTimers(); }
+  });
+
+  test("global accounting includes usage before limits are enabled", async () => {
+    const t = initTest();
+    const job = await start(t, { userId: "u" });
+    await t.mutation(api.lib.finishRequest, { requestId: job.requestId, costNanos: 100 });
+    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 });
+    expect((await start(t, { userId: "u", estimatedCostNanos: 1 })).allowed).toBe(false);
+  });
+
+  test("settlement does not write admission policy, but changing limits does", async () => {
+    const t = initTest();
+    const job = await start(t, { userId: "u" });
+    const before = await t.run(ctx => ctx.db.query("bucketPolicies").collect());
+    await t.mutation(api.lib.finishRequest, { requestId: job.requestId, costNanos: 100 });
+    await t.mutation(internal.lib.foldTotals, { requestId: job.requestId });
+    expect(await t.run(ctx => ctx.db.query("bucketPolicies").collect())).toEqual(before);
+    await setUserLimits(t, "u", { blocked: true });
+    expect((await start(t, { userId: "u" })).allowed).toBe(false);
+  });
+});
+
+test("legacy pending rows acquire deadlines without starving newer expired work", async () => {
+  vi.useFakeTimers();
+  try {
+    const t = initTest();
+    await t.run(async ctx => {
+      for (let i = 0; i < 201; i++) await ctx.db.insert("requests", {
+        userId: "legacy", model: MODEL, messages: [], status: "pending", reserveTtlMs: 86400_000,
+      });
+    });
+    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.run(ctx => ctx.db.get(job.requestId))).reservationExpired).toBe(true);
+    await t.mutation(internal.lib.reconcile, {});
+    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(); }
+});
+
+test("retention progresses past unresolved jobs", async () => {
+  vi.useFakeTimers();
+  try {
+    const t = initTest();
+    await t.run(async ctx => {
+      for (let i = 0; i < 501; i++) await ctx.db.insert("requests", {
+        userId: "long", model: MODEL, messages: [], status: "pending", expiresAt: Date.now() + 86400_000,
+      });
+    });
+    const job = await start(t, { userId: "short" });
+    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.run(ctx => ctx.db.get(job.requestId))).toBeNull();
+  } finally { vi.useRealTimers(); }
+});
+
+test("delayed folding attributes spend to completion day and leaves newer holds intact", async () => {
+  vi.useFakeTimers();
+  try {
+    vi.setSystemTime(new Date("2026-09-30T23:59:00Z"));
+    const t = initTest();
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 1000 });
+    const job = await start(t, { userId: "u", estimatedCostNanos: 100 });
+    await t.mutation(api.lib.finishRequest, { requestId: job.requestId, costNanos: 50 });
+    vi.setSystemTime(new Date("2026-10-01T00:01:00Z"));
+    await start(t, { userId: "u", estimatedCostNanos: 100 });
+    await t.mutation(internal.lib.foldTotals, { requestId: job.requestId });
+    const b = await userOf(t, "u");
+    expect(b.spendTodayNanos).toBe(0);
+    expect(b.reservedTodayNanos).toBe(100);
+    const history = await t.query(api.lib.usageHistory, { dimension: "user", value: "u", period: "day" });
+    expect(history[0].stamp).toBe("2026-09-30");
+    expect(history[0].spendNanos).toBe(50);
+  } finally { vi.useRealTimers(); }
+});