@convex-dev/ai-budget 0.0.2-alpha.0 → 0.0.2-alpha.12

package/src/component/_generated/component.ts

@@ -24,31 +24,43 @@ import type { FunctionReference } from "convex/server";
 export type ComponentApi<Name extends string | undefined = string | undefined> =
   {
     lib: {
-      bumpAction: FunctionReference<
+      adjustBucket: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number; name: string },
+        {
+          deltaNanos: number;
+          dimension: string;
+          reason?: string;
+          tokens?: number;
+          value: string;
+        },
         null,
         Name
       >;
-      bumpGlobal: FunctionReference<
+      bumpBucket: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number },
+        {
+          dailyNanos?: number;
+          dimension: string;
+          lifetimeNanos?: number;
+          monthlyNanos?: number;
+          value: string;
+        },
         null,
         Name
       >;
-      bumpUser: FunctionReference<
+      bumpGlobal: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number; userId: string },
+        { dailyNanos?: number; lifetimeNanos?: number; monthlyNanos?: number },
         null,
         Name
       >;
-      deleteUser: FunctionReference<
+      deleteBucket: FunctionReference<
         "mutation",
         "internal",
-        { userId: string },
+        { dimension: string; value: string },
         { deletedThisBatch: number; done: boolean },
         Name
       >;
@@ -58,23 +70,34 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         {
           cachedTokens?: number;
           completionTokens?: number;
+          costNanos?: number;
           error?: string;
           latencyMs?: number;
           promptTokens?: number;
           requestId: string;
           responseText?: string;
+          serverToolUses?: Record<string, number>;
         },
         { costNanos: number },
         Name
       >;
+      getBucket: FunctionReference<
+        "query",
+        "internal",
+        { dimension: string; value: string },
+        any,
+        Name
+      >;
       getGlobalStatus: FunctionReference<
         "query",
         "internal",
         {},
         {
           dailySpendLimitNanos: number | null;
+          defaultWarnAtPct: number | null;
           enforcement: "hard" | "soft";
           lifetimeSpendLimitNanos: number | null;
+          retentionMs: number | null;
           spentTodayNanos: number;
           spentTotalNanos: number;
         },
@@ -101,54 +124,70 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         any,
         Name
       >;
-      listActions: FunctionReference<"query", "internal", {}, any, Name>;
+      listAdjustments: FunctionReference<
+        "query",
+        "internal",
+        { dimension: string; limit?: number; value: string },
+        any,
+        Name
+      >;
+      listBuckets: FunctionReference<
+        "query",
+        "internal",
+        { dimension?: string },
+        any,
+        Name
+      >;
       listPrices: FunctionReference<"query", "internal", {}, any, Name>;
       listRequests: FunctionReference<
         "query",
         "internal",
-        { limit?: number; userId?: string },
+        { dimension?: string; limit?: number; userId?: string; value?: string },
+        any,
+        Name
+      >;
+      listServerToolPrices: FunctionReference<
+        "query",
+        "internal",
+        {},
         any,
         Name
       >;
-      listUsers: FunctionReference<"query", "internal", {}, any, Name>;
-      setActionLimits: FunctionReference<
+      setAlertDefaults: FunctionReference<
         "mutation",
         "internal",
-        {
-          dailySpendLimitNanos?: number;
-          dailyTokenLimit?: number;
-          disabled?: boolean;
-          enforcement?: "hard" | "soft";
-          lifetimeSpendLimitNanos?: number;
-          lifetimeTokenLimit?: number;
-          name: string;
-        },
+        { warnAtPct?: number },
         null,
         Name
       >;
-      setGlobalLimits: FunctionReference<
+      setBucketLimits: FunctionReference<
         "mutation",
         "internal",
         {
+          blocked?: boolean;
           dailySpendLimitNanos?: number;
+          dailyTokenLimit?: number;
+          dimension: string;
           enforcement?: "hard" | "soft";
           lifetimeSpendLimitNanos?: number;
+          lifetimeTokenLimit?: number;
+          maxConcurrent?: number;
+          monthlySpendLimitNanos?: number;
+          monthlyTokenLimit?: number;
+          requestsPerMinute?: number;
+          value: string;
+          warnAtPct?: number;
         },
         null,
         Name
       >;
-      setLimits: FunctionReference<
+      setGlobalLimits: FunctionReference<
         "mutation",
         "internal",
         {
-          blocked?: boolean;
           dailySpendLimitNanos?: number;
-          dailyTokenLimit?: number;
           enforcement?: "hard" | "soft";
           lifetimeSpendLimitNanos?: number;
-          lifetimeTokenLimit?: number;
-          requestsPerMinute?: number;
-          userId: string;
         },
         null,
         Name
@@ -164,6 +203,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         "mutation",
         "internal",
         {
+          cachedNanosPerMTok?: number;
           inputNanosPerMTok: number;
           model: string;
           outputNanosPerMTok: number;
@@ -178,6 +218,13 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         null,
         Name
       >;
+      setServerToolPrice: FunctionReference<
+        "mutation",
+        "internal",
+        { nanosPerCall: number; tool: string },
+        null,
+        Name
+      >;
       startRequest: FunctionReference<
         "mutation",
         "internal",
@@ -186,11 +233,29 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           messages: Array<{ content: string; role: string }>;
           model: string;
           rerunOf?: string;
+          tags?: Array<{ dimension: string; value: string }>;
           userId: string;
         },
-        | { allowed: true; requestId: string; warnings: Array<string> }
+        | {
+            allowed: true;
+            notices: Array<string>;
+            requestId: string;
+            warnings: Array<string>;
+          }
         | { allowed: false; code: string; reason: string },
         Name
       >;
+      usageHistory: FunctionReference<
+        "query",
+        "internal",
+        {
+          dimension: string;
+          limit?: number;
+          period: "day" | "month";
+          value: string;
+        },
+        any,
+        Name
+      >;
     };
   };

package/src/component/lib.test.ts

@@ -32,18 +32,24 @@ async function settle(t: any, requestId: any, p = 10, c = 5) {
   await t.finishAllScheduledFunctions(vi.runAllTimers);
   vi.useRealTimers();
 }
-const userOf = async (t: any, userId: string) =>
-  (await t.query(api.lib.listUsers, {})).find((u: any) => u.userId === userId);
+const setUserLimits = (t: any, userId: string, limits: any) =>
+  t.mutation(api.lib.setBucketLimits, {
+    dimension: "user",
+    value: userId,
+    ...limits,
+  });
+const bucketOf = async (t: any, dimension: string, value: string) =>
+  (await t.query(api.lib.listBuckets, { dimension })).find(
+    (b: any) => b.value === value
+  );
+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);
     // one gpt-4o-mini request reserves ~480_000 nanodollars ($0.00048); a
     // 1_000-nano ($0.000001) cap can't fit it.
-    await t.mutation(api.lib.setLimits, {
-      userId: "u",
-      dailySpendLimitNanos: 1_000,
-    });
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 1_000 });
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(false);
     expect(r.code).toBe("user_daily_spend_limit");
@@ -51,10 +57,7 @@ describe("reserve / settle spend caps", () => {
 
   test("reservation is released and settled to the real cost", async () => {
     const t = convexTest(schema, modules);
-    await t.mutation(api.lib.setLimits, {
-      userId: "u",
-      dailySpendLimitNanos: 1_000_000_000, // $1/day
-    });
+    await setUserLimits(t, "u", { dailySpendLimitNanos: 1_000_000_000 }); // $1/day
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(true);
     await settle(t, r.requestId, 1_000_000, 1_000_000); // 1M in, 1M out
@@ -67,6 +70,242 @@ describe("reserve / settle spend caps", () => {
   });
 });
 
+async function settleWith(t: any, requestId: any, fields: any) {
+  await t.mutation(api.lib.finishRequest, { requestId, ...fields });
+  vi.useFakeTimers();
+  await t.finishAllScheduledFunctions(vi.runAllTimers);
+  vi.useRealTimers();
+}
+
+describe("monthly budgets", () => {
+  test("a tiny monthly cap blocks up front", async () => {
+    const t = convexTest(schema, modules);
+    await setUserLimits(t, "u", { monthlySpendLimitNanos: 1_000 });
+    const r = await start(t, { userId: "u" });
+    expect(r.allowed).toBe(false);
+    expect(r.code).toBe("user_monthly_spend_limit");
+  });
+});
+
+describe("cache-aware pricing", () => {
+  test("cached prompt tokens are billed at the discount, not full input", async () => {
+    const t = convexTest(schema, modules);
+    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.
+    await settleWith(t, r.requestId, {
+      promptTokens: 1_000_000,
+      completionTokens: 0,
+      cachedTokens: 1_000_000,
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(15_000_000);
+    expect(req.cachedTokens).toBe(1_000_000);
+  });
+
+  test("an authoritative gateway cost overrides the token estimate", async () => {
+    const t = convexTest(schema, modules);
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, {
+      promptTokens: 1_000_000,
+      completionTokens: 1_000_000,
+      costNanos: 12_345,
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(12_345);
+  });
+});
+
+describe("server-tool pricing", () => {
+  test("server-tool uses add a per-call fee on top of tokens", async () => {
+    const t = convexTest(schema, modules);
+    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, {
+      promptTokens: 0,
+      completionTokens: 0,
+      serverToolUses: { web_search: 3 },
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(30_000_000);
+    expect(req.serverToolUses).toEqual({ web_search: 3 });
+  });
+
+  test("an override price is applied", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setServerToolPrice, {
+      tool: "web_search",
+      nanosPerCall: 12_000_000,
+    });
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, {
+      promptTokens: 0,
+      completionTokens: 0,
+      serverToolUses: { web_search: 2 },
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(24_000_000);
+  });
+
+  test("an authoritative cost already includes tool fees (not double-charged)", async () => {
+    const t = convexTest(schema, modules);
+    const r = await start(t, { userId: "u" });
+    await settleWith(t, r.requestId, {
+      promptTokens: 1_000_000,
+      completionTokens: 0,
+      serverToolUses: { web_search: 5 },
+      costNanos: 999,
+    });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
+    expect(req.costNanos).toBe(999);
+  });
+});
+
+describe("durable usage history", () => {
+  test("settled spend lands in a per-day usage row", async () => {
+    const t = convexTest(schema, modules);
+    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, {
+      dimension: "user",
+      value: "u",
+      period: "day",
+    });
+    expect(hist.length).toBe(1);
+    expect(hist[0].spendNanos).toBe(750_000_000); // $0.75
+    expect(hist[0].requests).toBe(1);
+  });
+});
+
+describe("manual adjustments", () => {
+  test("a credit reduces spend and is logged", async () => {
+    const t = convexTest(schema, modules);
+    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, {
+      dimension: "user",
+      value: "u",
+      deltaNanos: -250_000_000,
+      reason: "goodwill credit",
+    });
+    const u = await userOf(t, "u");
+    expect(u.totalSpendNanos).toBe(500_000_000); // 750M - 250M
+    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);
+  });
+});
+
+describe("threshold alerts", () => {
+  test("crossing warnAtPct returns a notice but still admits", async () => {
+    const t = convexTest(schema, modules);
+    // 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" });
+    expect(r.allowed).toBe(true);
+    expect(r.notices.length).toBeGreaterThan(0);
+  });
+});
+
+describe("concurrency cap", () => {
+  test("maxConcurrent blocks a second in-flight request", async () => {
+    const t = convexTest(schema, modules);
+    await setUserLimits(t, "u", { maxConcurrent: 1 });
+    const first = await start(t, { userId: "u" });
+    expect(first.allowed).toBe(true); // reserved, still pending
+    const second = await start(t, { userId: "u" });
+    expect(second.allowed).toBe(false);
+    expect(second.code).toBe("user_max_concurrent");
+  });
+});
+
+describe("per-bucket rate limits", () => {
+  test("the existing user rate limit remains compatible", async () => {
+    const t = convexTest(schema, modules);
+    await setUserLimits(t, "u", { requestsPerMinute: 1 });
+    const first = await start(t, { userId: "u" });
+    expect(first.allowed).toBe(true);
+    const second = await start(t, { userId: "u" });
+    expect(second.allowed).toBe(false);
+    expect(second.code).toBe("rate_limit");
+  });
+
+  test("an action rate limit blocks the next request for that action", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "action",
+      value: "ai:summarize",
+      requestsPerMinute: 1,
+    });
+    const first = await start(t, { userId: "u1", actionName: "ai:summarize" });
+    expect(first.allowed).toBe(true);
+    const second = await start(t, { userId: "u2", actionName: "ai:summarize" });
+    expect(second.allowed).toBe(false);
+    expect(second.code).toBe("action_rate_limit");
+  });
+
+  test("a custom-tag rate limit blocks the next request for that value", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "customer",
+      value: "acme",
+      requestsPerMinute: 1,
+    });
+    const tags = [{ dimension: "customer", value: "acme" }];
+    const first = await start(t, { userId: "u1", tags });
+    expect(first.allowed).toBe(true);
+    const second = await start(t, { userId: "u2", tags });
+    expect(second.allowed).toBe(false);
+    expect(second.code).toBe("customer_rate_limit");
+  });
+});
+
+describe("tag-filtered request log", () => {
+  test("listRequests filters by a custom tag dimension", async () => {
+    const t = convexTest(schema, modules);
+    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, {
+      dimension: "customer",
+      value: "acme",
+    });
+    expect(acme.length).toBe(1);
+    expect(acme[0].userId).toBe("u");
+  });
+});
+
+describe("tagged attribution buckets", () => {
+  test("a cap on a custom tag blocks, and settlement accrues to every bucket", async () => {
+    const t = convexTest(schema, modules);
+    // A tiny cap on customer "acme" — the user is uncapped.
+    await t.mutation(api.lib.setBucketLimits, {
+      dimension: "customer",
+      value: "acme",
+      dailySpendLimitNanos: 1_000,
+    });
+    const blocked = await start(t, {
+      userId: "u",
+      tags: [{ dimension: "customer", value: "acme" }],
+    });
+    expect(blocked.allowed).toBe(false);
+    expect(blocked.code).toBe("customer_daily_spend_limit");
+
+    // A different customer with no cap goes through, and the spend lands on
+    // BOTH the user bucket and the customer bucket.
+    const ok = await start(t, {
+      userId: "u",
+      tags: [{ dimension: "customer", value: "globex" }],
+    });
+    expect(ok.allowed).toBe(true);
+    await settle(t, ok.requestId, 1_000_000, 1_000_000); // $0.75
+    const user = await userOf(t, "u");
+    const cust = await bucketOf(t, "customer", "globex");
+    expect(user.totalSpendNanos).toBe(750_000_000);
+    expect(cust.totalSpendNanos).toBe(750_000_000);
+    expect(cust.totalRequests).toBe(1);
+  });
+});
+
 describe("D-00 exactly-once settlement", () => {
   test("a duplicate finishRequest does not double-count", async () => {
     const t = convexTest(schema, modules);
@@ -85,7 +324,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);
-    await t.mutation(api.lib.setLimits, { userId: "u", dailyTokenLimit: 10 });
+    await setUserLimits(t, "u", { dailyTokenLimit: 10 });
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(false);
     expect(r.code).toBe("user_daily_token_limit");
@@ -95,15 +334,14 @@ describe("token quotas", () => {
 describe("soft enforcement", () => {
   test("over a soft budget: allowed, warned, flagged overBudget", async () => {
     const t = convexTest(schema, modules);
-    await t.mutation(api.lib.setLimits, {
-      userId: "u",
+    await setUserLimits(t, "u", {
       dailySpendLimitNanos: 1, // 1 nanodollar — one estimate blows past it
       enforcement: "soft",
     });
     const r = await start(t, { userId: "u" });
     expect(r.allowed).toBe(true);
     expect(r.warnings.length).toBeGreaterThan(0);
-    const req = await t.query(api.lib.getRequest, { requestId: r.requestId });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
     expect(req.overBudget).toBe(true);
   });
 });
@@ -145,7 +383,7 @@ describe("F-04 fail-closed pricing", () => {
     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);
-    const req = await t.query(api.lib.getRequest, { requestId: r.requestId });
+    const req = (await t.query(api.lib.getRequest, { requestId: r.requestId }))!;
     expect(req.unpricedModel).toBe(true);
   });
 });