floe-guard 0.1.0 → 0.2.0

package/dist/index.d.ts

@@ -1,5 +1,3 @@
-import { LanguageModelV1Middleware } from 'ai';
-
 /**
  * Offline token pricing from a vendored LiteLLM cost map.
  *
@@ -51,6 +49,16 @@ declare namespace pricing {
  * 2. Call {@link BudgetGuard.record} AFTER every response, with the token usage.
  *    It prices the tokens offline and accrues the USD into a running total.
  *
+ * **Concurrency.** `check()` then `record()` is a check-then-act with an `await`
+ * in between. Fire several model calls at once (e.g. `Promise.all`) and they all
+ * `check()` against the same under-limit total before any `record()` lands, so
+ * the ceiling is blown (see issue #18). {@link BudgetGuard.reserve} /
+ * {@link BudgetGuard.settle} close that gap: `reserve()` holds the estimated cost
+ * in flight (synchronously, before the await), so parallel callers each take
+ * their own slice of the ceiling. JS is single-threaded, so an in-flight counter
+ * is enough — no lock needed. The middleware uses it; `check`/`record` are
+ * unchanged.
+ *
  * This is a faithful port of `src/floe_guard/guard.py` — same prediction logic,
  * same epsilon handling, same fail-closed default.
  */
@@ -70,15 +78,47 @@ interface BudgetGuardOptions {
      * `BUDGET EXCEEDED — call blocked` banner to stderr.
      */
     onBlock?: (spentUsd: number, limitUsd: number) => void;
+    /**
+     * Utilization (basis points, 0..10000) at which {@link BudgetGuard.advisory}
+     * flags `nearLimit` so an agent can taper before the hard-stop. Default 8000.
+     */
+    nearLimitBps?: number;
+}
+/**
+ * A context-aware spend signal for the single local budget.
+ *
+ * Mirrors the core fields of hosted Floe's `X-Floe-Budget-Advisory` header, so
+ * agent logic that reads it (taper as you approach the cap, stop at it) ports
+ * unchanged to the hosted path. Hosted adds what a local, single-budget guard
+ * cannot know: which of several caps is tightest (`scope` across
+ * `credit_line | session | task | api | vendor`), cross-vendor reasoning,
+ * server-truth balances, and rolling-window reset timing.
+ *
+ * This is a **soft** signal — the model may ignore it. The hard-stop
+ * ({@link BudgetGuard.check}) is what enforces the ceiling; the advisory is
+ * upside (let the agent finish on budget rather than be cut off).
+ */
+interface BudgetAdvisory {
+    nearLimit: boolean;
+    /** Utilization in basis points, 0..10000 (8500 = 85%). */
+    usedBps: number;
+    remainingUsd: number;
+    limitUsd: number;
+    spentUsd: number;
+    /** Hosted reports the tightest cap across all scopes; local is always "local". */
+    scope: "local";
 }
 declare class BudgetGuard {
     readonly limitUsd: number;
     spentUsd: number;
     priceOverrides?: Record<string, ManualPrice>;
     failClosed: boolean;
+    nearLimitBps: number;
     private readonly onBlock;
     /** Cost of the most recent priced call, used to predict the next one. */
     private lastCost;
+    /** USD held for in-flight calls (reserved, not yet settled). Counts toward the ceiling. */
+    private reserved;
     /**
      * @param limitUsd the spend ceiling, in USD. `0` blocks the very first call.
      */
@@ -88,10 +128,35 @@ declare class BudgetGuard {
      *
      * 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?: number): void;
+    /**
+     * Atomically check the ceiling AND hold the estimated cost in flight.
+     *
+     * 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.
+     */
+    reserve(estimatedCost?: number): number;
+    /**
+     * 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: string, promptTokens: number, completionTokens: number, options?: {
+        reserved?: number;
+        price?: ManualPrice;
+    }): number;
     /**
      * Price one response's tokens offline and add the cost to the total.
      *
@@ -102,8 +167,21 @@ declare class BudgetGuard {
     record(model: string, promptTokens: number, completionTokens: number, options?: {
         price?: ManualPrice;
     }): number;
-    /** USD left before the ceiling (never negative). */
+    /**
+     * Drop an in-flight reservation without recording spend (e.g. the call failed
+     * before producing usage). Safe to call with `0`.
+     */
+    release(reserved: number): void;
+    /** USD left before the ceiling, net of in-flight reservations (never negative). */
     get remainingUsd(): number;
+    /**
+     * 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(): BudgetAdvisory;
 }
 
 /**
@@ -148,18 +226,55 @@ declare class UnpriceableModelError extends FloeGuardError {
  * This is the TypeScript counterpart to the Python framework adapters. The AI SDK
  * is TypeScript-only, so it ships as its own npm package.
  *
- * Verified against `ai@4.3.19` (`LanguageModelV1Middleware`):
- *   - `wrapGenerate({ doGenerate, model })` — we `check()` (throws to hard-stop)
- *     BEFORE calling `doGenerate()`, then `record()` from `result.usage`.
- *   - `wrapStream({ doStream, model })` — we `check()` BEFORE `doStream()`, then
- *     read `usage` from the `finish` part as the stream drains.
+ * Works with BOTH `ai@4` (`LanguageModelV1Middleware`) and `ai@5`
+ * (`LanguageModelV2Middleware`). The two majors renamed the middleware type and
+ * the usage fields (`promptTokens`/`completionTokens` → `inputTokens`/
+ * `outputTokens`), so this module deliberately imports nothing from `ai`: the
+ * middleware is typed structurally against the surface both majors share, and
+ * usage is read from whichever field pair the installed SDK reports.
+ *
+ *   - `wrapGenerate({ doGenerate, model })` — we `reserve()` (throws to hard-stop)
+ *     BEFORE calling `doGenerate()`, hold the reservation across the await, then
+ *     `settle()` from `result.usage`.
+ *   - `wrapStream({ doStream, model })` — we `reserve()` BEFORE `doStream()`, then
+ *     `settle()` from the `finish` part as the stream drains.
+ *
+ * Reserving before the await is what makes parallel calls (`Promise.all` over
+ * several generations) honour the ceiling: each holds its slice instead of all
+ * reading the same stale total (issue #18). The reservation is released if the
+ * call throws, or if a stream ends without reporting usage.
  *
  * The model id used for pricing comes from `model.modelId`.
  */
 
 /**
- * Build a `LanguageModelV1Middleware` that hard-stops the model before a call
+ * The middleware call surface shared by `ai@4` and `ai@5`. Both majors invoke
+ * `wrapGenerate`/`wrapStream` with an options object carrying these members
+ * (plus richer `model` fields we don't read). `doGenerate`/`doStream` are
+ * declared optional so every 4.x/5.x minor's options type stays assignable —
+ * the SDK always provides the one each hook actually calls.
+ */
+interface MiddlewareCallOptions {
+    doGenerate?: () => PromiseLike<any>;
+    doStream?: () => PromiseLike<any>;
+    model: {
+        modelId: string;
+    };
+    params?: unknown;
+}
+/**
+ * Structural stand-in for `LanguageModelV1Middleware` (ai@4) and
+ * `LanguageModelV2Middleware` (ai@5) — assignable to the `middleware` option of
+ * `wrapLanguageModel` on either major.
+ */
+interface BudgetGuardMiddleware {
+    wrapGenerate: (options: MiddlewareCallOptions) => Promise<any>;
+    wrapStream: (options: MiddlewareCallOptions) => Promise<any>;
+}
+/**
+ * Build a budget-guard middleware that hard-stops the model before a call
  * crosses the guard's USD ceiling, and records priced token usage after.
+ * Compatible with `wrapLanguageModel` from both `ai@4` and `ai@5`.
  *
  * @example
  * import { wrapLanguageModel } from "ai";
@@ -172,6 +287,6 @@ declare class UnpriceableModelError extends FloeGuardError {
  *   middleware: budgetGuardMiddleware(guard),
  * });
  */
-declare function budgetGuardMiddleware(guard: BudgetGuard): LanguageModelV1Middleware;
+declare function budgetGuardMiddleware(guard: BudgetGuard): BudgetGuardMiddleware;
 
-export { BudgetExceeded, BudgetGuard, type BudgetGuardOptions, FloeGuardError, type ManualPrice, type PricedModel, UnpriceableModelError, budgetGuardMiddleware, priceTokens, pricing, resolvePrice };
+export { type BudgetAdvisory, BudgetExceeded, BudgetGuard, type BudgetGuardMiddleware, type BudgetGuardOptions, FloeGuardError, type ManualPrice, type PricedModel, UnpriceableModelError, budgetGuardMiddleware, priceTokens, pricing, resolvePrice };

package/dist/index.js

@@ -80,6 +80,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,
@@ -704,6 +710,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,
@@ -752,6 +776,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,
@@ -831,9 +861,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.
    */
@@ -843,35 +876,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 };
@@ -881,12 +960,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;
@@ -894,9 +983,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) {
@@ -908,33 +1043,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 };