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

package/dist/client/index.js

@@ -1,6 +1,8 @@
+import { httpActionGeneric, } from "convex/server";
 import { ConvexError } from "convex/values";
 import { generateText, wrapLanguageModel } from "ai";
 import { convexGateway } from "@convex-dev/ai-sdk-provider";
+import { DASHBOARD_HTML } from "./dashboard";
 // The calling Convex action's name (e.g. "ai:sendMessage"), unless overridden.
 async function resolveActionName(ctx, explicit) {
     if (explicit !== undefined)
@@ -39,14 +41,30 @@ function extractUsage(usage) {
     return {
         promptTokens: toTokenCount(usage?.inputTokens ?? usage?.promptTokens),
         completionTokens: toTokenCount(usage?.outputTokens ?? usage?.completionTokens),
-        // cached prompt tokens: AI SDK v5 `cachedInputTokens`, OpenAI-compat
-        // `prompt_tokens_details.cached_tokens` / `cached_tokens`.
-        cachedTokens: toTokenCount(usage?.cachedInputTokens ??
+        // cached prompt tokens. The Convex gateway reports these at
+        // `usage.inputTokenDetails.cacheReadTokens`; the other paths cover AI SDK v5
+        // (`cachedInputTokens`) and raw OpenAI-compatible shapes.
+        cachedTokens: toTokenCount(usage?.inputTokenDetails?.cacheReadTokens ??
+            usage?.cachedInputTokens ??
             usage?.promptTokensDetails?.cachedTokens ??
             usage?.prompt_tokens_details?.cached_tokens ??
             usage?.cached_tokens),
     };
 }
+// The gateway doesn't report a dollar cost today, but if a future response ever
+// carries an unambiguous nanodollar cost we pass it straight through as
+// authoritative (finishRequest prefers it over the token-based estimate). Only
+// an explicitly nano-denominated field is trusted — a bare `cost` could be in
+// dollars and silently mis-bill by 1e9×.
+function extractGatewayCostNanos(result) {
+    const meta = result?.providerMetadata?.convexGateway;
+    const candidates = [meta?.costNanos, result?.usage?.costNanos];
+    for (const c of candidates) {
+        if (typeof c === "number" && Number.isFinite(c) && c >= 0)
+            return c;
+    }
+    return undefined;
+}
 // Flatten an AI SDK prompt (roles + content parts) into simple storable messages.
 function simplifyPrompt(prompt) {
     if (!Array.isArray(prompt))
@@ -78,24 +96,46 @@ function extractText(result) {
     }
     return "";
 }
-// ---------- client ----------
 export class AIBudget {
     component;
     defaultModel;
     onSoftLimit;
+    onThreshold;
+    onLimitReached;
     constructor(component, options) {
         this.component = component;
         this.defaultModel = options?.defaultModel ?? "openai/gpt-4o-mini";
         this.onSoftLimit = options?.onSoftLimit;
+        this.onThreshold = options?.onThreshold;
+        this.onLimitReached = options?.onLimitReached;
+    }
+    // Fire the soft-limit + threshold callbacks from a startRequest result.
+    async fireBudgetEvents(base, warnings, notices) {
+        if (warnings.length > 0 && this.onSoftLimit) {
+            try {
+                await this.onSoftLimit({ ...base, messages: warnings, warnings });
+            }
+            catch {
+                /* never let a callback error break a request */
+            }
+        }
+        if (notices.length > 0 && this.onThreshold) {
+            try {
+                await this.onThreshold({ ...base, messages: notices });
+            }
+            catch {
+                /* swallow */
+            }
+        }
     }
-    async fireSoftLimit(info) {
-        if (info.warnings.length === 0 || !this.onSoftLimit)
+    async fireLimitReached(info) {
+        if (!this.onLimitReached)
             return;
         try {
-            await this.onSoftLimit(info);
+            await this.onLimitReached(info);
         }
         catch {
-            // never let a callback error break a request
+            /* swallow */
         }
     }
     /**
@@ -116,6 +156,14 @@ export class AIBudget {
             rerunOf: args.rerunOf,
         });
         if (!started.allowed) {
+            await this.fireLimitReached({
+                userId,
+                action: actionName,
+                tags: args.tags,
+                messages: [started.reason],
+                code: started.code,
+                reason: started.reason,
+            });
             throw new ConvexError({
                 kind: "AIBudgetLimit",
                 code: started.code,
@@ -124,13 +172,8 @@ export class AIBudget {
         }
         const requestId = started.requestId;
         const warnings = started.warnings;
-        await this.fireSoftLimit({
-            userId,
-            action: actionName,
-            tags: args.tags,
-            requestId,
-            warnings,
-        });
+        const notices = started.notices;
+        await this.fireBudgetEvents({ userId, action: actionName, tags: args.tags, requestId }, warnings, notices);
         const start = Date.now();
         try {
             // The full chain (incl. system) is stored on the request for audit/replay,
@@ -150,9 +193,10 @@ export class AIBudget {
                 requestId,
                 responseText: result.text,
                 ...usage,
+                costNanos: extractGatewayCostNanos(result),
                 latencyMs: Date.now() - start,
             });
-            return { text: result.text, requestId, costNanos, warnings, ...usage };
+            return { text: result.text, requestId, costNanos, warnings, notices, ...usage };
         }
         catch (e) {
             await ctx.runMutation(this.component.lib.finishRequest, {
@@ -172,10 +216,12 @@ export class AIBudget {
     languageModel(ctx, opts = {}) {
         const modelId = opts.model ?? this.defaultModel;
         const component = this.component;
-        const fireSoftLimit = this.fireSoftLimit.bind(this);
+        const fireBudgetEvents = this.fireBudgetEvents.bind(this);
+        const fireLimitReached = this.fireLimitReached.bind(this);
         const begin = async (params) => {
             const userId = await resolveUserId(ctx, opts.userId);
             const actionName = await resolveActionName(ctx, opts.action);
+            const base = { userId, action: actionName, tags: opts.tags };
             const started = await ctx.runMutation(component.lib.startRequest, {
                 userId,
                 actionName,
@@ -184,19 +230,19 @@ export class AIBudget {
                 messages: simplifyPrompt(params.prompt),
             });
             if (!started.allowed) {
+                await fireLimitReached({
+                    ...base,
+                    messages: [started.reason],
+                    code: started.code,
+                    reason: started.reason,
+                });
                 throw new ConvexError({
                     kind: "AIBudgetLimit",
                     code: started.code,
                     reason: started.reason,
                 });
             }
-            await fireSoftLimit({
-                userId,
-                action: actionName,
-                tags: opts.tags,
-                requestId: started.requestId,
-                warnings: started.warnings,
-            });
+            await fireBudgetEvents({ ...base, requestId: started.requestId }, started.warnings, started.notices);
             return started.requestId;
         };
         const finish = async (requestId, fields) => ctx.runMutation(component.lib.finishRequest, { requestId, ...fields });
@@ -211,6 +257,7 @@ export class AIBudget {
                         await finish(requestId, {
                             responseText: extractText(result),
                             ...extractUsage(result.usage),
+                            costNanos: extractGatewayCostNanos(result),
                             latencyMs: Date.now() - start,
                         });
                         return result;
@@ -295,6 +342,7 @@ export class AIBudget {
     get requests() {
         const c = this.component;
         return {
+            /** Filter by userId, or by any {dimension, value} (incl. custom tags). */
             list: (ctx, args = {}) => ctx.runQuery(c.lib.listRequests, args),
             /** Ancestors up to the original, plus direct re-runs. */
             lineage: (ctx, args) => ctx.runQuery(c.lib.lineage, { requestId: args.requestId }),
@@ -307,74 +355,73 @@ export class AIBudget {
      * generalization of `users`/`actions`. Give it any dimension name (team,
      * project, tenant, customer, env, feature, …) and set caps per value:
      *
-     *   ai.tag("customer").setLimits(ctx, { value: "acme", dailySpendLimitNanos });
-     *   ai.tag("customer").list(ctx);
+     *   ai.tag("customer").setLimits(ctx, { value: "acme", monthlySpendLimitNanos });
+     *   ai.tag("customer").history(ctx, { value: "acme", period: "day" });
      *
      * Attribute a call to it by passing `tags` to `chat`/`languageModel`.
      */
     tag(dimension) {
+        return this.dimensionApi(dimension, (a) => a.value);
+    }
+    // Shared implementation behind tag()/users/actions. `key` maps the namespace's
+    // id field (value/userId/name) to the bucket value.
+    dimensionApi(dimension, key) {
         const c = this.component;
         return {
             /** All buckets in this dimension. */
             list: (ctx) => ctx.runQuery(c.lib.listBuckets, { dimension }),
             /** One bucket's limits + spend (null if it has none yet). */
-            get: (ctx, args) => ctx.runQuery(c.lib.getBucket, { dimension, value: args.value }),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setBucketLimits, { dimension, ...args }),
-            /** One-time "approve another $X" bump (daily is today-only). */
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, { dimension, ...args }),
-            /** Delete the bucket (for "user", also its request rows). */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, { dimension, value: args.value }),
-        };
-    }
-    /** Per-user budgets and controls — sugar over the "user" dimension. */
-    get users() {
-        const c = this.component;
-        return {
-            list: (ctx) => ctx.runQuery(c.lib.listBuckets, { dimension: "user" }),
+            get: (ctx, args) => ctx.runQuery(c.lib.getBucket, { dimension, value: key(args) }),
             setLimits: (ctx, args) => {
-                const { userId, ...limits } = args;
+                const { value, userId, name, ...limits } = args;
                 return ctx.runMutation(c.lib.setBucketLimits, {
-                    dimension: "user",
-                    value: userId,
+                    dimension,
+                    value: key(args),
                     ...limits,
                 });
             },
-            /** One-time "approve another $X" bump (daily is today-only). */
+            /** One-time "approve another $X" bump (daily/monthly reset with the window). */
             bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
-                dimension: "user",
-                value: args.userId,
+                dimension,
+                value: key(args),
                 dailyNanos: args.dailyNanos,
+                monthlyNanos: args.monthlyNanos,
                 lifetimeNanos: args.lifetimeNanos,
             }),
-            /** Delete a user and all their request rows. */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, {
-                dimension: "user",
-                value: args.userId,
+            /** Manually credit (negative) or debit (positive) this bucket. */
+            adjust: (ctx, args) => ctx.runMutation(c.lib.adjustBucket, {
+                dimension,
+                value: key(args),
+                deltaNanos: args.deltaNanos,
+                tokens: args.tokens,
+                reason: args.reason,
+            }),
+            /** Durable spend history for this bucket (per day or per month). */
+            history: (ctx, args) => ctx.runQuery(c.lib.usageHistory, {
+                dimension,
+                value: key(args),
+                period: args.period ?? "day",
+                limit: args.limit,
             }),
+            /** Manual-adjustment audit log for this bucket. */
+            adjustments: (ctx, args) => ctx.runQuery(c.lib.listAdjustments, {
+                dimension,
+                value: key(args),
+                limit: args.limit,
+            }),
+            /** Delete the bucket (for "user", also its request rows). */
+            delete: (ctx, args) => ctx.runMutation(c.lib.deleteBucket, { dimension, value: key(args) }),
         };
     }
+    /** Per-user budgets and controls — sugar over the "user" dimension. */
+    get users() {
+        return this.dimensionApi("user", (a) => a.userId);
+    }
     /** Per-action (per-feature) budgets — sugar over the "action" dimension. */
     get actions() {
-        const c = this.component;
-        return {
-            list: (ctx) => ctx.runQuery(c.lib.listBuckets, { dimension: "action" }),
-            setLimits: (ctx, args) => {
-                const { name, ...limits } = args;
-                return ctx.runMutation(c.lib.setBucketLimits, {
-                    dimension: "action",
-                    value: name,
-                    ...limits,
-                });
-            },
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
-                dimension: "action",
-                value: args.name,
-                dailyNanos: args.dailyNanos,
-                lifetimeNanos: args.lifetimeNanos,
-            }),
-        };
+        return this.dimensionApi("action", (a) => a.name);
     }
-    /** The deployment-wide budget and retention config. */
+    /** The deployment-wide budget, alerts, and retention config. */
     get global() {
         const c = this.component;
         return {
@@ -383,6 +430,8 @@ export class AIBudget {
             /** A killswitch spend cap across all users/actions (enforced approximately). */
             setLimits: (ctx, args) => ctx.runMutation(c.lib.setGlobalLimits, args),
             bump: (ctx, args) => ctx.runMutation(c.lib.bumpGlobal, args),
+            /** Default approaching-limit alert threshold (fraction of a cap, e.g. 0.8). */
+            setAlertDefaults: (ctx, args) => ctx.runMutation(c.lib.setAlertDefaults, args),
             /** Request-row retention window in ms (default 1h; 0 disables). */
             setRetention: (ctx, args) => ctx.runMutation(c.lib.setRetention, args),
         };
@@ -404,6 +453,111 @@ export class AIBudget {
             set: (ctx, args) => ctx.runMutation(c.lib.setPrice, args),
         };
     }
+    /**
+     * Mount the built-in admin dashboard on your app's HTTP router with one call.
+     * Serves a self-contained HTML dashboard (buckets, requests, usage history,
+     * settings) plus a small JSON API, all backed by the component — no extra
+     * queries to write.
+     *
+     *   // convex/http.ts
+     *   import { httpRouter } from "convex/server";
+     *   const http = httpRouter();
+     *   ai.registerRoutes(http, { authorize: async (ctx) =>
+     *     (await ctx.auth.getUserIdentity())?.role === "admin" });
+     *   export default http;
+     *
+     * It then lives at `https://<deployment>.convex.site/aibudget`.
+     *
+     * SECURITY: the endpoint is public on the internet. You MUST gate it — either
+     * pass `authorize` (recommended: check the caller is a deployment admin) or
+     * set the `AI_BUDGET_DASHBOARD_TOKEN` env var (a bearer token / `?token=`).
+     * With neither, every route returns 401.
+     */
+    registerRoutes(http, opts = {}) {
+        const prefix = (opts.path ?? "/aibudget").replace(/\/+$/, "");
+        const c = this.component.lib;
+        const authorize = opts.authorize;
+        const guard = async (ctx, request) => {
+            if (authorize)
+                return { ok: await authorize(ctx, request), token: "" };
+            const token = globalThis.process?.env?.AI_BUDGET_DASHBOARD_TOKEN;
+            if (!token)
+                return { ok: false, token: "" };
+            const url = new URL(request.url);
+            const bearer = (request.headers.get("authorization") ?? "").replace(/^Bearer\s+/i, "");
+            const provided = bearer || url.searchParams.get("token") || "";
+            return { ok: provided === token, token };
+        };
+        const json = (data, status = 200) => new Response(JSON.stringify(data ?? null), {
+            status,
+            headers: { "content-type": "application/json" },
+        });
+        const handle = async (ctx, request) => {
+            const url = new URL(request.url);
+            const sub = url.pathname.slice(prefix.length) || "/";
+            const { ok, token } = await guard(ctx, request);
+            if (!ok) {
+                return new Response("Unauthorized. Pass `authorize` to registerRoutes or set AI_BUDGET_DASHBOARD_TOKEN.", { status: 401 });
+            }
+            if (sub.startsWith("/api/")) {
+                const route = request.method + " " + sub.slice(4); // strip "/api"
+                const p = {};
+                url.searchParams.forEach((v, k) => {
+                    p[k] = v;
+                });
+                const body = request.method === "POST"
+                    ? await request.json().catch(() => ({}))
+                    : {};
+                switch (route) {
+                    case "GET /buckets":
+                        return json(await ctx.runQuery(c.listBuckets, { dimension: p.dimension || undefined }));
+                    case "GET /requests":
+                        return json(await ctx.runQuery(c.listRequests, {
+                            userId: p.userId || undefined,
+                            dimension: p.dimension || undefined,
+                            value: p.value || undefined,
+                            limit: 100,
+                        }));
+                    case "GET /usage":
+                        return json(await ctx.runQuery(c.usageHistory, {
+                            dimension: p.dimension,
+                            value: p.value,
+                            period: p.period === "month" ? "month" : "day",
+                        }));
+                    case "GET /global":
+                        return json(await ctx.runQuery(c.getGlobalStatus, {}));
+                    case "GET /prices":
+                        return json(await ctx.runQuery(c.listPrices, {}));
+                    case "POST /setLimits":
+                        return json(await ctx.runMutation(c.setBucketLimits, body));
+                    case "POST /bump":
+                        return json(await ctx.runMutation(c.bumpBucket, body));
+                    case "POST /adjust":
+                        return json(await ctx.runMutation(c.adjustBucket, body));
+                    case "POST /delete":
+                        return json(await ctx.runMutation(c.deleteBucket, body));
+                    case "POST /global/setLimits":
+                        return json(await ctx.runMutation(c.setGlobalLimits, body));
+                    case "POST /global/setAlertDefaults":
+                        return json(await ctx.runMutation(c.setAlertDefaults, body));
+                    case "POST /global/setRetention":
+                        return json(await ctx.runMutation(c.setRetention, body));
+                    case "POST /setPrice":
+                        return json(await ctx.runMutation(c.setPrice, body));
+                    default:
+                        return json({ error: "not found" }, 404);
+                }
+            }
+            const html = DASHBOARD_HTML.replace(/__API_BASE__/g, `${prefix}/api`).replace(/__TOKEN__/g, token);
+            return new Response(html, {
+                headers: { "content-type": "text/html; charset=utf-8" },
+            });
+        };
+        const handler = httpActionGeneric(handle);
+        http.route({ path: prefix, method: "GET", handler });
+        http.route({ pathPrefix: `${prefix}/`, method: "GET", handler });
+        http.route({ pathPrefix: `${prefix}/`, method: "POST", handler });
+    }
 }
 /** @deprecated Renamed to `AIBudget`. */
 export const WorryFreeAI = AIBudget;

package/dist/component/_generated/component.d.ts

@@ -20,15 +20,24 @@ 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", {
             dailyNanos?: number;
             dimension: string;
             lifetimeNanos?: number;
+            monthlyNanos?: number;
             value: string;
         }, null, Name>;
         bumpGlobal: FunctionReference<"mutation", "internal", {
             dailyNanos?: number;
             lifetimeNanos?: number;
+            monthlyNanos?: number;
         }, null, Name>;
         deleteBucket: FunctionReference<"mutation", "internal", {
             dimension: string;
@@ -40,6 +49,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         finishRequest: FunctionReference<"mutation", "internal", {
             cachedTokens?: number;
             completionTokens?: number;
+            costNanos?: number;
             error?: string;
             latencyMs?: number;
             promptTokens?: number;
@@ -54,8 +64,10 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         }, 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;
         }, Name>;
@@ -69,14 +81,24 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
         lineage: FunctionReference<"query", "internal", {
             requestId: string;
         }, 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", {
+            dimension?: string;
             limit?: number;
             userId?: string;
+            value?: string;
         }, any, Name>;
+        setAlertDefaults: FunctionReference<"mutation", "internal", {
+            warnAtPct?: number;
+        }, null, Name>;
         setBucketLimits: FunctionReference<"mutation", "internal", {
             blocked?: boolean;
             dailySpendLimitNanos?: number;
@@ -85,8 +107,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>;
         setGlobalLimits: FunctionReference<"mutation", "internal", {
             dailySpendLimitNanos?: number;
@@ -98,6 +124,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             models: Array<string>;
         }, null, Name>;
         setPrice: FunctionReference<"mutation", "internal", {
+            cachedNanosPerMTok?: number;
             inputNanosPerMTok: number;
             model: string;
             outputNanosPerMTok: number;
@@ -120,6 +147,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             userId: string;
         }, {
             allowed: true;
+            notices: Array<string>;
             requestId: string;
             warnings: Array<string>;
         } | {
@@ -127,5 +155,11 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             code: string;
             reason: string;
         }, Name>;
+        usageHistory: FunctionReference<"query", "internal", {
+            dimension: string;
+            limit?: number;
+            period: "day" | "month";
+            value: string;
+        }, any, Name>;
     };
 };