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

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,35 @@ 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 AI Gateway reports the authoritative dollar cost of each request.
+// @convex-dev/ai-sdk-provider surfaces it at
+// `providerMetadata.convexGateway.cost` (USD); convert to nanodollars and pass
+// it through as authoritative (finishRequest prefers it over the token-based
+// estimate). No-op on older provider versions that don't surface it.
+function extractGatewayCostNanos(result) {
+    const meta = result?.providerMetadata?.convexGateway;
+    const costUsd = meta?.cost;
+    if (typeof costUsd === "number" && Number.isFinite(costUsd) && costUsd >= 0) {
+        return Math.round(costUsd * NANOS_PER_DOLLAR);
+    }
+    // Back-compat: honor an explicitly nano-denominated field if ever present.
+    const costNanos = meta?.costNanos ?? result?.usage?.costNanos;
+    if (typeof costNanos === "number" && Number.isFinite(costNanos) && costNanos >= 0) {
+        return Math.round(costNanos);
+    }
+    return undefined;
+}
+const NANOS_PER_DOLLAR = 1e9;
 // Flatten an AI SDK prompt (roles + content parts) into simple storable messages.
 function simplifyPrompt(prompt) {
     if (!Array.isArray(prompt))
@@ -79,42 +102,89 @@ function extractText(result) {
     return "";
 }
 // ---------- client ----------
+// Length-independent-branch string compare, so the dashboard token check
+// doesn't leak the token via response timing. (Length itself is not secret.)
+function timingSafeEqual(a, b) {
+    if (a.length !== b.length)
+        return false;
+    let r = 0;
+    for (let i = 0; i < a.length; i++)
+        r |= a.charCodeAt(i) ^ b.charCodeAt(i);
+    return r === 0;
+}
 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 */
         }
     }
     /**
-     * One-shot chat through the AI Gateway with tracking + limits.
-     * Call from an action. `userId` defaults to the authenticated caller.
+     * Meter ANY LLM call — gateway, a provider SDK, a raw fetch — with the same
+     * budgets, audit log, and cost tracking. Reserves before your `run` (throwing
+     * a ConvexError over a hard cap), runs it, then records the actual usage/cost.
+     * This is the provider-agnostic core; `chat` is sugar over it for the gateway.
+     *
+     * `run` returns what happened. Pass a raw provider `usage` object (auto-
+     * normalized) OR explicit `promptTokens`/`completionTokens`/`cachedTokens`,
+     * plus optional `serverToolUses` (e.g. `{ web_search: 3 }`, priced on top of
+     * tokens) and an authoritative `costNanos` (used verbatim if present).
      */
-    async chat(ctx, args = {}) {
-        const model = args.model ?? this.defaultModel;
-        const userId = await resolveUserId(ctx, args.userId);
-        const actionName = await resolveActionName(ctx, args.action);
-        const messages = args.messages ?? [{ role: "user", content: args.prompt ?? "" }];
+    async meter(ctx, opts, run) {
+        const userId = await resolveUserId(ctx, opts.userId);
+        const actionName = await resolveActionName(ctx, opts.action);
         const started = await ctx.runMutation(this.component.lib.startRequest, {
             userId,
             actionName,
-            model,
-            messages,
-            rerunOf: args.rerunOf,
+            tags: opts.tags,
+            model: opts.model,
+            messages: opts.messages,
+            rerunOf: opts.rerunOf,
         });
         if (!started.allowed) {
+            await this.fireLimitReached({
+                userId,
+                action: actionName,
+                tags: opts.tags,
+                messages: [started.reason],
+                code: started.code,
+                reason: started.reason,
+            });
             throw new ConvexError({
                 kind: "AIBudgetLimit",
                 code: started.code,
@@ -122,30 +192,30 @@ export class AIBudget {
             });
         }
         const requestId = started.requestId;
-        const warnings = started.warnings;
-        await this.fireSoftLimit({ userId, action: actionName, requestId, warnings });
+        const { warnings, notices } = started;
+        await this.fireBudgetEvents({ userId, action: actionName, tags: opts.tags, requestId }, warnings, notices);
         const start = Date.now();
         try {
-            // The full chain (incl. system) is stored on the request for audit/replay,
-            // but the AI SDK wants system prompts in the `system` option, not messages.
-            const system = messages
-                .filter((m) => m.role === "system")
-                .map((m) => m.content)
-                .join("\n\n") || undefined;
-            const convo = messages.filter((m) => m.role !== "system");
-            const result = await generateText({
-                model: convexGateway(model),
-                ...(system ? { system } : {}),
-                messages: convo,
-            });
-            const usage = extractUsage(result.usage);
+            const out = await run();
+            // Explicit token fields win; otherwise normalize a raw provider usage.
+            const usage = out.promptTokens !== undefined ||
+                out.completionTokens !== undefined ||
+                out.cachedTokens !== undefined
+                ? {
+                    promptTokens: out.promptTokens ?? 0,
+                    completionTokens: out.completionTokens ?? 0,
+                    cachedTokens: out.cachedTokens ?? 0,
+                }
+                : extractUsage(out.usage);
             const { costNanos } = await ctx.runMutation(this.component.lib.finishRequest, {
                 requestId,
-                responseText: result.text,
+                responseText: out.text,
                 ...usage,
+                serverToolUses: out.serverToolUses,
+                costNanos: out.costNanos,
                 latencyMs: Date.now() - start,
             });
-            return { text: result.text, requestId, costNanos, warnings, ...usage };
+            return { text: out.text ?? "", requestId, costNanos, warnings, notices, ...usage };
         }
         catch (e) {
             await ctx.runMutation(this.component.lib.finishRequest, {
@@ -156,6 +226,40 @@ export class AIBudget {
             throw e;
         }
     }
+    /**
+     * One-shot chat through the AI Gateway with tracking + limits — sugar over
+     * `meter`. Call from an action. `userId` defaults to the authenticated caller.
+     */
+    async chat(ctx, args = {}) {
+        const model = args.model ?? this.defaultModel;
+        const messages = args.messages ?? [{ role: "user", content: args.prompt ?? "" }];
+        return this.meter(ctx, {
+            model,
+            messages,
+            userId: args.userId,
+            action: args.action,
+            tags: args.tags,
+            rerunOf: args.rerunOf,
+        }, async () => {
+            // The full chain (incl. system) is stored for audit/replay, but the AI
+            // SDK wants system prompts in the `system` option, not messages.
+            const system = messages
+                .filter((m) => m.role === "system")
+                .map((m) => m.content)
+                .join("\n\n") || undefined;
+            const convo = messages.filter((m) => m.role !== "system");
+            const result = await generateText({
+                model: convexGateway(model),
+                ...(system ? { system } : {}),
+                messages: convo,
+            });
+            return {
+                text: result.text,
+                usage: result.usage,
+                costNanos: extractGatewayCostNanos(result),
+            };
+        });
+    }
     /**
      * An AI SDK LanguageModel that enforces limits and records usage/cost for
      * `userId` on every call. Drop it into `generateText`, `streamText`, or the
@@ -165,29 +269,33 @@ 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,
+                tags: opts.tags,
                 model: modelId,
                 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,
-                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 });
@@ -202,6 +310,7 @@ export class AIBudget {
                         await finish(requestId, {
                             responseText: extractText(result),
                             ...extractUsage(result.usage),
+                            costNanos: extractGatewayCostNanos(result),
                             latencyMs: Date.now() - start,
                         });
                         return result;
@@ -219,6 +328,7 @@ export class AIBudget {
                     const start = Date.now();
                     let text = "";
                     let usage = undefined;
+                    let providerMetadata = undefined;
                     try {
                         const result = await doStream();
                         // finishRequest is idempotent (terminal-guarded server-side), so
@@ -235,6 +345,7 @@ export class AIBudget {
                                 responseText: text,
                                 error,
                                 ...extractUsage(usage),
+                                costNanos: extractGatewayCostNanos({ providerMetadata }),
                                 latencyMs: Date.now() - start,
                             });
                         };
@@ -243,8 +354,10 @@ export class AIBudget {
                                 if (chunk?.type === "text-delta") {
                                     text += chunk.delta ?? chunk.textDelta ?? "";
                                 }
-                                if (chunk?.type === "finish")
+                                if (chunk?.type === "finish") {
                                     usage = chunk.usage;
+                                    providerMetadata = chunk.providerMetadata ?? providerMetadata;
+                                }
                                 if (chunk?.type === "error")
                                     void settle(String(chunk.error));
                                 controller.enqueue(chunk);
@@ -278,6 +391,7 @@ export class AIBudget {
             messages: args.messages ?? original.messages,
             rerunOf: args.requestId,
             action: original.actionName,
+            tags: original.tags,
         });
     }
     // ---------- namespaced admin API ----------
@@ -285,35 +399,88 @@ 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),
+            /** One request, including its stored prompt and response. */
+            get: (ctx, args) => ctx.runQuery(c.lib.getRequest, { requestId: args.requestId }),
             /** Ancestors up to the original, plus direct re-runs. */
             lineage: (ctx, args) => ctx.runQuery(c.lib.lineage, { requestId: args.requestId }),
             /** Replay a stored request, optionally with edited messages/model. */
             rerun: (ctx, args) => this.rerunImpl(ctx, args),
         };
     }
-    /** Per-user budgets and controls. */
-    get users() {
+    /**
+     * Budgets and controls for an arbitrary attribution dimension — the
+     * 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", 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 {
-            list: (ctx) => ctx.runQuery(c.lib.listUsers, {}),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setLimits, args),
-            /** One-time "approve another $X" bump (daily is today-only). */
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpUser, args),
-            /** Delete a user and all their request rows. */
-            delete: (ctx, args) => ctx.runMutation(c.lib.deleteUser, args),
+            /** 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: key(args) }),
+            setLimits: (ctx, args) => {
+                const { value, userId, name, ...limits } = args;
+                return ctx.runMutation(c.lib.setBucketLimits, {
+                    dimension,
+                    value: key(args),
+                    ...limits,
+                });
+            },
+            /** One-time "approve another $X" bump (daily/monthly reset with the window). */
+            bump: (ctx, args) => ctx.runMutation(c.lib.bumpBucket, {
+                dimension,
+                value: key(args),
+                dailyNanos: args.dailyNanos,
+                monthlyNanos: args.monthlyNanos,
+                lifetimeNanos: args.lifetimeNanos,
+            }),
+            /** 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-action (per-feature) budgets. */
+    /** 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.listActions, {}),
-            setLimits: (ctx, args) => ctx.runMutation(c.lib.setActionLimits, args),
-            bump: (ctx, args) => ctx.runMutation(c.lib.bumpAction, args),
-        };
+        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 {
@@ -322,6 +489,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),
         };
@@ -335,13 +504,128 @@ export class AIBudget {
             setPolicy: (ctx, args) => ctx.runMutation(c.lib.setModelPolicy, args),
         };
     }
-    /** Per-model prices (cents per million tokens). */
+    /** Per-model prices (nanodollars per million tokens) + server-tool fees. */
     get prices() {
         const c = this.component;
         return {
             list: (ctx) => ctx.runQuery(c.lib.listPrices, {}),
             set: (ctx, args) => ctx.runMutation(c.lib.setPrice, args),
+            /** Per-call fees for provider server tools (web search, etc.). */
+            listServerTools: (ctx) => ctx.runQuery(c.lib.listServerToolPrices, {}),
+            /** Set a server-tool's per-call price, e.g. { tool: "web_search", nanosPerCall }. */
+            setServerTool: (ctx, args) => ctx.runMutation(c.lib.setServerToolPrice, 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, "");
+            // `?token=` is accepted only for the initial page navigation (a browser
+            // GET can't set headers); the page strips it from the URL on load and the
+            // JSON API is called with the bearer header. Compared in constant time.
+            const provided = bearer || url.searchParams.get("token") || "";
+            return { ok: timingSafeEqual(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);
+                }
+            }
+            // Inject as JSON literals (function replacers so `$` in the value isn't
+            // treated as a replacement pattern). This keeps a token/prefix containing
+            // quotes, backslashes, or `</script>` from breaking out of the JS string.
+            const html = DASHBOARD_HTML.replace(/__API_BASE__/g, () => JSON.stringify(`${prefix}/api`)).replace(/__TOKEN__/g, () => JSON.stringify(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`. */