@convex-dev/ai-budget 0.0.2-alpha.4 → 0.0.2-alpha.6

package/src/component/_generated/component.ts

@@ -24,6 +24,19 @@ import type { FunctionReference } from "convex/server";
 export type ComponentApi<Name extends string | undefined = string | undefined> =
   {
     lib: {
+      adjustBucket: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          deltaNanos: number;
+          dimension: string;
+          reason?: string;
+          tokens?: number;
+          value: string;
+        },
+        null,
+        Name
+      >;
       bumpBucket: FunctionReference<
         "mutation",
         "internal",
@@ -31,6 +44,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           dailyNanos?: number;
           dimension: string;
           lifetimeNanos?: number;
+          monthlyNanos?: number;
           value: string;
         },
         null,
@@ -39,7 +53,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
       bumpGlobal: FunctionReference<
         "mutation",
         "internal",
-        { dailyNanos?: number; lifetimeNanos?: number },
+        { dailyNanos?: number; lifetimeNanos?: number; monthlyNanos?: number },
         null,
         Name
       >;
@@ -56,6 +70,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         {
           cachedTokens?: number;
           completionTokens?: number;
+          costNanos?: number;
           error?: string;
           latencyMs?: number;
           promptTokens?: number;
@@ -78,8 +93,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         {},
         {
           dailySpendLimitNanos: number | null;
+          defaultWarnAtPct: number | null;
           enforcement: "hard" | "soft";
           lifetimeSpendLimitNanos: number | null;
+          retentionMs: number | null;
           spentTodayNanos: number;
           spentTotalNanos: number;
         },
@@ -106,6 +123,13 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         any,
         Name
       >;
+      listAdjustments: FunctionReference<
+        "query",
+        "internal",
+        { dimension: string; limit?: number; value: string },
+        any,
+        Name
+      >;
       listBuckets: FunctionReference<
         "query",
         "internal",
@@ -117,10 +141,17 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
       listRequests: FunctionReference<
         "query",
         "internal",
-        { limit?: number; userId?: string },
+        { dimension?: string; limit?: number; userId?: string; value?: string },
         any,
         Name
       >;
+      setAlertDefaults: FunctionReference<
+        "mutation",
+        "internal",
+        { warnAtPct?: number },
+        null,
+        Name
+      >;
       setBucketLimits: FunctionReference<
         "mutation",
         "internal",
@@ -132,8 +163,12 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           enforcement?: "hard" | "soft";
           lifetimeSpendLimitNanos?: number;
           lifetimeTokenLimit?: number;
+          maxConcurrent?: number;
+          monthlySpendLimitNanos?: number;
+          monthlyTokenLimit?: number;
           requestsPerMinute?: number;
           value: string;
+          warnAtPct?: number;
         },
         null,
         Name
@@ -160,6 +195,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         "mutation",
         "internal",
         {
+          cachedNanosPerMTok?: number;
           inputNanosPerMTok: number;
           model: string;
           outputNanosPerMTok: number;
@@ -185,9 +221,26 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           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

@@ -70,6 +70,124 @@ 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("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("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);