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

package/README.md

@@ -604,14 +604,19 @@ endpoints verbatim. In production:
 
 `example/` is a full working demo: chat as different personas on the left; a live
 admin panel on the right — the request audit log (inspect → edit → re-run, with
-lineage), a users table (limits, soft toggle, block, bump), and per-action budgets.
+lineage), a users table (limits, soft toggle, block, bump), per-action budgets,
+and a **⚡ Burst** tab that fires N real concurrent AI requests against one
+tightly-capped budget: watch reservations appear live, some requests get
+admitted (with real settled costs), and the rest get atomically rejected by the
+cap — the reserve-then-settle admission design, visible.
 
 ![Users & Limits admin table](docs/users.png)
 ![Actions & Budgets admin table](docs/actions.png)
 
 ```sh
 cd example
-npm install
+npm install         # also links the repo-root node_modules the demo's
+                    # ../../src component imports resolve through (postinstall)
 npx convex dev      # terminal 1 — provisions a dev deployment
 npm run dev         # terminal 2 — Vite app
 ```

package/dist/component/lib.js

@@ -324,7 +324,7 @@ export const startRequest = mutation({
         // bloat the 60s rate-limit window read below.
         const reject = async (code, reason, persist = true) => {
             if (persist) {
-                await ctx.db.insert("requests", {
+                const requestId = await ctx.db.insert("requests", {
                     userId: args.userId,
                     actionName: args.actionName,
                     ...(extraTags.length ? { tags: extraTags } : {}),
@@ -334,6 +334,15 @@ export const startRequest = mutation({
                     status: "blocked",
                     error: reason,
                 });
+                // Reverse-index the blocked attempt too, so tag-filtered request logs
+                // show rejections alongside admitted traffic.
+                for (const t of extraTags) {
+                    await ctx.db.insert("requestTags", {
+                        dimension: t.dimension,
+                        value: t.value,
+                        requestId,
+                    });
+                }
             }
             return { allowed: false, code, reason };
         };
@@ -410,13 +419,17 @@ export const startRequest = mutation({
                 recentCount = recent.filter((r) => r.status !== "blocked").length;
             }
             else {
-                recentCount = (await ctx.db
+                // Tag rows also cover persisted blocked attempts; fetch each request
+                // to exclude them, matching the user/action paths above.
+                const tagRows = await ctx.db
                     .query("requestTags")
                     .withIndex("dim_value", (q) => q
                     .eq("dimension", b.dimension)
                     .eq("value", b.value)
                     .gt("_creationTime", rateCutoff))
-                    .take(limit)).length;
+                    .take(limit + 50);
+                const recent = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
+                recentCount = recent.filter((r) => r !== null && r.status !== "blocked").length;
             }
             if (recentCount >= limit) {
                 const code = b.dimension === USER_DIM ? "rate_limit" : `${b.dimension}_rate_limit`;

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

package/src/component/lib.test.ts

@@ -330,6 +330,51 @@ describe("tag-filtered request log", () => {
     expect(acme.length).toBe(1);
     expect(acme[0].userId).toBe("u");
   });
+
+  test("blocked attempts appear in the tag-filtered log", async () => {
+    const t = convexTest(schema, modules);
+    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).
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "burst",
+      value: "run-1",
+      lifetimeSpendLimitNanos: 1_000,
+    });
+    const r = await start(t, { userId: "u", tags });
+    expect(r.allowed).toBe(false);
+    const log = await t.query(api.lib.listRequests, {
+      dimension: "burst",
+      value: "run-1",
+    });
+    expect(log.length).toBe(1);
+    expect(log[0].status).toBe("blocked");
+  });
+
+  test("persisted blocked attempts don't consume a custom-tag rate limit", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "customer",
+      value: "acme",
+      requestsPerMinute: 1,
+      // also cap spend so attempts get budget-blocked (persisted) first
+      lifetimeSpendLimitNanos: 1_000,
+    });
+    const tags = [{ dimension: "customer", value: "acme" }];
+    // A budget-blocked (persisted) attempt writes a requestTags row…
+    const blocked = await start(t, { userId: "u1", tags });
+    expect(blocked.allowed).toBe(false);
+    expect(blocked.code).toBe("customer_lifetime_spend_limit");
+    // …which must NOT count toward the 1/min rate limit. Lift the spend cap:
+    // with no admitted requests in the window, the next request goes through.
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "customer",
+      value: "acme",
+      lifetimeSpendLimitNanos: 1_000_000_000,
+    });
+    const next = await start(t, { userId: "u2", tags });
+    expect(next.allowed).toBe(true);
+  });
 });
 
 describe("tagged attribution buckets", () => {

package/src/component/lib.ts

@@ -440,7 +440,7 @@ export const startRequest = mutation({
     // bloat the 60s rate-limit window read below.
     const reject = async (code: string, reason: string, persist = true) => {
       if (persist) {
-        await ctx.db.insert("requests", {
+        const requestId = await ctx.db.insert("requests", {
           userId: args.userId,
           actionName: args.actionName,
           ...(extraTags.length ? { tags: extraTags } : {}),
@@ -450,6 +450,15 @@ export const startRequest = mutation({
           status: "blocked" as const,
           error: reason,
         });
+        // Reverse-index the blocked attempt too, so tag-filtered request logs
+        // show rejections alongside admitted traffic.
+        for (const t of extraTags) {
+          await ctx.db.insert("requestTags", {
+            dimension: t.dimension,
+            value: t.value,
+            requestId,
+          });
+        }
       }
       return { allowed: false as const, code, reason };
     };
@@ -545,17 +554,19 @@ export const startRequest = mutation({
           .take(limit + 50);
         recentCount = recent.filter((r) => r.status !== "blocked").length;
       } else {
-        recentCount = (
-          await ctx.db
-            .query("requestTags")
-            .withIndex("dim_value", (q) =>
-              q
-                .eq("dimension", b.dimension)
-                .eq("value", b.value)
-                .gt("_creationTime", rateCutoff)
-            )
-            .take(limit)
-        ).length;
+        // Tag rows also cover persisted blocked attempts; fetch each request
+        // to exclude them, matching the user/action paths above.
+        const tagRows = await ctx.db
+          .query("requestTags")
+          .withIndex("dim_value", (q) =>
+            q
+              .eq("dimension", b.dimension)
+              .eq("value", b.value)
+              .gt("_creationTime", rateCutoff)
+          )
+          .take(limit + 50);
+        const recent = await Promise.all(tagRows.map((t) => ctx.db.get(t.requestId)));
+        recentCount = recent.filter((r) => r !== null && r.status !== "blocked").length;
       }
 
       if (recentCount >= limit) {