@convex-dev/ai-budget 0.0.2-alpha.10 → 0.0.2-alpha.13

package/src/component/lib.ts

@@ -43,6 +43,15 @@ const DEFAULT_PRICES: Record<string, { input: number; output: number }> = {
   "openai/gpt-5-mini": { input: 250_000_000, output: 2_000_000_000 },
 };
 
+// Per-call price (nanodollars) for provider server-side tools that bill a fee on
+// top of tokens — e.g. Anthropic web search at ~$0.01/call. Keyed by the tool
+// name the caller reports in `serverToolUses` (e.g. { web_search: 3 }). Used
+// only when a request settles WITHOUT an authoritative gateway cost; if you pass
+// `costNanos`, that already includes tool fees. Override via setServerToolPrice.
+const DEFAULT_SERVER_TOOL_PRICES: Record<string, number> = {
+  web_search: 10_000_000, // $0.01 per search
+};
+
 // Pessimistic assumed output length when reserving budget up front. This makes
 // concurrent admission atomic against the estimate; a response that exceeds the
 // estimate can still settle above the cap by the estimation delta.
@@ -136,6 +145,22 @@ const settleCost = (
   );
 };
 
+// Per-call fees for provider server tools (web search, etc.), merging the
+// defaults with any deployment overrides. Unknown tools price at 0 (recorded
+// but not charged) rather than guessing.
+const serverToolCost = (
+  uses: Record<string, number> | undefined,
+  overrides: Record<string, number> | undefined
+) => {
+  if (!uses) return 0;
+  const prices = { ...DEFAULT_SERVER_TOOL_PRICES, ...(overrides ?? {}) };
+  let total = 0;
+  for (const [tool, count] of Object.entries(uses)) {
+    if (count > 0 && prices[tool] > 0) total += Math.round(count * prices[tool]);
+  }
+  return total;
+};
+
 // Upsert-add a settled amount into the durable per-(bucket, period) usage row.
 // These rows are never swept by request retention, so spend history survives.
 async function addUsage(
@@ -394,6 +419,15 @@ export const startRequest = mutation({
     tags: v.optional(v.array(vTag)),
     model: v.string(),
     messages: v.array(vMessage),
+    // Reserve this exact amount (nanodollars) instead of the token-based
+    // estimate. Use it whenever the cost is known up front — image generation
+    // (n × per-image), audio, per-call APIs — so a hard cap reserves the real
+    // amount rather than a meaningless token guess.
+    estimatedCostNanos: v.optional(v.number()),
+    // Hold the reservation this long (ms) before the reconciler may reap it —
+    // for long async jobs (video) that settle minutes later. Extends the 30-min
+    // default floor.
+    reserveTtlMs: v.optional(v.number()),
     rerunOf: v.optional(v.id("requests")),
   },
   returns: vStartResult,
@@ -424,6 +458,11 @@ export const startRequest = mutation({
     const month = monthStamp();
     const priceInfo = await getPrice(ctx, args.model);
     const est = estimateUsage(args.messages, priceInfo);
+    // A caller-supplied known cost (image gen, audio, per-call APIs) reserves
+    // the real amount up front; the token estimate stays as the token reserve.
+    if (args.estimatedCostNanos !== undefined && args.estimatedCostNanos >= 0) {
+      est.cost = Math.round(args.estimatedCostNanos);
+    }
     const warnings: string[] = [];
     const notices: string[] = [];
 
@@ -657,6 +696,7 @@ export const startRequest = mutation({
       model: args.model,
       messages: args.messages,
       rerunOf: args.rerunOf,
+      ...(args.reserveTtlMs !== undefined ? { reserveTtlMs: args.reserveTtlMs } : {}),
       status: "pending",
       estimatedNanos: est.cost,
       estimatedTokens: est.tokens,
@@ -684,9 +724,13 @@ export const finishRequest = mutation({
     promptTokens: v.optional(v.number()),
     completionTokens: v.optional(v.number()),
     cachedTokens: v.optional(v.number()),
-    // Authoritative cost from the gateway, if it ever reports one. When present
-    // it's recorded verbatim (no token-based estimate); when absent we price
-    // from tokens (cache-aware). Wired now so adopting a real cost is one line.
+    // Provider server-tool invocations that bill a per-call fee (e.g.
+    // { web_search: 3 }). Added to the token cost when no authoritative cost is
+    // supplied; recorded either way.
+    serverToolUses: v.optional(v.record(v.string(), v.number())),
+    // Authoritative cost from the gateway/provider, if reported. When present
+    // it's recorded verbatim (already includes any tool fees); when absent we
+    // price from tokens (cache-aware) plus server-tool fees.
     costNanos: v.optional(v.number()),
     latencyMs: v.optional(v.number()),
   },
@@ -711,17 +755,22 @@ export const finishRequest = mutation({
     const promptTokens = Math.max(0, args.promptTokens ?? 0);
     const completionTokens = Math.max(0, args.completionTokens ?? 0);
     const cachedTokens = Math.min(promptTokens, Math.max(0, args.cachedTokens ?? 0));
-    // Prefer an authoritative gateway cost when supplied; otherwise price from
-    // tokens, discounting the cached (prompt-cache-read) slice of the prompt.
-    const costNanos =
-      args.costNanos !== undefined && args.costNanos >= 0
-        ? Math.round(args.costNanos)
-        : settleCost(
-            promptTokens,
-            cachedTokens,
-            completionTokens,
-            await getPrice(ctx, request.model)
-          );
+    // Prefer an authoritative gateway cost when supplied (it already includes
+    // tool fees); otherwise price from tokens — discounting the cached
+    // (prompt-cache-read) slice — plus any server-tool per-call fees.
+    let costNanos: number;
+    if (args.costNanos !== undefined && args.costNanos >= 0) {
+      costNanos = Math.round(args.costNanos);
+    } else {
+      const settings = await getSettings(ctx);
+      costNanos =
+        settleCost(
+          promptTokens,
+          cachedTokens,
+          completionTokens,
+          await getPrice(ctx, request.model)
+        ) + serverToolCost(args.serverToolUses, settings?.serverToolPrices);
+    }
 
     // Durable write to the request's OWN row only — uncontended, so it always
     // lands. `settled: false` hands it to the fold step; the row is never left
@@ -733,6 +782,7 @@ export const finishRequest = mutation({
       promptTokens,
       completionTokens,
       ...(cachedTokens > 0 ? { cachedTokens } : {}),
+      ...(args.serverToolUses ? { serverToolUses: args.serverToolUses } : {}),
       costNanos,
       latencyMs: args.latencyMs,
       settled: false,
@@ -833,14 +883,20 @@ export const reconcile = internalMutation({
       .take(200);
     for (const req of toFold) await foldOne(ctx, req);
 
+    // Reap dead reservations: pending rows older than the default floor, but a
+    // per-request reserveTtlMs (set for long async jobs like video) holds the
+    // reservation until *its* deadline so a still-running job isn't reaped.
     const cutoff = Date.now() - STALE_PENDING_MS;
-    const stale = await ctx.db
+    const candidates = await ctx.db
       .query("requests")
       .withIndex("status", (q) =>
         q.eq("status", "pending").lt("_creationTime", cutoff)
       )
       .take(200);
-    for (const req of stale) {
+    let expired = 0;
+    for (const req of candidates) {
+      const ttl = req.reserveTtlMs ?? STALE_PENDING_MS;
+      if (Date.now() - req._creationTime <= ttl) continue; // still within its window
       await ctx.db.patch(req._id, {
         status: "error",
         error: "Timed out before settling; reservation released",
@@ -848,6 +904,7 @@ export const reconcile = internalMutation({
         settled: false,
       });
       await foldOne(ctx, await ctx.db.get(req._id));
+      expired++;
     }
 
     // Retention: delete terminal, fully-accounted request rows past the window.
@@ -872,7 +929,7 @@ export const reconcile = internalMutation({
         }
       }
     }
-    return { folded: toFold.length, expired: stale.length, purged };
+    return { folded: toFold.length, expired, purged };
   },
 });
 
@@ -1347,3 +1404,26 @@ export const listPrices = query({
     return merged;
   },
 });
+
+// Per-call fees for provider server tools (web search, etc.), defaults merged
+// with any deployment overrides.
+export const listServerToolPrices = query({
+  args: {},
+  handler: async (ctx) => {
+    const s = await getSettings(ctx as any);
+    return { ...DEFAULT_SERVER_TOOL_PRICES, ...(s?.serverToolPrices ?? {}) };
+  },
+});
+
+export const setServerToolPrice = mutation({
+  args: { tool: v.string(), nanosPerCall: v.number() },
+  returns: v.null(),
+  handler: async (ctx, { tool, nanosPerCall }) => {
+    if (nanosPerCall < 0) throw new Error("Prices must be non-negative");
+    const s = await getSettings(ctx);
+    const serverToolPrices = { ...(s?.serverToolPrices ?? {}), [tool]: nanosPerCall };
+    if (s) await ctx.db.patch(s._id, { serverToolPrices });
+    else await ctx.db.insert("settings", { key: "singleton", serverToolPrices });
+    return null;
+  },
+});

package/src/component/schema.ts

@@ -132,6 +132,14 @@ export default defineSchema({
     completionTokens: v.optional(v.number()),
     // subset of promptTokens served from the provider's prompt cache (cheaper).
     cachedTokens: v.optional(v.number()),
+    // server-side tool invocations that bill a per-call fee on top of tokens
+    // (e.g. { web_search: 3 }). Priced via serverToolPrices at settle.
+    serverToolUses: v.optional(v.record(v.string(), v.number())),
+    // How long the reservation may stay held before the reconciler reaps it as
+    // dead (ms). For long async jobs (video generation) set this to the job's
+    // max duration so the hold isn't released mid-flight. Extends the default
+    // 30-min floor; only stored while pending.
+    reserveTtlMs: v.optional(v.number()),
     costNanos: v.optional(v.number()),
     latencyMs: v.optional(v.number()),
     rerunOf: v.optional(v.id("requests")),
@@ -184,5 +192,8 @@ export default defineSchema({
     // default approaching-limit alert threshold (fraction of a cap) for buckets
     // that don't set their own warnAtPct. 0/unset disables threshold alerts.
     defaultWarnAtPct: v.optional(v.number()),
+    // per-call price (nanodollars) overrides for provider server tools, keyed by
+    // tool name (e.g. { web_search: 12_000_000 }). Merged over the defaults.
+    serverToolPrices: v.optional(v.record(v.string(), v.number())),
   }).index("key", ["key"]),
 });