floe-guard 0.7.0 → 0.9.0

package/dist/index.js

@@ -15,15 +15,31 @@ var FloeGuardError = class extends Error {
 var BudgetExceeded = class extends FloeGuardError {
   spentUsd;
   limitUsd;
-  constructor(spentUsd, limitUsd) {
+  constructor(spentUsd, limitUsd, message) {
     super(
-      `BUDGET EXCEEDED \u2014 call blocked (spent $${spentUsd.toFixed(6)} of $${limitUsd.toFixed(6)} ceiling)`
+      message ?? `BUDGET EXCEEDED \u2014 call blocked (spent $${spentUsd.toFixed(6)} of $${limitUsd.toFixed(6)} ceiling)`
     );
     this.name = "BudgetExceeded";
     this.spentUsd = spentUsd;
     this.limitUsd = limitUsd;
   }
 };
+var TokenBudgetExceeded = class extends BudgetExceeded {
+  spentTokens;
+  limitTokens;
+  scope;
+  constructor(spentTokens, limitTokens, scope) {
+    super(
+      0,
+      0,
+      `TOKEN BUDGET EXCEEDED \u2014 call blocked (${scope}: ${spentTokens} of ${limitTokens} token ceiling)`
+    );
+    this.name = "TokenBudgetExceeded";
+    this.spentTokens = spentTokens;
+    this.limitTokens = limitTokens;
+    this.scope = scope;
+  }
+};
 var UnpriceableModelError = class extends FloeGuardError {
   model;
   constructor(model) {
@@ -1088,12 +1104,23 @@ function priceTokens(priced, promptTokens, completionTokens) {
 
 // src/guard.ts
 var EPS = 1e-12;
+function isReservation(h) {
+  return typeof h === "object" && h !== null && typeof h.usd === "number" && typeof h.tokens === "number";
+}
 var BudgetGuard = class {
   limitUsd;
   spentUsd = 0;
   priceOverrides;
   failClosed;
   nearLimitBps;
+  /** Aggregate token ceiling, or null when the token dimension is disabled. */
+  tokenLimit;
+  /** Aggregate tokens accrued (prompt + completion + cache buckets). */
+  spentTokens = 0;
+  /** Tokens held in-flight — the token twin of `reserved`. */
+  reservedTokens = 0;
+  /** Step stack: innermost step is last. Empty when no step() is active. */
+  steps = [];
   onBlock;
   /**
    * Costs of the most recent priced LLM call and tool call, tracked
@@ -1135,7 +1162,13 @@ var BudgetGuard = class {
         `maxLogEvents must be a non-negative integer, got ${options.maxLogEvents}`
       );
     }
+    if (options.tokenLimit !== void 0 && (!Number.isInteger(options.tokenLimit) || options.tokenLimit < 0)) {
+      throw new RangeError(
+        `tokenLimit must be a non-negative integer, got ${options.tokenLimit}`
+      );
+    }
     this.limitUsd = limitUsd;
+    this.tokenLimit = options.tokenLimit ?? null;
     this.maxLogEvents = options.maxLogEvents;
     this.priceOverrides = options.priceOverrides;
     this.failClosed = options.failClosed ?? true;
@@ -1155,7 +1188,7 @@ var BudgetGuard = class {
    * Note: `check` is a non-binding peek. For parallel calls, use `reserve()` /
    * `settle()`, which hold the estimate across the await.
    */
-  check(estimatedNextCost) {
+  check(estimatedNextCost, options = {}) {
     const rawEstimate = estimatedNextCost === void 0 ? this.defaultEstimate() : estimatedNextCost;
     if (!Number.isFinite(rawEstimate)) {
       throw new RangeError(
@@ -1163,11 +1196,10 @@ var BudgetGuard = class {
       );
     }
     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.validateTokenEstimate(options.estimatedTokens);
+    const tokens = Math.max(0, options.estimatedTokens ?? 0);
+    const blocked = this.blockingCross(estimate, tokens);
+    if (blocked !== null) this.raiseBlock(blocked);
   }
   /**
    * Atomically check the ceiling AND hold the estimated cost in flight.
@@ -1180,7 +1212,7 @@ var BudgetGuard = class {
    * `estimatedCost` defaults to the costlier of the last LLM call and the last
    * tool call.
    */
-  reserve(estimatedCost) {
+  reserve(estimatedCost, options = {}) {
     const rawEstimate = estimatedCost === void 0 ? this.defaultEstimate() : estimatedCost;
     if (!Number.isFinite(rawEstimate)) {
       throw new RangeError(
@@ -1188,13 +1220,19 @@ var BudgetGuard = class {
       );
     }
     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.validateTokenEstimate(options.estimatedTokens);
+    const tokens = Math.max(0, options.estimatedTokens ?? 0);
+    const blocked = this.blockingCross(estimate, tokens);
+    if (blocked !== null) this.raiseBlock(blocked);
     this.reserved += estimate;
-    return estimate;
+    this.reservedTokens += tokens;
+    const step = this.steps.length ? this.steps[this.steps.length - 1] : null;
+    if (step !== null) {
+      step.reservedUsd += estimate;
+      step.reservedTokens += tokens;
+    }
+    if (tokens === 0 && step === null) return estimate;
+    return { usd: estimate, tokens };
   }
   /**
    * Release a reservation and record the actual cost. `record` is `settle` with
@@ -1207,9 +1245,7 @@ var BudgetGuard = class {
    */
   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}`);
-    }
+    const reservedUsd = this.reservedUsdOf(reserved);
     let overrides = this.priceOverrides;
     if (options.price !== void 0) {
       overrides = { ...overrides ?? {}, [model]: options.price };
@@ -1235,7 +1271,10 @@ var BudgetGuard = class {
     if (reserved) {
       this.consumeReservation(reserved);
     }
+    const accruedTokens = Math.max(0, promptTokens) + Math.max(0, completionTokens);
     this.spentUsd += cost;
+    this.spentTokens += accruedTokens;
+    this.accrueStep(cost, accruedTokens);
     if (this.spentUsd - this.limitUsd > 0 && this.spentUsd - this.limitUsd < EPS) {
       this.spentUsd = this.limitUsd;
     }
@@ -1249,8 +1288,9 @@ var BudgetGuard = class {
       costUsd: cost,
       ...options.label !== void 0 ? { label: options.label } : {},
       // 0 means "no reservation" (the plain record() path) — omit rather than
-      // log a meaningless zero.
-      ...reserved ? { reserved } : {}
+      // log a meaningless zero. Log the USD amount so the ledger schema stays
+      // number-shaped even for a BudgetReservation.
+      ...reservedUsd ? { reserved: reservedUsd } : {}
     });
     return cost;
   }
@@ -1315,13 +1355,12 @@ var BudgetGuard = class {
       throw new RangeError(`costUsd must be a finite, non-negative number, got ${costUsd}`);
     }
     const reserved = options.reserved ?? 0;
-    if (!Number.isFinite(reserved) || reserved < 0) {
-      throw new RangeError(`reserved must be a finite, non-negative number, got ${reserved}`);
-    }
+    const reservedUsd = this.reservedUsdOf(reserved);
     if (reserved) {
       this.consumeReservation(reserved);
     }
     this.spentUsd += costUsd;
+    this.accrueStep(costUsd, 0);
     if (this.spentUsd - this.limitUsd > 0 && this.spentUsd - this.limitUsd < EPS) {
       this.spentUsd = this.limitUsd;
     }
@@ -1335,7 +1374,7 @@ var BudgetGuard = class {
       completionTokens: null,
       costUsd,
       ...options.label !== void 0 ? { label: options.label } : {},
-      ...reserved ? { reserved } : {}
+      ...reservedUsd ? { reserved: reservedUsd } : {}
     });
     return costUsd;
   }
@@ -1355,9 +1394,7 @@ var BudgetGuard = class {
    * 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}`);
-    }
+    this.reservedUsdOf(reserved);
     if (!reserved) return;
     this.consumeReservation(reserved);
   }
@@ -1429,12 +1466,171 @@ var BudgetGuard = class {
    * responsibility.
    */
   consumeReservation(reserved) {
-    if (reserved > this.reserved + EPS) {
+    const usd = isReservation(reserved) ? reserved.usd : reserved;
+    const tokens = isReservation(reserved) ? reserved.tokens : 0;
+    if (usd > this.reserved + EPS) {
       throw new RangeError(
-        `reserved handle (${reserved}) exceeds total in-flight reservations (${this.reserved}) \u2014 a handle must come from a matching reserve()`
+        `reserved handle (${usd}) exceeds total in-flight reservations (${this.reserved}) \u2014 a handle must come from a matching reserve()`
       );
     }
-    this.reserved = Math.max(0, this.reserved - reserved);
+    if (tokens > this.reservedTokens) {
+      throw new RangeError(
+        `reserved token handle (${tokens}) exceeds total in-flight token reservations (${this.reservedTokens}) \u2014 a handle must come from a matching reserve()`
+      );
+    }
+    this.reserved = Math.max(0, this.reserved - usd);
+    this.reservedTokens = Math.max(0, this.reservedTokens - tokens);
+    const step = this.steps.length ? this.steps[this.steps.length - 1] : null;
+    if (step !== null) {
+      step.reservedUsd = Math.max(0, step.reservedUsd - usd);
+      step.reservedTokens = Math.max(0, step.reservedTokens - tokens);
+    }
+  }
+  /**
+   * The USD amount of a handle, validated. Both shapes are checked — a raw
+   * number and a {@link BudgetReservation}'s `usd`/`tokens` fields — so a bad
+   * hand-rolled handle can't corrupt the in-flight tally.
+   */
+  /**
+   * Validate a caller-supplied token estimate. Rejects a fraction, NaN,
+   * Infinity, or boolean (via `Number.isInteger`) so it can't disable the
+   * integer token hard-stops or corrupt the in-flight token tally. `undefined`
+   * (the default) is fine; a negative integer is clamped by `Math.max(0, ...)`,
+   * matching the lenient USD estimate. Mirrors Python's `_validate_token_estimate`.
+   */
+  validateTokenEstimate(estimatedTokens) {
+    if (estimatedTokens !== void 0 && !Number.isInteger(estimatedTokens)) {
+      throw new RangeError(`estimatedTokens must be an integer, got ${estimatedTokens}`);
+    }
+  }
+  reservedUsdOf(reserved) {
+    if (isReservation(reserved)) {
+      if (!Number.isFinite(reserved.usd) || reserved.usd < 0) {
+        throw new RangeError(
+          `reserved.usd must be a finite, non-negative number, got ${reserved.usd}`
+        );
+      }
+      if (!Number.isInteger(reserved.tokens) || reserved.tokens < 0) {
+        throw new RangeError(
+          `reserved.tokens must be a non-negative integer, got ${reserved.tokens}`
+        );
+      }
+      return reserved.usd;
+    }
+    if (!Number.isFinite(reserved) || reserved < 0) {
+      throw new RangeError(`reserved must be a finite, non-negative number, got ${reserved}`);
+    }
+    return reserved;
+  }
+  /**
+   * Accrue a settled call into the innermost active step (no-op when none).
+   * Sequential-loop contract: the innermost step owns the call.
+   */
+  accrueStep(cost, tokens) {
+    const step = this.steps.length ? this.steps[this.steps.length - 1] : null;
+    if (step !== null) {
+      step.spentUsd += cost;
+      step.spentTokens += tokens;
+    }
+  }
+  /**
+   * The ONE choke point, across two dimensions and two scopes. Returns the first
+   * ceiling that blocks as `[dimension, scope]` — dimension `"usd" | "tokens"`,
+   * scope `"aggregate" | "step"` — or `null` if the call fits everywhere.
+   * Aggregate is checked before step (the guard-wide ceiling is the hard limit).
+   */
+  blockingCross(estimateUsd, estimateTokens) {
+    const committed = this.spentUsd + this.reserved;
+    if (committed > this.limitUsd - EPS || committed + estimateUsd > this.limitUsd + EPS) {
+      return ["usd", "aggregate"];
+    }
+    if (this.tokenLimit !== null) {
+      const committedT = this.spentTokens + this.reservedTokens;
+      if (committedT >= this.tokenLimit || committedT + estimateTokens > this.tokenLimit) {
+        return ["tokens", "aggregate"];
+      }
+    }
+    const step = this.steps.length ? this.steps[this.steps.length - 1] : null;
+    if (step !== null) {
+      if (step.maxUsd !== null) {
+        const sCommitted = step.spentUsd + step.reservedUsd;
+        if (sCommitted > step.maxUsd - EPS || sCommitted + estimateUsd > step.maxUsd + EPS) {
+          return ["usd", "step"];
+        }
+      }
+      if (step.maxTokens !== null) {
+        const sCommittedT = step.spentTokens + step.reservedTokens;
+        if (sCommittedT >= step.maxTokens || sCommittedT + estimateTokens > step.maxTokens) {
+          return ["tokens", "step"];
+        }
+      }
+    }
+    return null;
+  }
+  /** Notify + throw the right error for a [dimension, scope] block. */
+  raiseBlock(blocked) {
+    const [dimension, scope] = blocked;
+    if (dimension === "usd") {
+      this.onBlock(this.spentUsd, this.limitUsd);
+      throw new BudgetExceeded(this.spentUsd, this.limitUsd);
+    }
+    let spentT;
+    let limitT;
+    if (scope === "step") {
+      const step = this.steps[this.steps.length - 1];
+      spentT = step.spentTokens + step.reservedTokens;
+      limitT = step.maxTokens ?? 0;
+    } else {
+      spentT = this.spentTokens + this.reservedTokens;
+      limitT = this.tokenLimit ?? 0;
+    }
+    throw new TokenBudgetExceeded(spentT, limitT, scope);
+  }
+  /**
+   * Scope a per-step USD and/or token cap for a **sequential agent loop**.
+   *
+   * Runs `fn` with a step pushed onto the guard's stack; the
+   * enforcement/accrual path honours the innermost active step *on top of* the
+   * aggregate ceilings. A call that would cross the step's `maxUsd` / `maxTokens`
+   * is hard-blocked ({@link BudgetExceeded} / {@link TokenBudgetExceeded} with
+   * `scope: "step"`) even if the aggregate budget has room. `fn` receives the
+   * SAME guard, so no adapter needs to know about steps. The step is popped when
+   * `fn` settles or throws. **Not for concurrent parallel steps on one guard** —
+   * that's out of scope for issue #46; use one guard per parallel branch.
+   */
+  step(options, fn) {
+    const { maxUsd, maxTokens } = options;
+    if (maxUsd !== void 0 && (!Number.isFinite(maxUsd) || maxUsd < 0)) {
+      throw new RangeError(`maxUsd must be a finite, non-negative number, got ${maxUsd}`);
+    }
+    if (maxTokens !== void 0 && (!Number.isInteger(maxTokens) || maxTokens < 0)) {
+      throw new RangeError(`maxTokens must be a non-negative integer, got ${maxTokens}`);
+    }
+    const state = {
+      maxUsd: maxUsd ?? null,
+      maxTokens: maxTokens ?? null,
+      spentUsd: 0,
+      spentTokens: 0,
+      reservedUsd: 0,
+      reservedTokens: 0
+    };
+    this.steps.push(state);
+    const pop = () => {
+      const idx = this.steps.lastIndexOf(state);
+      if (idx !== -1) this.steps.splice(idx, 1);
+    };
+    let result;
+    try {
+      result = fn(this);
+    } catch (err) {
+      pop();
+      throw err;
+    }
+    if (result != null && typeof result.then === "function") {
+      return result.finally(pop);
+    }
+    pop();
+    return result;
   }
   appendEvent(event) {
     this.spendEvents.push(Object.freeze(event));
@@ -1451,17 +1647,63 @@ var BudgetGuard = class {
    */
   advisory() {
     const usedBps = this.limitUsd <= 0 ? 1e4 : Math.max(0, Math.min(1e4, Math.floor(this.spentUsd / this.limitUsd * 1e4 + 1e-9)));
+    const remainingUsd = Math.max(0, this.limitUsd - this.spentUsd);
+    const expectedCost = Math.max(this.lastLlmCost, this.lastToolCost);
+    const estCallsRemaining = expectedCost > 0 ? Math.floor(remainingUsd / expectedCost + 1e-9) : null;
+    let tokenUsedBps = null;
+    let remainingTokens = null;
+    let nearToken = false;
+    if (this.tokenLimit !== null) {
+      tokenUsedBps = this.tokenLimit <= 0 ? 1e4 : Math.max(
+        0,
+        Math.min(1e4, Math.floor(this.spentTokens / this.tokenLimit * 1e4 + 1e-9))
+      );
+      remainingTokens = Math.max(0, this.tokenLimit - this.spentTokens);
+      nearToken = tokenUsedBps >= this.nearLimitBps;
+    }
+    const step = this.steps.length ? this.steps[this.steps.length - 1] : null;
+    let stepRemainingUsd = null;
+    let stepRemainingTokens = null;
+    let nearStep = false;
+    if (step !== null) {
+      if (step.maxUsd !== null) {
+        stepRemainingUsd = Math.max(0, step.maxUsd - step.spentUsd);
+        if (step.maxUsd <= 0) {
+          nearStep = true;
+        } else {
+          const sBps = Math.floor(step.spentUsd / step.maxUsd * 1e4 + 1e-9);
+          nearStep = nearStep || sBps >= this.nearLimitBps;
+        }
+      }
+      if (step.maxTokens !== null) {
+        stepRemainingTokens = Math.max(0, step.maxTokens - step.spentTokens);
+        if (step.maxTokens <= 0) {
+          nearStep = true;
+        } else {
+          const sTBps = Math.floor(step.spentTokens / step.maxTokens * 1e4 + 1e-9);
+          nearStep = nearStep || sTBps >= this.nearLimitBps;
+        }
+      }
+    }
     return {
-      nearLimit: usedBps >= this.nearLimitBps,
+      // nearLimit now also flips when a token ceiling or the active step is near
+      // its cap, so a router can downshift before ANY hard-stop.
+      nearLimit: usedBps >= this.nearLimitBps || nearToken || nearStep,
       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),
+      remainingUsd,
       limitUsd: this.limitUsd,
       spentUsd: this.spentUsd,
-      scope: "local"
+      scope: "local",
+      expectedCost,
+      estCallsRemaining,
+      tokenUsedBps,
+      remainingTokens,
+      stepRemainingUsd,
+      stepRemainingTokens
     };
   }
 };
@@ -1666,6 +1908,7 @@ export {
   DeadlineExceeded,
   FloeGuardError,
   LatencyBudget,
+  TokenBudgetExceeded,
   UnpriceableModelError,
   budgetGuardMiddleware,
   priceTokens,