floe-guard (Vercel AI SDK)

A local budget guardrail for AI agents — the TypeScript counterpart to the Python floe-guard. It hard-stops your agent before its next LLM or paid tool call when it would cross a USD spend ceiling — tokens and tool calls under one local ceiling. No account, no signup, no network. Runs in your process.

Works with both AI SDK v4 and v5 (ai@4 / ai@5).

npm i floe-guard ai@5 @ai-sdk/openai@2

This example uses AI SDK 5 with its compatible OpenAI provider, version 2. For AI SDK 4, use ai@4 @ai-sdk/openai@1. Keep the SDK and provider major versions compatible; installing the latest provider can select a model API that this middleware does not support.

import { wrapLanguageModel } from "ai";
import { openai } from "@ai-sdk/openai";
import { BudgetGuard, budgetGuardMiddleware } from "floe-guard";

const guard = new BudgetGuard(5.0); // your ceiling, in USD

const model = wrapLanguageModel({
  model: openai("gpt-4o"),
  middleware: budgetGuardMiddleware(guard),
});
// generateText / streamText with `model` now stop at $5 — the call that would
// cross the ceiling throws `BudgetExceeded` BEFORE it runs.

The middleware sits in the call path: it check()s before doGenerate / doStream (throwing BudgetExceeded to halt the run) and record()s priced token usage after — for streaming it reads usage from the finish part.

Mid-stream budget enforcement

StreamGuard and guardStream port Python's streaming USD guard. They price each text delta before it reaches the consumer. When a chunk crosses the shared ceiling, its partial spend is recorded in guard.spendLog before BudgetExceeded is thrown. Active streams on the same guard share the ceiling.

import { BudgetGuard, guardStream } from "floe-guard";

const guard = new BudgetGuard(0.01);
// textDeltas is your provider's Iterable<string> or AsyncIterable<string>.
for await (const text of guardStream(guard, "gpt-4o", textDeltas)) {
  consume(text);
}

For structured chunks, supply { getText: chunk => chunk.delta.text ?? "" } using your provider's actual shape. Without an extractor, non-string chunks throw instead of silently recording zero. Synchronous inputs return a synchronous iterator; asynchronous inputs return an asynchronous iterator.

Use StreamGuard directly when the provider reports final usage:

import { StreamGuard } from "floe-guard";

const reserved = guard.reserve(guard.estimateCall("gpt-4o", 100, 200));
const stream = new StreamGuard(guard, "gpt-4o", { promptTokens: 100, reserved });
try {
  for await (const text of textDeltas) {
    stream.feedText(text); // or feedTokens(n) with known per-chunk counts
    consume(text);
  }
  stream.finish({ promptTokens: reportedPromptTokens, completionTokens: reportedCompletionTokens });
} finally {
  stream.close(); // idempotent; settles estimates if finish() was not reached
}

• Options: promptTokens, reserved, price, label, and countTokens(delta). guardStream additionally accepts getText(chunk). Configuration is captured at construction; changing the caller's options or manual price object later does not change an existing stream.
• approxTokens defaults to roughly four Unicode characters per token, with a minimum of one for a non-empty delta. A custom tokenizer can replace it. finish() reconciles to reported usage; completionTokens exposes the running estimate. Mid-stream checks cover the aggregate USD ceiling, matching Python; token reservations are reconciled at settlement.
• An unpriceable model fails at construction and releases its reservation. With failClosed: false, unpriceable streams pass through and settlement warns and skips accounting, matching the existing guard policy.
• The wrapper settles on exhaustion, source/consumer errors, and early break, and closes the source iterator when iteration ends early. Direct users must call close() in finally. If a wrapper is never iterated, its reservation remains the caller's responsibility: call guard.release(reserved).
• The crossing chunk has already been generated. With accurate counts, the local cutoff can overshoot by that chunk; heuristic error, parallel streams, provider buffering and delayed cancellation can cause additional actual spend. Stopping iteration requests iterator cleanup, not guaranteed remote cancellation. Connect cleanup to your provider's abort mechanism where needed.
• Existing Vapi, LiveKit and middleware streaming behavior is unchanged; this is the standalone primitive requested in issue #124.

Run the no-key example from a repository checkout:

cd js
npm ci
npm run build
node ../examples/streaming_guard.mjs

Pricing

Tokens are priced offline from a bundled LiteLLM cost map. A model that isn't in the map (and has no manual price) fails closed: record throws UnpriceableModelError rather than silently treating spend as free — you can't cap spend you can't measure.

const guard = new BudgetGuard(5.0, {
  priceOverrides: {
    "my-self-hosted-model": { inputCostPerToken: 1e-6, outputCostPerToken: 2e-6 },
  },
  // or failClosed: false to warn-and-skip for models you accept un-metered.
});

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.

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");

Request-sized estimates

To ensure the ceiling is enforced on the first run or for a call much larger than the previous one, you can price the actual incoming request using estimateCall() and pass the estimate to reserve() or check():

const est = guard.estimateCall("gpt-4o", 12_000, 4_096);
const handle = guard.reserve(est); // throws BudgetExceeded NOW if this call alone would cross
try {
  const response = await callYourLlm({ model: "gpt-4o", ... });
  guard.settle("gpt-4o", response.usage.promptTokens, response.usage.completionTokens, { reserved: handle });
} catch (err) {
  guard.release(handle);
  throw err;
}

If the model is unpriceable, estimateCall() returns undefined and reserve(undefined) / check(undefined) fall back gracefully to the last-cost prediction.

Tool spend under the same ceiling

Paid tool calls (Apollo, Exa, scrapers) burn the same budget as tokens. The full reserve/settle contract applies — and the price is known before the call, so the pre-call hard-stop is exact:

const handle = guard.reserveTool(0.02); // throws BudgetExceeded BEFORE the call
const result = await apollo.peopleLookup(...);
guard.settleTool("apollo.people_lookup", 0.02, { reserved: handle });

guard.recordTool("exa.search", 0.004); // post-hoc, for metered APIs
guard.toolCosts; // { "apollo.people_lookup": 0.42, "exa.search": 0.11 }

Token ceilings and per-step budgets

Cap total token usage — every bucket the guard counts: prompt, completion, and cache (a token ceiling) and keep one step of a sequential loop from starving the rest (a per-step cap) — a second dimension on the same reserve/settle machinery, not a second guard:

import { BudgetGuard, TokenBudgetExceeded } from "floe-guard";

// aggregate token ceiling alongside the USD one
const guard = new BudgetGuard(100, { tokenLimit: 20_000 });
guard.check(undefined, { estimatedTokens: 1_200 }); // throws TokenBudgetExceeded if it'd cross
guard.record("gpt-4o", 800, 400); // tokens accrue for free from the counts

// a per-step cap for one step of a sequential loop (callback — g IS guard)
guard.step({ maxTokens: 5_000 }, (g) => {
  g.record("gpt-4o", 3_000, 1_500);
  g.check(undefined, { estimatedTokens: 1_000 }); // 4_500 + 1_000 > 5_000 → scope "step"
});

const adv = guard.advisory();
adv.tokenUsedBps; // aggregate token utilization (null if no tokenLimit)
adv.remainingTokens; // tokens left before the ceiling (null if no tokenLimit)
adv.stepRemainingTokens; // active step's headroom (null if no step, or its token cap is unset)

TokenBudgetExceeded extends BudgetExceeded, so budget-aware retry treats a token block as terminal automatically. With no tokenLimit and no step(), USD enforcement is unchanged and reserve() still returns a plain number — a BudgetReservation handle appears only when tokens are actually reserved or a step is active. (advisory() gains the token/step fields above; additive and null when their dimension is unused.)

Per-call spend log

The guard keeps a typed, in-memory ledger of everything it priced: each record() / settle() appends one SpendEvent, and recordTool() lets paid non-LLM calls spend the same budget and land in the same log. The events sum to spentUsd (unless a maxLogEvents ring buffer has evicted old ones).

const guard = new BudgetGuard(1.0); // { maxLogEvents: N } caps memory
guard.record("gpt-4o", 1_200, 350, { label: "researcher" });
guard.recordTool("serpapi.search", 0.01, { label: "researcher" });

guard.spendLog; // [{ timestamp, kind: "llm", modelOrTool: "gpt-4o", … }, …]
process.stdout.write(guard.exportLog()); // JSONL, one event per line

exportLog() emits a stable snake_case schema — {timestamp, kind: llm|tool, model_or_tool, prompt_tokens, completion_tokens, cost_usd, label?, reserved?} — identical to the Python package's export_log(), so every agent produces the same shape regardless of stack.

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.

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

npm install
npm run build
npm test
npm run typecheck

License

MIT — see ../LICENSE.