floe-guard 0.1.0 → 0.2.1

package/dist/index.d.ts

@@ -1,5 +1,3 @@
-import { LanguageModelV1Middleware } from 'ai';
-
 /**
  * Offline token pricing from a vendored LiteLLM cost map.
  *
@@ -25,8 +23,9 @@ interface PricedModel {
 /**
  * Resolve a model to its per-token price, or `null` if it cannot be priced.
  *
- * Overrides win, then the bundled cost map (looked up by bare name, then the raw
- * field). Fail-closed: both prices must be finite, else `null`.
+ * Per specificity group (exact forms first, date-stripped fallbacks second):
+ * overrides win, then the bundled cost map. Fail-closed: the first matching
+ * entry must have finite prices, else `null`.
  */
 declare function resolvePrice(model: string, overrides?: Record<string, ManualPrice>): PricedModel | null;
 /** USD cost for token usage. Negative counts are clamped to zero. */
@@ -51,6 +50,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 +79,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 +129,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 +168,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 +227,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 +288,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 };