How do you enforce token and spend budgets on AI agents?
You put a token bucket at the tool call rather than at the HTTP entry point, keyed on the identity that you are protecting. Each call draws the bucket down in proportion to its cost. Arcjet shares the bucket by key, not by correlationId. Key on a user and every tool call for that user draws from the same bucket, including concurrent runs. You are limiting total consumption by that identity, not how many times someone starts a workflow.
A runaway agent is a budget problem before it is a security problem, and it is the failure mode that most production agents meet first.
The first agent incident is usually a bill
Before anyone exploits your agent, someone runs up an unexpected bill with it. The usual causes are ordinary:
- A loop that does not stop
- A user calling an expensive tool repeatedly
- An injection that turns one request into 200
Cost explosion is the most common way production agents fail, and it is the one your finance team notices first. It is also containable, provided the limit sits where the spend happens.
Why per-request limits do not work for agents
A rate limit on your HTTP endpoint caps how many times a user can start a workflow. It says nothing about how much that workflow consumes once running. One request can trigger 50 tool calls. The limit fires once, and the spend happens 50 times.
| Limit placed at | What it counts | What a runaway loop costs you |
|---|---|---|
| HTTP entry point | Workflow starts | Unbounded, because one start can fan out arbitrarily |
| Model provider account | Total spend across every tenant | Bounded, but the cap is global, so one bad actor degrades everyone |
| The tool call, keyed per identity | Actual consumption by that identity | Bounded per identity, with no collateral damage |
The budget has to sit where the spend happens.
Token buckets at the tool
A token bucket refills at a set rate, and each call draws from it in proportion to cost. Callers within their allowance never notice. A runaway loop is throttled once the bucket empties. There is no global cap penalizing everyone.
import { launchArcjet, tokenBucket } from "@arcjet/guard";
const arcjet = launchArcjet({ key: process.env.ARCJET_KEY! });
const budget = tokenBucket({ refillRate: 2_000, intervalSeconds: 3_600, maxTokens: 5_000,});
const decision = await arcjet.guard({ label: "tools.generate-report", actor: session.userId, correlationId: workflowRunId, rules: [ budget({ key: session.userId, requested: estimatedTokens, metadata: { model: "claude-opus-5" }, }), ],});
if (decision.hasFailedOpen()) { throw new Error("Budget check unavailable");}if (decision.conclusion === "DENY" && decision.reason === "RATE_LIMIT") { // Over budget, surface a retry-after rather than a generic failure}Because Arcjet shares the bucket by key, the rule constrains total consumption for that identity rather than only the HTTP entry point. The sample keys on session.userId, so it is per-user, including concurrent runs. correlationId tags the workflow for later inspection; it does not scope the bucket. If you want a per-run budget, pass the run ID as key and be explicit that that is what you are protecting.
The requested value is where cost proportionality happens. Charge an expensive model or a large context more tokens than a cheap lookup, and the same bucket limits spend rather than call count.
Get the key right
The key decides what you are protecting.
- Per user or session is the default for anything customer-facing, and the right starting point when you are unsure. The preceding sample does this: one bucket for the user, across every tool call.
- Per run is a different choice. Key on the workflow run ID if you want one loop bounded independently of the user's other sessions.
- Per plan or tier is right where budget is part of what someone bought, so the limit expresses an entitlement rather than a safety valve.
- Per deployment or instance is for a single-tenant worker or a stdio MCP server with no user context. Use a stable identifier that you control, and be explicit that this is what you are keying on, because it is easy to assume a per-user limit exists where none does.
A limit keyed on something that the caller controls is not a limit. If the key comes from a request header or a model-supplied argument, an attacker rotates it and the bucket never empties.
What happens when the limit is reached
Return a distinguishable error. "Rate limited, retry in 12 seconds" and "input flagged as prompt injection" are different conditions, and the caller must be able to tell them apart. That caller might be a user, a retrying client, or another agent. Acting on decision.reason rather than a generic failure is what makes that possible.
For agent workflows specifically, consider whether exceeding the budget ends the run or degrades it. A research agent that reaches its ceiling can often return partial results usefully. A payment workflow that reaches its ceiling mid-sequence must stop, because a half-completed sequence is worse than a refused one.
Operate budgets over time
Budgets change. A customer upgrades. An incident calls for temporary tightening. A different model changes the cost per call.
Token-bucket refill and capacity live in application code. They cannot be authored as HTTP remote rules. HTTP remote rules on protect() (bots, Shield, filters, fixed-window and sliding-window rate limits) are site-wide and do not support token bucket. Guard remote policies do not include rate limits either. What a security team can change without a deploy on a Guard path is allow/deny lists, length, membership, prompt injection, and local sensitive-info.
Watch the effect of a code-defined budget in the Arcjet Console. Start in dry run if you are unsure where the ceiling belongs, because dry run shows what would have been throttled against real traffic rather than making you guess. Arcjet records dry-run results without changing the allow/deny conclusion, so the measurement is evidence rather than an impression.
You decide where the budget check sits, because that is part of building the tool. Token-bucket numbers stay in code, so a budget change ships with the application.
Checklist
Check the following before you enforce a budget:
- Put the budget at the tool call. An HTTP entry point limit counts workflow starts, not consumption.
- Key it on a stable identity that the caller cannot control. The sample keys on user ID (per-user, including concurrent runs); key on the run ID only if you want a per-run cap.
- Make
requestedproportional to real cost rather than one per call. - Remember sharing is by
key, not bycorrelationId. - Return a distinguishable error with retry information.
- On sensitive tools, deny when
hasFailedOpen()is true; otherwise a timeout lets the call proceed. - Decide whether exceeding the budget ends or degrades the workflow.
- Deploy in dry run, measure against real traffic, then enforce.
- Keep Guard remote-policy changes (lists, injection, local PII) out of the release cycle; budget numbers stay in code.
Where Arcjet fits
Budget enforcement sits awkwardly between categories. FinOps tooling reports spend after it happens. Model gateways cap tokens at the provider boundary, which bounds the bill but not per-tenant fairness. Neither is positioned to stop the fiftieth tool call in a runaway loop.
Arcjet enforces the budget in the application, at the tool call, keyed on an identity that you control. Sharing is by that key, across every matching guard() call, not by correlationId. The decision returns before the call proceeds, so the output is a call that did not happen rather than an alert about one that did. A direct Guard call fails open (allow with error codes). Check hasFailedOpen() where spending without a complete check is unacceptable. Vercel AI SDK and LangChain wrappers fail closed unless you opt into continuing on error.
Rate limiting is a cloud-backed rule, because distributed counters need shared state across your instances. For more information about the trade-offs of local versus cloud evaluation, see keeping security inspection local.
Learn more: AI budget control · Rate limiting algorithms
Frequently asked questions
How do I enforce token and spend budgets on AI agents?
Use a token bucket at the tool call, keyed on the identity you are protecting, with each call drawing in proportion to its cost. Sharing is by that key, not by correlationId. A limit at the HTTP entrypoint counts workflow starts rather than consumption.
Why do rate limits on my API endpoint not control agent cost?
An endpoint limit caps how many times a user can start a workflow. One request can trigger fifty tool calls, so the limit fires once while the spend happens fifty times. The budget has to sit where the spend happens.
What should I key an agent budget on?
Per user or session for customer-facing work, per plan or tier where budget is part of what someone bought, and per deployment or instance for a single-tenant worker or stdio MCP server with no user context. Never key on a value the caller supplies, because an attacker rotates it and the bucket never empties.
Should exceeding a budget stop the agent or degrade it?
It depends on the workflow. A research agent can often return partial results usefully. A payment workflow should stop, because a half-completed sequence is worse than a refused one.
AI runtime security in your code
Protect your AI agent workflows with Arcjet
A token-bucket budget at the tool call, keyed on an identity you control. Get a deny before the fiftieth call in a runaway loop.