A token budget for AI agents is, to me, the brake I install before I hit the accelerator for the first time. One question keeps bugging me: who stops the loop when the model simply keeps going? An agent calls a tool, gets a result, thinks, and calls the next tool. That is exactly what we want. However, every round costs money, and each round costs more than the one before.
That is why this article builds a support assistant for a fictional bookshop called “Seitenwind”. It runs in a Lambda function and talks to Amazon Bedrock through the Converse API. On top of that, it gets a fixed budget per job and per customer per day. The counters live in DynamoDB, and the metrics end up in CloudWatch. This is how I would build it today, with Serverless Framework v4 and TypeScript.
In short
Agent loops resend the whole history to the model on every call, so costs grow faster than the number of steps. Before every Bedrock call I reserve tokens in DynamoDB, cap maxTokens at the budget left and settle the actual usage afterwards. When the budget runs out, the job ends cleanly with a partial result, and a CloudWatch alarm flags unusual spikes.
What you will learn in this article
- Why tool use loops make token costs grow disproportionately
- Which usage fields the Converse API returns, including cache tokens
- Reserving and settling atomically in DynamoDB without races
- How maxTokens follows from the budget left and the agent stops cleanly
- How to write metrics with the Embedded Metric Format and set alarms
Why do agent loops make costs explode?
A token is the smallest unit of text a language model works with, roughly a piece of a word. Bedrock bills per token, separately for input, output and cache. For a simple chat, that stays manageable: question in, answer out.
An agent works differently. The Converse API is stateless, so you send the complete history with every call. That includes the system prompt, the tool descriptions, every model reply and every tool result. As a result, the tenth call contains everything from the nine calls before it. The sum of input tokens therefore grows roughly quadratically with the number of steps.
Here is a small calculation from my test: the first call has 1,000 input tokens, and each further call adds 600. After three rounds, that is 1,000 + 1,600 + 2,200 = 4,800 input tokens. After ten rounds, it would already be 37,000, even though each single step barely grew. A tool that returns a long order list as JSON speeds this up even more.
On top of that, there are situations where the agent never finishes. For example, it calls the same tool again and again with slightly different parameters. Or a tool returns an error, and the model bravely keeps trying. Without a limit, such a loop runs until the Lambda timeout hits. Until then, you have still paid for every single call.
What does a token budget for AI agents look like?
I think of the budget in three levels. Each level catches a different kind of outlier. In addition, there is a hard cap on the number of loop iterations as a last line of defense.
| Level | What it limits | Where it applies |
|---|---|---|
| Per call | Length of a single model reply | inferenceConfig.maxTokens in the Converse request |
| Per job | Sum of all calls for one customer request | Counter JOB#<id> in DynamoDB |
| Per user and day | Many small jobs from the same person | Counter USER#<id>#DAY#<date> in DynamoDB |
| Iterations | Loops that do not end despite the budget | maxIterations in the code |
I count the budget in “weighted tokens”. Input, output, cache reads and cache writes each have their own price per model. Instead of hard-coding prices, I store weights that reflect the price ratio. You take the values from the Bedrock pricing page for your model, your region and your service tier. There are several tiers there, such as Standard, Flex, Priority and Reserved. If you enter the prices yourself, the code also calculates a cost estimate in US dollars.
Which numbers does Bedrock return after each call?
The response from Converse contains a usage field of type TokenUsage. According to the API reference, it holds inputTokens, outputTokens and totalTokens. Optionally, cacheReadInputTokens and cacheWriteInputTokens appear as soon as prompt caching is involved. A newer addition is cacheDetails, which breaks down cache writes by time to live.
The reference does not say explicitly whether inputTokens already includes the cache tokens. I therefore treat the four values as separate items and weight each one individually. This also matches the pricing page, which lists the token types separately. Also important is stopReason. The value tool_use means “keep going”, end_turn means “done”, and max_tokens tells you that your limit cut the reply short.
Before the call, you only know the input tokens roughly. For some models, there is the free CountTokens API for this. According to the docs, it returns exactly the number that would be billed. However, not every model supports it. In the example, I therefore use a simple estimate: the history of the last call plus its reply plus the tool result divided by four characters per token.
Why reserve instead of counting afterwards?
The obvious approach would be: call the model, then increment the usage with ADD, and check before the next call. That works for a single job. For the daily budget, though, three jobs from the same customer may run in parallel. All three read “2,000 tokens left”, all three call Bedrock, and in the end you have overrun the budget three times over.
That is why I reserve before every call. A single UpdateItem raises the committed counter by the planned amount, but only if the limit still holds afterwards. DynamoDB applies such updates to an item one after another, so the check and the write happen in one step. In a condition, however, DynamoDB can only compare, not calculate. I therefore compute the limit minus the reservation in the code first and pass the value as :maxBefore.
After the call, I settle: the counter receives the difference between actual usage and the reservation. This settlement deliberately runs without a condition. The tokens are used and paid for, so they must end up in the counter, even if that pushes the budget slightly over.
🎬 Token budget inside the agent loop – this animated diagram only runs on my site: watch it on sascha-ms.de
How do I build a token budget for AI agents in Lambda?
Prerequisites
- An AWS account with access to a Bedrock model that supports tool use through the Converse API
- Node.js and npm locally, plus Serverless Framework v4 (
npm i -g serverless) - Basic knowledge of TypeScript, Lambda and DynamoDB
I tested the code with @aws-sdk/client-bedrock-runtime and @aws-sdk/client-dynamodb version 3.1161.0, TypeScript 7.0.2 and tsx 4.23.15. As the runtime, I use nodejs24.x. According to the docs, Node.js 26 on Lambda is currently only a public preview, so it is not meant for production.
Step 1: Create the project
mkdir seitenwind-support-agent && cd seitenwind-support-agent npm init -y npm pkg set type=module npm pkg set scripts.typecheck="tsc --noEmit" npm pkg set scripts.test="node --import tsx --test test/*.test.ts" npm install @aws-sdk/client-bedrock-runtime @aws-sdk/client-dynamodb npm install -D typescript tsx @types/node@24 mkdir src test The tsconfig.json is deliberately strict. With allowImportingTsExtensions, you may import files with the .ts extension. tsx handles that in the tests, and esbuild builds it during deployment as well.
// tsconfig.json { "compilerOptions": { "target": "es2023", "module": "nodenext", "moduleResolution": "nodenext", "strict": true, "noUncheckedIndexedAccess": true, "allowImportingTsExtensions": true, "noEmit": true, "skipLibCheck": true, "types": ["node"] }, "include": ["src", "test"] } Step 2: Budget logic as pure functions
The heart of it all has no AWS dependency at all. planNextCall receives the estimated input and both counter values. From that, the function calculates how many output tokens still fit. If that is below a sensible minimum, it returns “stop”; otherwise it returns the values for maxTokens and the reservation.
// src/budget.ts // Pure budget logic: no AWS call, therefore easy to test. export interface Usage { inputTokens: number; outputTokens: number; cacheReadInputTokens: number; cacheWriteInputTokens: number; } // Weights relative to a regular input token. Derive them from the // Bedrock pricing page for your model instead of guessing. export interface Weights { input: number; output: number; cacheRead: number; cacheWrite: number; } export interface BudgetConfig { jobLimit: number; // weighted tokens per job dailyUserLimit: number; // weighted tokens per user and day maxOutputPerCall: number; // upper bound for maxTokens minOutputPerCall: number; // below this, another call is not worth it maxIterations: number; // hard limit for the loop weights: Weights; } export const EMPTY_USAGE: Usage = { inputTokens: 0, outputTokens: 0, cacheReadInputTokens: 0, cacheWriteInputTokens: 0, }; // Bedrock only returns the cache fields when caching is involved. export function normalizeUsage(raw: Partial<Usage> | undefined): Usage { return { inputTokens: raw?.inputTokens ?? 0, outputTokens: raw?.outputTokens ?? 0, cacheReadInputTokens: raw?.cacheReadInputTokens ?? 0, cacheWriteInputTokens: raw?.cacheWriteInputTokens ?? 0, }; } export function addUsage(a: Usage, b: Usage): Usage { return { inputTokens: a.inputTokens + b.inputTokens, outputTokens: a.outputTokens + b.outputTokens, cacheReadInputTokens: a.cacheReadInputTokens + b.cacheReadInputTokens, cacheWriteInputTokens: a.cacheWriteInputTokens + b.cacheWriteInputTokens, }; } export function weightedTokens(u: Usage, w: Weights): number { return Math.ceil( u.inputTokens * w.input + u.outputTokens * w.output + u.cacheReadInputTokens * w.cacheRead + u.cacheWriteInputTokens * w.cacheWrite, ); } // Rough estimate of the next prompt: everything already in the history, // plus the new model reply plus the tool results (about 4 characters per token). export function estimateNextInput(last: Usage, toolResultChars: number): number { const previousPrompt = last.inputTokens + last.cacheReadInputTokens + last.cacheWriteInputTokens; return previousPrompt + last.outputTokens + Math.ceil(toolResultChars / 4); } export type Plan = | { action: "call"; maxTokens: number; reserve: number } | { action: "stop"; reason: "job_budget" | "daily_budget" | "max_iterations" }; export interface PlanInput { iteration: number; estimatedInput: number; jobCommitted: number; // used + reserved in this job userCommitted: number; // used + reserved today } export function planNextCall(p: PlanInput, cfg: BudgetConfig): Plan { if (p.iteration >= cfg.maxIterations) { return { action: "stop", reason: "max_iterations" }; } const jobLeft = cfg.jobLimit - p.jobCommitted; const userLeft = cfg.dailyUserLimit - p.userCommitted; const left = Math.min(jobLeft, userLeft); const inputCost = Math.ceil(p.estimatedInput * cfg.weights.input); const outputRoom = Math.floor((left - inputCost) / cfg.weights.output); const maxTokens = Math.min(cfg.maxOutputPerCall, outputRoom); if (maxTokens < cfg.minOutputPerCall) { return { action: "stop", reason: userLeft < jobLeft ? "daily_budget" : "job_budget", }; } const reserve = inputCost + Math.ceil(maxTokens * cfg.weights.output); return { action: "call", maxTokens, reserve }; } // Optional: cost estimate with prices YOU copy from the pricing page. export interface PricesPerMillion { input: number; output: number; cacheRead: number; cacheWrite: number; } export function estimateCost(u: Usage, p: PricesPerMillion): number { const sum = u.inputTokens * p.input + u.outputTokens * p.output + u.cacheReadInputTokens * p.cacheRead + u.cacheWriteInputTokens * p.cacheWrite; return sum / 1_000_000; } // Day key in the shop's time zone, e.g. "2026-10-31". export function dayKey(now: Date, timeZone: string): string { return new Intl.DateTimeFormat("en-CA", { timeZone, year: "numeric", month: "2-digit", day: "2-digit", }).format(now); } I build the day key in the shop's time zone. Otherwise, the new budget day for customers in Berlin would start in the middle of the night. Depending on daylight saving time, that is one or two o'clock. That sounds pedantic, but it is exactly the kind of bug that ends up in the support inbox.
Step 3: Reserve and settle counters in DynamoDB
The table has only a partition key pk. One item per job, one item per user and day. ADD creates missing attributes and even the whole item, so you do not need a prior PutItem. The expiresAt attribute uses DynamoDB TTL so that old daily counters disappear on their own.
// src/ledger.ts // Counters in DynamoDB: reserve before the call, settle afterwards. import { ConditionalCheckFailedException, GetItemCommand, UpdateItemCommand, type DynamoDBClient, } from "@aws-sdk/client-dynamodb"; import type { Usage } from "./budget.ts"; export type ReserveResult = { ok: boolean; committed: number }; export interface Ledger { current(key: string): Promise<number>; reserve(key: string, amount: number, limit: number, ttlSec: number): Promise<ReserveResult>; settle(key: string, delta: number, usage: Usage): Promise<number>; } type Sender = Pick<DynamoDBClient, "send">; const n = (v: number) => ({ N: String(v) }); export function createLedger(db: Sender, tableName: string): Ledger { return { async current(key) { const res = await db.send( new GetItemCommand({ TableName: tableName, Key: { pk: { S: key } }, ProjectionExpression: "committed", ConsistentRead: true, }), ); return Number(res.Item?.committed?.N ?? 0); }, async reserve(key, amount, limit, ttlSec) { try { const res = await db.send( new UpdateItemCommand({ TableName: tableName, Key: { pk: { S: key } }, // Check and increment in ONE operation: no race // between two parallel Lambda invocations. UpdateExpression: "ADD committed :amount, calls :one SET expiresAt = if_not_exists(expiresAt, :ttl)", ConditionExpression: "attribute_not_exists(committed) OR committed <= :maxBefore", ExpressionAttributeValues: { ":amount": n(amount), ":one": n(1), ":ttl": n(ttlSec), ":maxBefore": n(limit - amount), }, ReturnValues: "UPDATED_NEW", ReturnValuesOnConditionCheckFailure: "ALL_OLD", }), ); return { ok: true, committed: Number(res.Attributes?.committed?.N ?? amount) }; } catch (err) { if (err instanceof ConditionalCheckFailedException) { return { ok: false, committed: Number(err.Item?.committed?.N ?? limit) }; } throw err; } }, async settle(key, delta, usage) { // Deliberately WITHOUT a condition: used tokens are paid for // and must always end up in the counter. const res = await db.send( new UpdateItemCommand({ TableName: tableName, Key: { pk: { S: key } }, UpdateExpression: "ADD committed :delta, inputTokens :in, outputTokens :out, cacheReadTokens :cr, cacheWriteTokens :cw", ExpressionAttributeValues: { ":delta": n(delta), ":in": n(usage.inputTokens), ":out": n(usage.outputTokens), ":cr": n(usage.cacheReadInputTokens), ":cw": n(usage.cacheWriteInputTokens), }, ReturnValues: "UPDATED_NEW", }), ); return Number(res.Attributes?.committed?.N ?? 0); }, }; } With ReturnValuesOnConditionCheckFailure: "ALL_OLD", a rejected reservation gives you the current counter value right in the exception. As a result, you save an extra GetItem.
Reservation and settlement per call, below it the test run with 8,000 weighted tokens.
Step 4: The agent loop with a budget
First, the tools. The bookshop gets a catalog search and a status lookup for orders. All data is static, because this is about the loop and not about inventory management.
// src/tools.ts // Two tools for the support assistant of a fictional bookshop. import type { Tool } from "@aws-sdk/client-bedrock-runtime"; const CATALOG = [ { isbn: "978-3-0000-0001-1", title: "Cooking with Serverless", stock: 4 }, { isbn: "978-3-0000-0002-8", title: "The Bookworm in the Data Center", stock: 0 }, ]; const ORDERS: Source：https://dev.to/sascha-ms/token-budget-for-ai-agents-on-amazon-bedrock-2mo0Related posts
A Practical Workflow for AI-Assisted SEO Analysis
Stablecoin subscriptions: separate renewal from paid access
How to Merge Multiple PDF Files Online Without Losing Quality or Hitting Paywalls
How to Convert PDF Technical Specs and Whitepapers into Clean Markdown for GitHub
Continuation of the discussion about Gemini
