Skip to content
reaatechREAATECH

@reaatech/agent-budget-middleware

npm v0.1.0

Enforces LLM spending limits by providing Express/Fastify middleware and a `BudgetInterceptor` class that dynamically filters tools, downgrades models, and records token usage. It requires a `BudgetController` instance from the `@reaatech/agent-budget-engine` package to manage budget state and policy logic.

@reaatech/agent-budget-middleware

npm version License: MIT CI

Status: Pre-1.0 — APIs may change in minor versions. Pin to a specific version in production.

Express/Fastify-compatible middleware and a direct SDK wrapper (BudgetInterceptor) that integrates budget enforcement into any agent server. Injects budget response headers, intercepts model selection for auto-downgrade, and filters expensive tools — all before the LLM call reaches your application code.

Installation

terminal
npm install @reaatech/agent-budget-middleware
# or
pnpm add @reaatech/agent-budget-middleware

Feature Overview

  • BudgetInterceptor — direct SDK wrapper for non-HTTP agents (worker threads, CLI agents, job queues)
  • createBudgetMiddleware — Express/Fastify-compatible HTTP middleware factory
  • Pre-request budget check — validates estimated cost against remaining budget before the LLM call
  • Auto-downgrade interception — swaps model selection if budget has tightened
  • Tool filtering — removes expensive tools from the request when budget is constrained
  • Post-request spend recording — records actual token usage and cost after the LLM call
  • Budget response headers — injects X-Budget-Remaining, X-Budget-Status, X-Budget-Suggested-Model
  • Scope extraction from HTTP — reads x-budget-scope-type and x-budget-scope-key headers

Quick Start

With Express

typescript
import express from 'express';
import { createBudgetMiddleware } from '@reaatech/agent-budget-middleware';
import { BudgetController } from '@reaatech/agent-budget-engine';
import { SpendStore } from '@reaatech/agent-budget-spend-tracker';
import { BudgetScope } from '@reaatech/agent-budget-types';
 
const store = new SpendStore();
const controller = new BudgetController({ spendTracker: store });
 
await controller.defineBudget({
  scopeType: BudgetScope.User,
  scopeKey: 'user-42',
  limit: 10.0,
  policy: { softCap: 0.8, hardCap: 1.0 },
});
 
const middleware = createBudgetMiddleware({ controller });
 
const app = express();
app.use(express.json());
app.use('/agent', middleware.beforeStep);
app.use('/agent', middleware.afterStep);
 
app.post('/agent', (req, res) => {
  // req.budgetContext contains { scope, allowed, modelId, tools, ... }
  if (!req.budgetContext.allowed) {
    return res.status(402).json({ error: 'Budget exceeded' });
  }
  // Use req.budgetContext.modelId and req.budgetContext.tools for the LLM call
  res.json({ status: 'ok' });
});
 
app.listen(3000);

With Fastify

typescript
import Fastify from 'fastify';
import { createBudgetMiddleware } from '@reaatech/agent-budget-middleware';
 
const fastify = Fastify();
const middleware = createBudgetMiddleware({ controller });
 
fastify.addHook('preHandler', middleware.beforeStep);
fastify.addHook('onResponse', middleware.afterStep);

Direct SDK Usage (non-HTTP)

typescript
import { BudgetInterceptor } from '@reaatech/agent-budget-middleware';
import { BudgetController } from '@reaatech/agent-budget-engine';
import { SpendStore } from '@reaatech/agent-budget-spend-tracker';
import { BudgetScope } from '@reaatech/agent-budget-types';
 
const store = new SpendStore();
const controller = new BudgetController({ spendTracker: store });
const interceptor = new BudgetInterceptor({ controller });
 
// Before LLM call
const ctx = await interceptor.beforeStep({
  scopeType: BudgetScope.User,
  scopeKey: 'user-42',
  modelId: 'claude-opus-4-1',
  estimatedCost: 0.15,
  tools: ['web-search', 'code-interpreter'],
});
 
if (!ctx.allowed) throw new Error(`Budget exceeded: ${ctx.reason}`);
 
// Use ctx.modelId and ctx.tools for the actual LLM call
const llmResult = await callLLM(ctx.modelId, ctx.tools);
 
// After LLM call — record actual spend
await interceptor.afterStep({
  ...ctx,
  actualCost: llmResult.cost,
  inputTokens: llmResult.inputTokens,
  outputTokens: llmResult.outputTokens,
  requestId: 'req-abc123',
});

API Reference

BudgetInterceptor

MethodDescription
beforeStep(params)Pre-flight budget check. Returns InterceptorContext with potentially downgraded model and filtered tools. Throws BudgetExceededError if the request is blocked.
afterStep(params)Records actual spend after the LLM call completes. Accepts InterceptorAfterContext with actual cost and token counts.

BudgetInterceptorOptions

typescript
interface BudgetInterceptorOptions {
  controller: BudgetController;
}

createBudgetMiddleware

Factory function that returns Express/Fastify middleware handlers.

typescript
function createBudgetMiddleware(options: { controller: BudgetController }): {
  beforeStep: (req, res, next) => Promise<void>;
  afterStep: (req, res, next) => Promise<void>;
};

The beforeStep handler reads scope from HTTP headers:

HeaderValue
x-budget-scope-typeOne of task, user, session, org
x-budget-scope-keyThe scope identifier (e.g., user-42)

It injects the following response headers:

HeaderDescription
X-Budget-RemainingDollars remaining in the budget
X-Budget-StatusCurrent state: active, warned, degraded, stopped
X-Budget-LimitTotal budget limit
X-Budget-SpentDollars spent so far
X-Budget-Suggested-ModelModel to use (if downgrade was applied)

The afterStep handler records the actual spend and cleans up context.

Usage Patterns

Custom Scope Extraction

typescript
const interceptor = new BudgetInterceptor({ controller });
const scope = { scopeType: BudgetScope.Org, scopeKey: jwtPayload.orgId };
 
const ctx = await interceptor.beforeStep({
  ...scope,
  modelId: 'claude-sonnet-4',
  estimatedCost: 0.05,
});

Advanced: Progressive Response

typescript
app.post('/agent', middleware.beforeStep, (req, res, next) => {
  const ctx = req.budgetContext;
 
  if (!ctx.allowed) {
    return res.status(402).json({
      error: 'Budget exhausted',
      budget: {
        limit: res.get('X-Budget-Limit'),
        spent: res.get('X-Budget-Spent'),
        suggestedModel: res.get('X-Budget-Suggested-Model'),
      },
    });
  }
 
  if (ctx.warning) {
    res.set('X-Budget-Warning', ctx.warning);
  }
 
  next();
}, middleware.afterStep);

Error Handling

The interceptor throws typed errors from @reaatech/agent-budget-types:

ErrorWhen
BudgetExceededErrorThe request would exceed the hard cap
BudgetValidationErrorMissing or invalid scope identifiers
typescript
try {
  const ctx = await interceptor.beforeStep({ ... });
} catch (err) {
  if (err instanceof BudgetExceededError) {
    res.status(402).json({ error: err.message, remaining: err.remaining });
  }
}

License

MIT