floe-guard 0.1.0 → 0.2.0

package/README.md

@@ -1,12 +1,17 @@
 # floe-guard (Vercel AI SDK)
 
+[![npm version](https://img.shields.io/npm/v/floe-guard.svg)](https://www.npmjs.com/package/floe-guard)
+[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](../LICENSE)
+
 **A local budget guardrail for AI agents** — the TypeScript counterpart to the
 [Python `floe-guard`](../README.md). It hard-stops your agent *before its next LLM
 call* when it would cross a USD spend ceiling. No account, no signup, no network.
 Runs in your process.
 
+Works with both **AI SDK v4 and v5** (`ai@4` / `ai@5`).
+
 ```bash
-npm i floe-guard ai@4 @ai-sdk/openai
+npm i floe-guard ai @ai-sdk/openai
 ```
 
 ```ts
@@ -44,10 +49,33 @@ const guard = new BudgetGuard(5.0, {
 });
 ```
 
-## Verified against
+## Context-aware budgeting
+
+`guard.advisory()` returns a soft signal you can act on before a call — `nearLimit`,
+`usedBps`, `remainingUsd` — so the agent can taper near the cap instead of being
+cut off. The hard-stop (`check`) is still the guarantee. This is the same shape
+the Python package exposes and that hosted Floe returns on every proxied call
+(`X-Floe-Budget-Advisory`), so the logic ports unchanged to the hosted path.
+
+```ts
+const guard = new BudgetGuard(0.1, { nearLimitBps: 7000 }); // flag at 70% used
+const adv = guard.advisory();
+const model = adv.nearLimit ? openai("gpt-4o-mini") : openai("gpt-4o");
+```
+
+## Compatibility
+
+`ai` is declared as a peer dependency with the range `>=4.0.0 <6.0.0`:
+
+- **`ai@4`** — `LanguageModelV1Middleware` via `wrapLanguageModel` /
+  `experimental_wrapLanguageModel`; usage read from
+  `promptTokens`/`completionTokens`.
+- **`ai@5`** — `LanguageModelV2Middleware` via `wrapLanguageModel`; usage read
+  from `inputTokens`/`outputTokens`.
 
-`ai@4` (`LanguageModelV1Middleware` via `wrapLanguageModel` /
-`experimental_wrapLanguageModel`). Declared as a peer dependency.
+The middleware imports nothing from `ai` at runtime or in its types, so one
+build serves both majors. If a provider reports no usable token counts, the
+call is rejected (fail-closed) rather than metered as $0.
 
 ## Development
 

package/dist/index.cjs

@@ -107,6 +107,12 @@ var cost_map_default = {
     litellm_provider: "anthropic",
     mode: "chat"
   },
+  "claude-fable-5": {
+    input_cost_per_token: 1e-5,
+    output_cost_per_token: 5e-5,
+    litellm_provider: "anthropic",
+    mode: "chat"
+  },
   "claude-haiku-4-5": {
     input_cost_per_token: 1e-6,
     output_cost_per_token: 5e-6,
@@ -731,6 +737,24 @@ var cost_map_default = {
     litellm_provider: "openai",
     mode: "chat"
   },
+  "llama-3.1-8b-instant": {
+    input_cost_per_token: 5e-8,
+    output_cost_per_token: 8e-8,
+    litellm_provider: "groq",
+    mode: "chat"
+  },
+  "llama-3.3-70b-versatile": {
+    input_cost_per_token: 59e-8,
+    output_cost_per_token: 79e-8,
+    litellm_provider: "groq",
+    mode: "chat"
+  },
+  "meta-llama/llama-4-scout-17b-16e-instruct": {
+    input_cost_per_token: 11e-8,
+    output_cost_per_token: 34e-8,
+    litellm_provider: "groq",
+    mode: "chat"
+  },
   o1: {
     input_cost_per_token: 15e-6,
     output_cost_per_token: 6e-5,
@@ -779,6 +803,12 @@ var cost_map_default = {
     litellm_provider: "openai",
     mode: "chat"
   },
+  "qwen/qwen3-32b": {
+    input_cost_per_token: 29e-8,
+    output_cost_per_token: 59e-8,
+    litellm_provider: "groq",
+    mode: "chat"
+  },
   "text-embedding-3-large": {
     input_cost_per_token: 13e-8,
     output_cost_per_token: 0,
@@ -858,9 +888,12 @@ var BudgetGuard = class {
   spentUsd = 0;
   priceOverrides;
   failClosed;
+  nearLimitBps;
   onBlock;
   /** Cost of the most recent priced call, used to predict the next one. */
   lastCost = 0;
+  /** USD held for in-flight calls (reserved, not yet settled). Counts toward the ceiling. */
+  reserved = 0;
   /**
    * @param limitUsd the spend ceiling, in USD. `0` blocks the very first call.
    */
@@ -870,35 +903,81 @@ var BudgetGuard = class {
         `limitUsd must be a finite, non-negative number, got ${limitUsd}`
       );
     }
+    const nearLimitBps = options.nearLimitBps === void 0 ? 8e3 : options.nearLimitBps;
+    if (!Number.isInteger(nearLimitBps) || nearLimitBps < 0 || nearLimitBps > 1e4) {
+      throw new RangeError(
+        `nearLimitBps must be an integer in 0..10000, got ${nearLimitBps}`
+      );
+    }
     this.limitUsd = limitUsd;
     this.priceOverrides = options.priceOverrides;
     this.failClosed = options.failClosed ?? true;
     this.onBlock = options.onBlock ?? defaultOnBlock;
+    this.nearLimitBps = nearLimitBps;
   }
   /**
    * Throw {@link BudgetExceeded} if the next call would cross the ceiling.
    *
    * Call this immediately before each LLM request. The "next call" is estimated
    * from the last recorded call's cost (override with `estimatedNextCost`); the
-   * first call is always allowed unless the ceiling is already met. A check on
-   * the running total catches an overshoot if the estimate was too low.
+   * first call is always allowed unless the ceiling is already met. In-flight
+   * reservations count toward the total, so this stays correct alongside
+   * {@link BudgetGuard.reserve}.
+   *
+   * Note: `check` is a non-binding peek. For parallel calls, use `reserve()` /
+   * `settle()`, which hold the estimate across the await.
    */
   check(estimatedNextCost) {
-    const estimate = estimatedNextCost === void 0 ? this.lastCost : Math.max(0, estimatedNextCost);
-    const projected = this.spentUsd + estimate;
-    if (this.spentUsd > this.limitUsd - EPS || projected > this.limitUsd + EPS) {
+    const rawEstimate = estimatedNextCost === void 0 ? this.lastCost : estimatedNextCost;
+    if (!Number.isFinite(rawEstimate)) {
+      throw new RangeError(
+        `estimatedNextCost must be a finite number, got ${rawEstimate}`
+      );
+    }
+    const estimate = Math.max(0, rawEstimate);
+    const committed = this.spentUsd + this.reserved;
+    if (committed > this.limitUsd - EPS || committed + estimate > this.limitUsd + EPS) {
       this.onBlock(this.spentUsd, this.limitUsd);
       throw new BudgetExceeded(this.spentUsd, this.limitUsd);
     }
   }
   /**
-   * Price one response's tokens offline and add the cost to the total.
+   * Atomically check the ceiling AND hold the estimated cost in flight.
    *
-   * Returns the USD cost of this call. If the model is unpriceable and no `price`
-   * is given, behaviour depends on `failClosed`: warn + throw (default), or
-   * warn + skip accrual.
+   * The concurrency-safe enforcement path: call before the request and hold the
+   * returned reservation across the await, so parallel callers can't all clear
+   * the same stale total. Throws {@link BudgetExceeded} (without reserving) if
+   * the reservation would cross the ceiling. Returns the reservation handle to
+   * pass to {@link BudgetGuard.settle} (or {@link BudgetGuard.release} on error).
+   * `estimatedCost` defaults to the last call's cost.
    */
-  record(model, promptTokens, completionTokens, options = {}) {
+  reserve(estimatedCost) {
+    const rawEstimate = estimatedCost === void 0 ? this.lastCost : estimatedCost;
+    if (!Number.isFinite(rawEstimate)) {
+      throw new RangeError(
+        `estimatedCost must be a finite number, got ${rawEstimate}`
+      );
+    }
+    const estimate = Math.max(0, rawEstimate);
+    const committed = this.spentUsd + this.reserved;
+    if (committed > this.limitUsd - EPS || committed + estimate > this.limitUsd + EPS) {
+      this.onBlock(this.spentUsd, this.limitUsd);
+      throw new BudgetExceeded(this.spentUsd, this.limitUsd);
+    }
+    this.reserved += estimate;
+    return estimate;
+  }
+  /**
+   * Release a reservation and record the actual cost. `record` is `settle` with
+   * no reservation. Returns the USD cost of this call; unpriceable-model handling
+   * matches {@link BudgetGuard.record}, and any held reservation is released even
+   * on the warn-and-skip path.
+   */
+  settle(model, promptTokens, completionTokens, options = {}) {
+    const reserved = options.reserved ?? 0;
+    if (!Number.isFinite(reserved) || reserved < 0) {
+      throw new RangeError(`reserved must be a finite, non-negative number, got ${reserved}`);
+    }
     let overrides = this.priceOverrides;
     if (options.price !== void 0) {
       overrides = { ...overrides ?? {}, [model]: options.price };
@@ -908,12 +987,22 @@ var BudgetGuard = class {
       console.warn(
         `Cannot price model '${model}': not in the bundled cost map and no manual price given. The budget guard cannot enforce a ceiling on spend it cannot measure \u2014 pass { price } or set it in priceOverrides.`
       );
+      this.release(reserved);
       if (this.failClosed) {
         throw new UnpriceableModelError(model);
       }
       return 0;
     }
-    const cost = priceTokens(priced, promptTokens, completionTokens);
+    let cost;
+    try {
+      cost = priceTokens(priced, promptTokens, completionTokens);
+    } catch (err) {
+      this.release(reserved);
+      throw err;
+    }
+    if (reserved) {
+      this.reserved = Math.max(0, this.reserved - reserved);
+    }
     this.spentUsd += cost;
     if (this.spentUsd - this.limitUsd > 0 && this.spentUsd - this.limitUsd < EPS) {
       this.spentUsd = this.limitUsd;
@@ -921,9 +1010,55 @@ var BudgetGuard = class {
     this.lastCost = cost;
     return cost;
   }
-  /** USD left before the ceiling (never negative). */
+  /**
+   * Price one response's tokens offline and add the cost to the total.
+   *
+   * Returns the USD cost of this call. If the model is unpriceable and no `price`
+   * is given, behaviour depends on `failClosed`: warn + throw (default), or
+   * warn + skip accrual.
+   */
+  record(model, promptTokens, completionTokens, options = {}) {
+    return this.settle(model, promptTokens, completionTokens, {
+      reserved: 0,
+      price: options.price
+    });
+  }
+  /**
+   * Drop an in-flight reservation without recording spend (e.g. the call failed
+   * before producing usage). Safe to call with `0`.
+   */
+  release(reserved) {
+    if (!Number.isFinite(reserved) || reserved < 0) {
+      throw new RangeError(`reserved must be a finite, non-negative number, got ${reserved}`);
+    }
+    if (!reserved) return;
+    this.reserved = Math.max(0, this.reserved - reserved);
+  }
+  /** USD left before the ceiling, net of in-flight reservations (never negative). */
   get remainingUsd() {
-    return Math.max(0, this.limitUsd - this.spentUsd);
+    return Math.max(0, this.limitUsd - this.spentUsd - this.reserved);
+  }
+  /**
+   * Context-aware spend advisory for this budget — see {@link BudgetAdvisory}.
+   *
+   * `nearLimit` flips once utilization reaches `nearLimitBps` (default 80%), so an
+   * agent can taper *before* the hard-stop. Advisory only: read it to adapt;
+   * {@link BudgetGuard.check} is what enforces the ceiling.
+   */
+  advisory() {
+    const usedBps = this.limitUsd <= 0 ? 1e4 : Math.max(0, Math.min(1e4, Math.floor(this.spentUsd / this.limitUsd * 1e4 + 1e-9)));
+    return {
+      nearLimit: usedBps >= this.nearLimitBps,
+      usedBps,
+      // Settled budget: limit minus accrued spend, deliberately NOT net of
+      // in-flight reservations. Unlike the remainingUsd getter (which subtracts
+      // `reserved`), the advisory is a soft utilization signal about money already
+      // spent, while the getter reports what a new call can still claim.
+      remainingUsd: Math.max(0, this.limitUsd - this.spentUsd),
+      limitUsd: this.limitUsd,
+      spentUsd: this.spentUsd,
+      scope: "local"
+    };
   }
 };
 function defaultOnBlock(spentUsd, limitUsd) {
@@ -935,33 +1070,88 @@ function defaultOnBlock(spentUsd, limitUsd) {
 }
 
 // src/middleware.ts
+function usageTokens(modelId, usage) {
+  const u = usage;
+  const promptTokens = u?.promptTokens ?? u?.inputTokens;
+  const completionTokens = u?.completionTokens ?? u?.outputTokens;
+  if (typeof promptTokens !== "number" || typeof completionTokens !== "number") {
+    throw new Error(
+      `Model '${modelId}' reported no token usage \u2014 the budget guard cannot meter spend it cannot see, so this call is rejected rather than treated as free.`
+    );
+  }
+  return { promptTokens, completionTokens };
+}
 function budgetGuardMiddleware(guard) {
   return {
     async wrapGenerate({ doGenerate, model }) {
-      guard.check();
-      const result = await doGenerate();
-      guard.record(
-        model.modelId,
-        result.usage.promptTokens,
-        result.usage.completionTokens
-      );
-      return result;
+      const reserved = guard.reserve();
+      let result;
+      try {
+        result = await doGenerate();
+      } catch (err) {
+        guard.release(reserved);
+        throw err;
+      }
+      let handled = false;
+      try {
+        const { promptTokens, completionTokens } = usageTokens(
+          model.modelId,
+          result?.usage
+        );
+        handled = true;
+        guard.settle(model.modelId, promptTokens, completionTokens, { reserved });
+        return result;
+      } catch (err) {
+        if (!handled) guard.release(reserved);
+        throw err;
+      }
     },
     async wrapStream({ doStream, model }) {
-      guard.check();
-      const { stream, ...rest } = await doStream();
+      const reserved = guard.reserve();
+      let streamResult;
+      try {
+        streamResult = await doStream();
+      } catch (err) {
+        guard.release(reserved);
+        throw err;
+      }
+      const { stream, ...rest } = streamResult;
+      let handled = false;
       const guarded = stream.pipeThrough(
         new TransformStream({
           transform(chunk, controller) {
-            if (chunk.type === "finish") {
-              guard.record(
-                model.modelId,
-                chunk.usage.promptTokens,
-                chunk.usage.completionTokens
-              );
+            if (chunk?.type === "finish" && !handled) {
+              try {
+                const { promptTokens, completionTokens } = usageTokens(
+                  model.modelId,
+                  chunk.usage
+                );
+                handled = true;
+                guard.settle(model.modelId, promptTokens, completionTokens, { reserved });
+              } catch (err) {
+                if (!handled) {
+                  handled = true;
+                  guard.release(reserved);
+                }
+                throw err;
+              }
             }
             controller.enqueue(chunk);
+          },
+          flush() {
+            if (!handled) {
+              handled = true;
+              guard.release(reserved);
+            }
+          },
+          cancel() {
+            if (!handled) {
+              handled = true;
+              guard.release(reserved);
+            }
           }
+          // `cancel` is valid per the Streams spec and supported in Node 18+, but
+          // TS's Transformer lib type lags and omits it — cast to keep the type check.
         })
       );
       return { stream: guarded, ...rest };