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

package/src/component/_generated/component.ts

@@ -0,0 +1,196 @@
+/* eslint-disable */
+/**
+ * Generated `ComponentApi` utility.
+ *
+ * THIS CODE IS AUTOMATICALLY GENERATED.
+ *
+ * To regenerate, run `npx convex dev`.
+ * @module
+ */
+
+import type { FunctionReference } from "convex/server";
+
+/**
+ * A utility for referencing a Convex component's exposed API.
+ *
+ * Useful when expecting a parameter like `components.myComponent`.
+ * Usage:
+ * ```ts
+ * async function myFunction(ctx: QueryCtx, component: ComponentApi) {
+ *   return ctx.runQuery(component.someFile.someQuery, { ...args });
+ * }
+ * ```
+ */
+export type ComponentApi<Name extends string | undefined = string | undefined> =
+  {
+    lib: {
+      bumpAction: FunctionReference<
+        "mutation",
+        "internal",
+        { dailyNanos?: number; lifetimeNanos?: number; name: string },
+        null,
+        Name
+      >;
+      bumpGlobal: FunctionReference<
+        "mutation",
+        "internal",
+        { dailyNanos?: number; lifetimeNanos?: number },
+        null,
+        Name
+      >;
+      bumpUser: FunctionReference<
+        "mutation",
+        "internal",
+        { dailyNanos?: number; lifetimeNanos?: number; userId: string },
+        null,
+        Name
+      >;
+      deleteUser: FunctionReference<
+        "mutation",
+        "internal",
+        { userId: string },
+        { deletedThisBatch: number; done: boolean },
+        Name
+      >;
+      finishRequest: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          cachedTokens?: number;
+          completionTokens?: number;
+          error?: string;
+          latencyMs?: number;
+          promptTokens?: number;
+          requestId: string;
+          responseText?: string;
+        },
+        { costNanos: number },
+        Name
+      >;
+      getGlobalStatus: FunctionReference<
+        "query",
+        "internal",
+        {},
+        {
+          dailySpendLimitNanos: number | null;
+          enforcement: "hard" | "soft";
+          lifetimeSpendLimitNanos: number | null;
+          spentTodayNanos: number;
+          spentTotalNanos: number;
+        },
+        Name
+      >;
+      getModelPolicy: FunctionReference<
+        "query",
+        "internal",
+        {},
+        { mode: "open" | "allowlist" | "denylist"; models: Array<string> },
+        Name
+      >;
+      getRequest: FunctionReference<
+        "query",
+        "internal",
+        { requestId: string },
+        any,
+        Name
+      >;
+      lineage: FunctionReference<
+        "query",
+        "internal",
+        { requestId: string },
+        any,
+        Name
+      >;
+      listActions: FunctionReference<"query", "internal", {}, any, Name>;
+      listPrices: FunctionReference<"query", "internal", {}, any, Name>;
+      listRequests: FunctionReference<
+        "query",
+        "internal",
+        { limit?: number; userId?: string },
+        any,
+        Name
+      >;
+      listUsers: FunctionReference<"query", "internal", {}, any, Name>;
+      setActionLimits: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          dailySpendLimitNanos?: number;
+          dailyTokenLimit?: number;
+          disabled?: boolean;
+          enforcement?: "hard" | "soft";
+          lifetimeSpendLimitNanos?: number;
+          lifetimeTokenLimit?: number;
+          name: string;
+        },
+        null,
+        Name
+      >;
+      setGlobalLimits: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          dailySpendLimitNanos?: number;
+          enforcement?: "hard" | "soft";
+          lifetimeSpendLimitNanos?: number;
+        },
+        null,
+        Name
+      >;
+      setLimits: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          blocked?: boolean;
+          dailySpendLimitNanos?: number;
+          dailyTokenLimit?: number;
+          enforcement?: "hard" | "soft";
+          lifetimeSpendLimitNanos?: number;
+          lifetimeTokenLimit?: number;
+          requestsPerMinute?: number;
+          userId: string;
+        },
+        null,
+        Name
+      >;
+      setModelPolicy: FunctionReference<
+        "mutation",
+        "internal",
+        { mode: "open" | "allowlist" | "denylist"; models: Array<string> },
+        null,
+        Name
+      >;
+      setPrice: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          inputNanosPerMTok: number;
+          model: string;
+          outputNanosPerMTok: number;
+        },
+        null,
+        Name
+      >;
+      setRetention: FunctionReference<
+        "mutation",
+        "internal",
+        { retentionMs: number },
+        null,
+        Name
+      >;
+      startRequest: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          actionName?: string;
+          messages: Array<{ content: string; role: string }>;
+          model: string;
+          rerunOf?: string;
+          userId: string;
+        },
+        | { allowed: true; requestId: string; warnings: Array<string> }
+        | { allowed: false; code: string; reason: string },
+        Name
+      >;
+    };
+  };

package/src/component/_generated/dataModel.ts

@@ -0,0 +1,60 @@
+/* eslint-disable */
+/**
+ * Generated data model types.
+ *
+ * THIS CODE IS AUTOMATICALLY GENERATED.
+ *
+ * To regenerate, run `npx convex dev`.
+ * @module
+ */
+
+import type {
+  DataModelFromSchemaDefinition,
+  DocumentByName,
+  TableNamesInDataModel,
+  SystemTableNames,
+} from "convex/server";
+import type { GenericId } from "convex/values";
+import schema from "../schema.js";
+
+/**
+ * The names of all of your Convex tables.
+ */
+export type TableNames = TableNamesInDataModel<DataModel>;
+
+/**
+ * The type of a document stored in Convex.
+ *
+ * @typeParam TableName - A string literal type of the table name (like "users").
+ */
+export type Doc<TableName extends TableNames> = DocumentByName<
+  DataModel,
+  TableName
+>;
+
+/**
+ * An identifier for a document in Convex.
+ *
+ * Convex documents are uniquely identified by their `Id`, which is accessible
+ * on the `_id` field. To learn more, see [Document IDs](https://docs.convex.dev/using/document-ids).
+ *
+ * Documents can be loaded using `db.get(tableName, id)` in query and mutation functions.
+ *
+ * IDs are just strings at runtime, but this type can be used to distinguish them from other
+ * strings when type checking.
+ *
+ * @typeParam TableName - A string literal type of the table name (like "users").
+ */
+export type Id<TableName extends TableNames | SystemTableNames> =
+  GenericId<TableName>;
+
+/**
+ * A type describing your Convex data model.
+ *
+ * This type includes information about what tables you have, the type of
+ * documents stored in those tables, and the indexes defined on them.
+ *
+ * This type is used to parameterize methods like `queryGeneric` and
+ * `mutationGeneric` to make them type-safe.
+ */
+export type DataModel = DataModelFromSchemaDefinition<typeof schema>;

package/src/component/_generated/server.ts

@@ -0,0 +1,169 @@
+/* eslint-disable */
+/**
+ * Generated utilities for implementing server-side Convex query and mutation functions.
+ *
+ * THIS CODE IS AUTOMATICALLY GENERATED.
+ *
+ * To regenerate, run `npx convex dev`.
+ * @module
+ */
+
+import type {
+  ActionBuilder,
+  HttpActionBuilder,
+  MutationBuilder,
+  QueryBuilder,
+  GenericActionCtx,
+  GenericMutationCtx,
+  GenericQueryCtx,
+  GenericDatabaseReader,
+  GenericDatabaseWriter,
+} from "convex/server";
+import {
+  actionGeneric,
+  httpActionGeneric,
+  queryGeneric,
+  mutationGeneric,
+  internalActionGeneric,
+  internalMutationGeneric,
+  internalQueryGeneric,
+} from "convex/server";
+import type { DataModel } from "./dataModel.js";
+
+/**
+ * Typesafe environment variables.
+ *
+ * This includes platform-provided env vars and any variables declared in
+ * `convex.config.ts`.
+ */
+type Env = {
+  readonly CONVEX_CLOUD_URL: string;
+  readonly CONVEX_SITE_URL: string;
+};
+
+/**
+ * Define a query in this Convex app's public API.
+ *
+ * This function will be allowed to read your Convex database and will be accessible from the client.
+ *
+ * @param func - The query function. It receives a {@link QueryCtx} as its first argument.
+ * @returns The wrapped query. Include this as an `export` to name it and make it accessible.
+ */
+export const query: QueryBuilder<DataModel, "public"> = queryGeneric;
+
+/**
+ * Define a query that is only accessible from other Convex functions (but not from the client).
+ *
+ * This function will be allowed to read from your Convex database. It will not be accessible from the client.
+ *
+ * @param func - The query function. It receives a {@link QueryCtx} as its first argument.
+ * @returns The wrapped query. Include this as an `export` to name it and make it accessible.
+ */
+export const internalQuery: QueryBuilder<DataModel, "internal"> =
+  internalQueryGeneric;
+
+/**
+ * Define a mutation in this Convex app's public API.
+ *
+ * This function will be allowed to modify your Convex database and will be accessible from the client.
+ *
+ * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument.
+ * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible.
+ */
+export const mutation: MutationBuilder<DataModel, "public"> = mutationGeneric;
+
+/**
+ * Define a mutation that is only accessible from other Convex functions (but not from the client).
+ *
+ * This function will be allowed to modify your Convex database. It will not be accessible from the client.
+ *
+ * @param func - The mutation function. It receives a {@link MutationCtx} as its first argument.
+ * @returns The wrapped mutation. Include this as an `export` to name it and make it accessible.
+ */
+export const internalMutation: MutationBuilder<DataModel, "internal"> =
+  internalMutationGeneric;
+
+/**
+ * Define an action in this Convex app's public API.
+ *
+ * An action is a function which can execute any JavaScript code, including non-deterministic
+ * code and code with side-effects, like calling third-party services.
+ * They can be run in Convex's JavaScript environment or in Node.js using the "use node" directive.
+ * They can interact with the database indirectly by calling queries and mutations using the {@link ActionCtx}.
+ *
+ * @param func - The action. It receives an {@link ActionCtx} as its first argument.
+ * @returns The wrapped action. Include this as an `export` to name it and make it accessible.
+ */
+export const action: ActionBuilder<DataModel, "public"> = actionGeneric;
+
+/**
+ * Define an action that is only accessible from other Convex functions (but not from the client).
+ *
+ * @param func - The function. It receives an {@link ActionCtx} as its first argument.
+ * @returns The wrapped function. Include this as an `export` to name it and make it accessible.
+ */
+export const internalAction: ActionBuilder<DataModel, "internal"> =
+  internalActionGeneric;
+
+/**
+ * Define an HTTP action.
+ *
+ * The wrapped function will be used to respond to HTTP requests received
+ * by a Convex deployment if the requests matches the path and method where
+ * this action is routed. Be sure to route your httpAction in `convex/http.js`.
+ *
+ * @param func - The function. It receives an {@link ActionCtx} as its first argument
+ * and a Fetch API `Request` object as its second.
+ * @returns The wrapped function. Import this function from `convex/http.js` and route it to hook it up.
+ */
+export const httpAction: HttpActionBuilder = httpActionGeneric;
+export const env: Env = (globalThis as unknown as { process: { env: Env } })
+  .process.env;
+
+/**
+ * A set of services for use within Convex query functions.
+ *
+ * The query context is passed as the first argument to any Convex query
+ * function run on the server.
+ *
+ * If you're using code generation, use the `QueryCtx` type in `convex/_generated/server.d.ts` instead.
+ */
+export type QueryCtx = GenericQueryCtx<DataModel>;
+
+/**
+ * A set of services for use within Convex mutation functions.
+ *
+ * The mutation context is passed as the first argument to any Convex mutation
+ * function run on the server.
+ *
+ * If you're using code generation, use the `MutationCtx` type in `convex/_generated/server.d.ts` instead.
+ */
+export type MutationCtx = GenericMutationCtx<DataModel>;
+
+/**
+ * A set of services for use within Convex action functions.
+ *
+ * The action context is passed as the first argument to any Convex action
+ * function run on the server.
+ */
+export type ActionCtx = GenericActionCtx<DataModel>;
+
+/**
+ * An interface to read from the database within Convex query functions.
+ *
+ * The two entry points are {@link DatabaseReader.get}, which fetches a single
+ * document by its {@link Id}, or {@link DatabaseReader.query}, which starts
+ * building a query.
+ */
+export type DatabaseReader = GenericDatabaseReader<DataModel>;
+
+/**
+ * An interface to read from and write to the database within Convex mutation
+ * functions.
+ *
+ * Convex guarantees that all writes within a single mutation are
+ * executed atomically, so you never have to worry about partial writes leaving
+ * your data in an inconsistent state. See [the Convex Guide](https://docs.convex.dev/understanding/convex-fundamentals/functions#atomicity-and-optimistic-concurrency-control)
+ * for the guarantees Convex provides your functions.
+ */
+export type DatabaseWriter = GenericDatabaseWriter<DataModel>;

package/src/component/convex.config.ts

@@ -0,0 +1,7 @@
+import { defineComponent } from "convex/server";
+import shardedCounter from "@convex-dev/sharded-counter/convex.config";
+
+const component = defineComponent("aiBudget");
+// Global spend totals use a sharded counter for high write throughput.
+component.use(shardedCounter);
+export default component;

package/src/component/crons.ts

@@ -0,0 +1,15 @@
+import { cronJobs } from "convex/server";
+import { internal } from "./_generated/api";
+
+const crons = cronJobs();
+
+// Backstop reservation/settlement: fold any finished-but-unfolded requests and
+// release reservations for requests that never settled.
+crons.interval(
+  "reconcile ai gateway spend",
+  { minutes: 1 },
+  internal.lib.reconcile,
+  {}
+);
+
+export default crons;

package/src/component/lib.test.ts

@@ -0,0 +1,151 @@
+import { convexTest } from "convex-test";
+import { describe, expect, test, vi } from "vitest";
+import schema from "./schema";
+import { api } from "./_generated/api";
+
+// convex-test loads the component's own modules; exclude convex.config (not a
+// function module) and the test files themselves.
+const modules = import.meta.glob([
+  "./**/*.ts",
+  "!./**/*.test.ts",
+  "!./**/convex.config.ts",
+]);
+
+const MODEL = "openai/gpt-4o-mini";
+const msg = (content: string) => [{ role: "user", content }];
+
+async function start(t: any, args: any) {
+  return t.mutation(api.lib.startRequest, {
+    model: MODEL,
+    messages: msg("hi"),
+    ...args,
+  });
+}
+async function settle(t: any, requestId: any, p = 10, c = 5) {
+  await t.mutation(api.lib.finishRequest, {
+    requestId,
+    promptTokens: p,
+    completionTokens: c,
+  });
+  // finishRequest schedules the fold via runAfter(0); drain it.
+  vi.useFakeTimers();
+  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);
+
+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,
+    });
+    const r = await start(t, { userId: "u" });
+    expect(r.allowed).toBe(false);
+    expect(r.code).toBe("user_daily_spend_limit");
+  });
+
+  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
+    });
+    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
+    const u = await userOf(t, "u");
+    // gpt-4o-mini: $0.15/Mtok in + $0.60/Mtok out = $0.75 = 750_000_000 nano.
+    expect(u.totalSpendNanos).toBe(750_000_000);
+    expect(u.reservedTotalNanos ?? 0).toBe(0);
+    expect(u.pendingCount ?? 0).toBe(0);
+    expect(u.totalRequests).toBe(1);
+  });
+});
+
+describe("D-00 exactly-once settlement", () => {
+  test("a duplicate finishRequest does not double-count", async () => {
+    const t = convexTest(schema, modules);
+    const r = await start(t, { userId: "u" });
+    await settle(t, r.requestId); // first settle
+    const before = await userOf(t, "u");
+    await settle(t, r.requestId); // late duplicate — must be a no-op
+    const after = await userOf(t, "u");
+    expect(after.totalRequests).toBe(1);
+    expect(after.totalRequests).toBe(before.totalRequests);
+    expect(after.totalSpendNanos).toBeCloseTo(before.totalSpendNanos, 9);
+    expect(after.reservedTotalNanos ?? 0).toBeCloseTo(0, 9);
+  });
+});
+
+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 });
+    const r = await start(t, { userId: "u" });
+    expect(r.allowed).toBe(false);
+    expect(r.code).toBe("user_daily_token_limit");
+  });
+});
+
+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",
+      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 });
+    expect(req.overBudget).toBe(true);
+  });
+});
+
+describe("model policy", () => {
+  test("allowlist blocks an off-list model", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.lib.setModelPolicy, {
+      mode: "allowlist",
+      models: ["openai/gpt-4o-mini"],
+    });
+    const blocked = await start(t, { userId: "u", model: "openai/gpt-4o" });
+    expect(blocked.allowed).toBe(false);
+    expect(blocked.code).toBe("model_not_allowed");
+    const ok = await start(t, { userId: "u", model: "openai/gpt-4o-mini" });
+    expect(ok.allowed).toBe(true);
+  });
+});
+
+describe("D-02 pricing validation", () => {
+  test("setPrice rejects negative rates", async () => {
+    const t = convexTest(schema, modules);
+    await expect(
+      t.mutation(api.lib.setPrice, {
+        model: "x/y",
+        inputNanosPerMTok: -1,
+        outputNanosPerMTok: 5,
+      })
+    ).rejects.toThrow(/non-negative/);
+  });
+});
+
+describe("F-04 fail-closed pricing", () => {
+  test("an unknown model is charged the conservative max, not zero", async () => {
+    const t = convexTest(schema, modules);
+    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);
+    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 });
+    expect(req.unpricedModel).toBe(true);
+  });
+});