DEV Community

Cover image for Token Budget for AI Agents on Amazon Bedrock
Sascha Mirow-Stanitzok
Sascha Mirow-Stanitzok

Posted on Originally published at sascha-ms.de

Token Budget for AI Agents on Amazon Bedrock

#agents #ai #aws #serverless

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
Enter fullscreen mode Exit fullscreen mode

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"]
}
Enter fullscreen mode Exit fullscreen mode

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);
}
Enter fullscreen mode Exit fullscreen mode

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);
    },
  };
}
Enter fullscreen mode Exit fullscreen mode

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.

Diagram: reserve, call Converse, settle; budget bar with three calls and a stop before the fourth

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: Record<string, string> = {
  "B-1001": "shipped",
  "B-1002": "processing",
};

export interface AgentTool {
  spec: Tool;
  run(input: Record<string, unknown>): Promise<unknown>;
}

export const tools: Record<string, AgentTool> = {
  searchBooks: {
    spec: {
      toolSpec: {
        name: "searchBooks",
        description: "Searches the catalog for books with a keyword in the title.",
        inputSchema: {
          json: {
            type: "object",
            properties: { query: { type: "string" } },
            required: ["query"],
          },
        },
      },
    },
    async run(input) {
      const q = String(input.query ?? "").toLowerCase();
      return CATALOG.filter((b) => b.title.toLowerCase().includes(q));
    },
  },
  getOrderStatus: {
    spec: {
      toolSpec: {
        name: "getOrderStatus",
        description: "Returns the status of an order, e.g. B-1001.",
        inputSchema: {
          json: {
            type: "object",
            properties: { orderId: { type: "string" } },
            required: ["orderId"],
          },
        },
      },
    },
    async run(input) {
      const id = String(input.orderId ?? "");
      return { orderId: id, status: ORDERS[id] ?? "unknown" };
    },
  },
};
Enter fullscreen mode Exit fullscreen mode

The loop itself always follows the same pattern: plan, reserve, call, settle, run tools. If the Bedrock call fails, the code releases the reservation and rethrows the error. When the budget stops the job, finish returns all model texts so far as the partial result.

// src/agent.ts
// Tool use loop with a token budget per job and per user and day.
import {
  ConverseCommand,
  type ContentBlock,
  type ConverseCommandOutput,
  type Message,
} from "@aws-sdk/client-bedrock-runtime";
import {
  EMPTY_USAGE,
  addUsage,
  dayKey,
  estimateNextInput,
  normalizeUsage,
  planNextCall,
  weightedTokens,
  type BudgetConfig,
  type Usage,
} from "./budget.ts";
import type { Ledger } from "./ledger.ts";
import type { AgentTool } from "./tools.ts";

export interface ConverseSender {
  send(command: ConverseCommand): Promise<ConverseCommandOutput>;
}

export interface AgentDeps {
  bedrock: ConverseSender;
  ledger: Ledger;
  tools: Record<string, AgentTool>;
  config: BudgetConfig;
  modelId: string;
  systemPrompt: string;
  timeZone: string;
  now?: () => Date;
}

export type AgentStatus =
  | "completed"
  | "truncated"
  | "job_budget"
  | "daily_budget"
  | "max_iterations";

export interface AgentResult {
  status: AgentStatus;
  answer: string;       // when stopped early: the partial result
  toolsUsed: string[];
  usage: Usage;
  weightedTokens: number;
  modelCalls: number;
}

const DAY = 86_400;

export async function runAgent(
  job: { jobId: string; userId: string; question: string },
  deps: AgentDeps,
): Promise<AgentResult> {
  const { config: cfg, ledger } = deps;
  const now = (deps.now ?? (() => new Date()))();
  const nowSec = Math.floor(now.getTime() / 1000);
  const jobKey = `JOB#${job.jobId}`;
  const userKey = `USER#${job.userId}#DAY#${dayKey(now, deps.timeZone)}`;

  let jobCommitted = await ledger.current(jobKey);
  let userCommitted = await ledger.current(userKey);

  const messages: Message[] = [{ role: "user", content: [{ text: job.question }] }];
  const toolSpecs = Object.values(deps.tools).map((t) => t.spec);
  const texts: string[] = [];
  const toolsUsed: string[] = [];
  let usage = EMPTY_USAGE;
  let last = EMPTY_USAGE;
  let toolChars = 0;
  let modelCalls = 0;

  const finish = (status: AgentStatus): AgentResult => ({
    status,
    answer: texts.join("\n").trim(),
    toolsUsed,
    usage,
    weightedTokens: weightedTokens(usage, cfg.weights),
    modelCalls,
  });

  for (let i = 0; ; i++) {
    const firstGuess = Math.ceil(
      (deps.systemPrompt.length + job.question.length + JSON.stringify(toolSpecs).length) / 4,
    );
    const estimatedInput = i === 0 ? firstGuess : estimateNextInput(last, toolChars);
    const plan = planNextCall({ iteration: i, estimatedInput, jobCommitted, userCommitted }, cfg);
    if (plan.action === "stop") return finish(plan.reason);

    // 1. Reserve: first the job, then the user's daily counter.
    const jobRes = await ledger.reserve(jobKey, plan.reserve, cfg.jobLimit, nowSec + 30 * DAY);
    if (!jobRes.ok) return finish("job_budget");
    const userRes = await ledger.reserve(userKey, plan.reserve, cfg.dailyUserLimit, nowSec + 3 * DAY);
    if (!userRes.ok) {
      await ledger.settle(jobKey, -plan.reserve, EMPTY_USAGE);
      return finish("daily_budget");
    }

    // 2. Call the model, maxTokens from the budget left.
    let res: ConverseCommandOutput;
    try {
      res = await deps.bedrock.send(
        new ConverseCommand({
          modelId: deps.modelId,
          system: [{ text: deps.systemPrompt }],
          messages,
          toolConfig: { tools: toolSpecs },
          inferenceConfig: { maxTokens: plan.maxTokens },
        }),
      );
    } catch (err) {
      // Failed call: release the reservation.
      await ledger.settle(jobKey, -plan.reserve, EMPTY_USAGE);
      await ledger.settle(userKey, -plan.reserve, EMPTY_USAGE);
      throw err;
    }

    // 3. Settle: difference between actual usage and reservation.
    const u = normalizeUsage(res.usage);
    const delta = weightedTokens(u, cfg.weights) - plan.reserve;
    jobCommitted = await ledger.settle(jobKey, delta, u);
    userCommitted = await ledger.settle(userKey, delta, u);
    usage = addUsage(usage, u);
    last = u;
    modelCalls++;

    const reply = res.output?.message;
    if (reply) messages.push(reply);
    for (const block of reply?.content ?? []) {
      if (block.text) texts.push(block.text);
    }

    if (res.stopReason === "max_tokens") return finish("truncated");
    if (res.stopReason !== "tool_use") return finish("completed");

    // 4. Run the tools and send the results back.
    const results: ContentBlock[] = [];
    for (const block of reply?.content ?? []) {
      const call = block.toolUse;
      if (!call?.toolUseId || !call.name) continue;
      const tool = deps.tools[call.name];
      const output = tool
        ? await tool.run((call.input ?? {}) as Record<string, unknown>)
        : { error: `Unknown tool ${call.name}` };
      toolsUsed.push(call.name);
      results.push({
        toolResult: { toolUseId: call.toolUseId, content: [{ text: JSON.stringify(output) }] },
      });
    }
    toolChars = results.reduce((s, r) => s + (r.toolResult?.content?.[0]?.text?.length ?? 0), 0);
    messages.push({ role: "user", content: results });
  }
}
Enter fullscreen mode Exit fullscreen mode

Step 5: Metrics in the Embedded Metric Format

The Embedded Metric Format (EMF) is a JSON structure that CloudWatch Logs reads as an instruction: extract metrics from this log line. You therefore need neither a PutMetricData call nor an extra permission. I write the user ID only as a property, not as a dimension. Every new dimension value creates a separate CloudWatch metric, and with thousands of customers that quickly becomes a mess.

// src/metrics.ts
// CloudWatch Embedded Metric Format: one JSON line in the log becomes a metric.

export interface JobMetrics {
  weightedTokens: number;
  modelCalls: number;
  budgetExhausted: 0 | 1;
  estimatedCostUsd?: number;
}

export function emfLine(
  namespace: string,
  service: string,
  m: JobMetrics,
  props: Record<string, string>,
  now = Date.now(),
): string {
  const metrics = [
    { Name: "WeightedTokens", Unit: "Count" },
    { Name: "ModelCalls", Unit: "Count" },
    { Name: "BudgetExhausted", Unit: "Count" },
  ];
  if (m.estimatedCostUsd !== undefined) {
    metrics.push({ Name: "EstimatedCostUsd", Unit: "None" });
  }
  return JSON.stringify({
    _aws: {
      Timestamp: now,
      CloudWatchMetrics: [
        { Namespace: namespace, Dimensions: [["Service"]], Metrics: metrics },
      ],
    },
    Service: service,
    WeightedTokens: m.weightedTokens,
    ModelCalls: m.modelCalls,
    BudgetExhausted: m.budgetExhausted,
    ...(m.estimatedCostUsd !== undefined ? { EstimatedCostUsd: m.estimatedCostUsd } : {}),
    // User and job ID only as properties, never as dimensions:
    // otherwise every user gets a separate metric.
    ...props,
  });
}
Enter fullscreen mode Exit fullscreen mode

In the handler, I write the line with process.stdout.write. The Lambda docs explicitly recommend this for the case where you switch the function's log format to JSON. With console.log, CloudWatch would no longer recognize the metric in that case.

// src/handler.ts
// Lambda entry point: accept the job, run the agent, write the metric.
import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime";
import { DynamoDBClient } from "@aws-sdk/client-dynamodb";
import { estimateCost, type BudgetConfig, type PricesPerMillion } from "./budget.ts";
import { createLedger } from "./ledger.ts";
import { emfLine } from "./metrics.ts";
import { runAgent } from "./agent.ts";
import { tools } from "./tools.ts";

const env = (name: string): string => {
  const v = process.env[name];
  if (!v) throw new Error(`Environment variable ${name} is missing`);
  return v;
};

const config: BudgetConfig = {
  jobLimit: Number(env("JOB_TOKEN_LIMIT")),
  dailyUserLimit: Number(env("DAILY_USER_TOKEN_LIMIT")),
  maxOutputPerCall: Number(env("MAX_OUTPUT_PER_CALL")),
  minOutputPerCall: 256,
  maxIterations: 8,
  weights: JSON.parse(env("TOKEN_WEIGHTS")),
};

// Prices are optional, e.g. {"input":..,"output":..,"cacheRead":..,"cacheWrite":..}
// in USD per 1M tokens, copied from the Bedrock pricing page.
const prices: PricesPerMillion | undefined = process.env.PRICES_PER_MILLION
  ? JSON.parse(process.env.PRICES_PER_MILLION)
  : undefined;

const bedrock = new BedrockRuntimeClient({});
const ledger = createLedger(new DynamoDBClient({}), env("TABLE_NAME"));

export const handler = async (event: { jobId: string; userId: string; question: string }) => {
  const result = await runAgent(event, {
    bedrock,
    ledger,
    tools,
    config,
    modelId: env("MODEL_ID"),
    timeZone: env("SHOP_TIME_ZONE"),
    systemPrompt:
      "You are the support assistant of the Seitenwind bookshop. Answer briefly and use the tools.",
  });

  const exhausted = result.status === "job_budget" || result.status === "daily_budget";
  // process.stdout.write instead of console.log: the line stays pure JSON,
  // even if you switch the function's log format to JSON.
  process.stdout.write(
    emfLine(
      "Seitenwind/Agent",
      "support-agent",
      {
        weightedTokens: result.weightedTokens,
        modelCalls: result.modelCalls,
        budgetExhausted: exhausted ? 1 : 0,
        ...(prices ? { estimatedCostUsd: estimateCost(result.usage, prices) } : {}),
      },
      { jobId: event.jobId, userId: event.userId, status: result.status },
    ) + "\n",
  );

  if (exhausted) {
    return {
      ...result,
      answer:
        (result.answer ? result.answer + "\n\n" : "") +
        "I have reached the budget for this request. A person from the team will look at the rest.",
    };
  }
  return result;
};
Enter fullscreen mode Exit fullscreen mode

Step 6: Infrastructure with serverless.yml

Serverless Framework v4 ships with esbuild and builds TypeScript handlers during deployment without a plugin. By default, it leaves out the packages under @aws-sdk/*, because the Lambda runtime provides them. However, esbuild does not type-check, so tsc --noEmit runs separately beforehand.

# serverless.yml
service: seitenwind-support-agent
frameworkVersion: "4"

params:
  default:
    # Model or inference profile ID from the Bedrock console of your account
    modelId: ${env:MODEL_ID}
    alarmEmail: ${env:ALARM_EMAIL}

provider:
  name: aws
  runtime: nodejs24.x
  architecture: arm64
  region: eu-central-1
  memorySize: 512
  timeout: 120
  environment:
    TABLE_NAME: !Ref BudgetTable
    MODEL_ID: ${param:modelId}
    JOB_TOKEN_LIMIT: "60000"
    DAILY_USER_TOKEN_LIMIT: "400000"
    MAX_OUTPUT_PER_CALL: "2000"
    # Price ratio per token type, derived from the Bedrock pricing page
    TOKEN_WEIGHTS: '{"input":1,"output":1,"cacheRead":1,"cacheWrite":1}'
    SHOP_TIME_ZONE: Europe/Berlin
  iam:
    role:
      statements:
        - Effect: Allow
          Action:
            - bedrock:InvokeModel
          Resource:
            # Broad for getting started; narrow it to your model later
            - arn:aws:bedrock:*::foundation-model/*
            - arn:aws:bedrock:*:${aws:accountId}:inference-profile/*
        - Effect: Allow
          Action:
            - dynamodb:GetItem
            - dynamodb:UpdateItem
          Resource: !GetAtt BudgetTable.Arn

functions:
  agent:
    handler: src/handler.handler

resources:
  Resources:
    BudgetTable:
      Type: AWS::DynamoDB::Table
      Properties:
        BillingMode: PAY_PER_REQUEST
        AttributeDefinitions:
          - AttributeName: pk
            AttributeType: S
        KeySchema:
          - AttributeName: pk
            KeyType: HASH
        TimeToLiveSpecification:
          AttributeName: expiresAt
          Enabled: true

    AlarmTopic:
      Type: AWS::SNS::Topic
      Properties:
        Subscription:
          - Protocol: email
            Endpoint: ${param:alarmEmail}

    BudgetExhaustedAlarm:
      Type: AWS::CloudWatch::Alarm
      Properties:
        AlarmDescription: Unusually many jobs run into the token budget
        Namespace: Seitenwind/Agent
        MetricName: BudgetExhausted
        Dimensions:
          - Name: Service
            Value: support-agent
        Statistic: Sum
        Period: 300
        EvaluationPeriods: 1
        Threshold: 5
        ComparisonOperator: GreaterThanOrEqualToThreshold
        TreatMissingData: notBreaching
        AlarmActions:
          - !Ref AlarmTopic

    TokenSpikeAlarm:
      Type: AWS::CloudWatch::Alarm
      Properties:
        AlarmDescription: Weighted tokens per hour above the expected level
        Namespace: Seitenwind/Agent
        MetricName: WeightedTokens
        Dimensions:
          - Name: Service
            Value: support-agent
        Statistic: Sum
        Period: 3600
        EvaluationPeriods: 1
        Threshold: 2000000
        ComparisonOperator: GreaterThanThreshold
        TreatMissingData: notBreaching
        AlarmActions:
          - !Ref AlarmTopic
Enter fullscreen mode Exit fullscreen mode

The limits are example values. Start on the tight side and watch the WeightedTokens metric for a week. After that, you know how much a normal job needs, and you can set the limit and alarm threshold with a reason behind them. In the example, all weights are set to 1 until you have derived them from the pricing page.

Step 7: Test with a mocked Bedrock client

You do not need an AWS account for the tests. The agent only expects an object with a send method, so I give it a fake that plays back prepared replies. In the test, the ledger runs in memory with the same rule as in DynamoDB. The weights in the test are chosen freely and are not real prices.

// test/budget.test.ts
import { test } from "node:test";
import assert from "node:assert/strict";
import {
  dayKey,
  estimateCost,
  estimateNextInput,
  planNextCall,
  weightedTokens,
  type BudgetConfig,
} from "../src/budget.ts";

const cfg: BudgetConfig = {
  jobLimit: 10_000,
  dailyUserLimit: 50_000,
  maxOutputPerCall: 1_000,
  minOutputPerCall: 256,
  maxIterations: 8,
  weights: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
};

test("weighting accounts for cache tokens", () => {
  const w = weightedTokens(
    { inputTokens: 100, outputTokens: 10, cacheReadInputTokens: 1000, cacheWriteInputTokens: 0 },
    cfg.weights,
  );
  assert.equal(w, 100 + 50 + 100);
});

test("plenty of budget: maxTokens stays at the upper bound", () => {
  const plan = planNextCall({ iteration: 0, estimatedInput: 500, jobCommitted: 0, userCommitted: 0 }, cfg);
  assert.deepEqual(plan, { action: "call", maxTokens: 1_000, reserve: 500 + 5_000 });
});

test("tight budget: maxTokens shrinks to what is left", () => {
  const plan = planNextCall({ iteration: 2, estimatedInput: 3_000, jobCommitted: 4_000, userCommitted: 0 }, cfg);
  // 6000 left, input 3000 -> 3000 / 5 = 600 output tokens
  assert.deepEqual(plan, { action: "call", maxTokens: 600, reserve: 6_000 });
});

test("too little left: stop with job_budget", () => {
  const plan = planNextCall({ iteration: 3, estimatedInput: 5_000, jobCommitted: 4_000, userCommitted: 0 }, cfg);
  assert.deepEqual(plan, { action: "stop", reason: "job_budget" });
});

test("daily budget kicks in before the job budget", () => {
  const plan = planNextCall({ iteration: 0, estimatedInput: 500, jobCommitted: 0, userCommitted: 49_000 }, cfg);
  assert.deepEqual(plan, { action: "stop", reason: "daily_budget" });
});

test("maxIterations stops every loop", () => {
  const plan = planNextCall({ iteration: 8, estimatedInput: 1, jobCommitted: 0, userCommitted: 0 }, cfg);
  assert.deepEqual(plan, { action: "stop", reason: "max_iterations" });
});

test("estimate grows with history and tool result", () => {
  const est = estimateNextInput(
    { inputTokens: 800, outputTokens: 120, cacheReadInputTokens: 2_000, cacheWriteInputTokens: 0 },
    400,
  );
  assert.equal(est, 800 + 2_000 + 120 + 100);
});

test("cost only with prices you entered yourself", () => {
  const cost = estimateCost(
    { inputTokens: 1_000_000, outputTokens: 0, cacheReadInputTokens: 0, cacheWriteInputTokens: 0 },
    { input: 2, output: 10, cacheRead: 0.2, cacheWrite: 2.5 },
  );
  assert.equal(cost, 2);
});

test("day key follows the shop's time zone", () => {
  const lateEvening = new Date("2026-10-31T23:30:00Z");
  assert.equal(dayKey(lateEvening, "UTC"), "2026-10-31");
  assert.equal(dayKey(lateEvening, "Europe/Berlin"), "2026-11-01");
});
Enter fullscreen mode Exit fullscreen mode 
// test/agent.test.ts
import { test } from "node:test";
import assert from "node:assert/strict";
import type { ConverseCommand, ConverseCommandOutput } from "@aws-sdk/client-bedrock-runtime";
import { runAgent, type AgentDeps } from "../src/agent.ts";
import type { BudgetConfig } from "../src/budget.ts";
import type { Ledger } from "../src/ledger.ts";
import { tools } from "../src/tools.ts";

// In-memory ledger with the same semantics as the DynamoDB version.
function memoryLedger(start: Record<string, number> = {}) {
  const store = new Map(Object.entries(start));
  const ledger: Ledger = {
    async current(key) {
      return store.get(key) ?? 0;
    },
    async reserve(key, amount, limit) {
      const before = store.get(key) ?? 0;
      if (before > limit - amount) return { ok: false, committed: before };
      store.set(key, before + amount);
      return { ok: true, committed: before + amount };
    },
    async settle(key, delta) {
      const after = (store.get(key) ?? 0) + delta;
      store.set(key, after);
      return after;
    },
  };
  return { ledger, store };
}

// Mocked Bedrock client: plays back prepared replies.
function fakeBedrock(replies: ConverseCommandOutput[]) {
  const seen: ConverseCommand[] = [];
  return {
    seen,
    client: {
      async send(cmd: ConverseCommand) {
        seen.push(cmd);
        const next = replies.shift();
        if (!next) throw new Error("no more replies prepared");
        return next;
      },
    },
  };
}

const toolCall = (id: string, input: Record<string, unknown>, usage = { inputTokens: 900, outputTokens: 80 }) =>
  ({
    $metadata: {},
    stopReason: "tool_use",
    usage: { ...usage, totalTokens: usage.inputTokens + usage.outputTokens },
    output: {
      message: {
        role: "assistant",
        content: [{ text: "Let me check the catalog." }, { toolUse: { toolUseId: id, name: "searchBooks", input } }],
      },
    },
    metrics: { latencyMs: 1 },
  }) as ConverseCommandOutput;

const finalAnswer = (text: string) =>
  ({
    $metadata: {},
    stopReason: "end_turn",
    usage: { inputTokens: 1_100, outputTokens: 60, totalTokens: 1_160 },
    output: { message: { role: "assistant", content: [{ text }] } },
    metrics: { latencyMs: 1 },
  }) as ConverseCommandOutput;

const cfg: BudgetConfig = {
  jobLimit: 8_000,
  dailyUserLimit: 40_000,
  maxOutputPerCall: 1_000,
  minOutputPerCall: 256,
  maxIterations: 8,
  weights: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
};

const job = { jobId: "a-1", userId: "customer-7", question: "Do you have Cooking with Serverless?" };
const now = () => new Date("2026-10-31T10:00:00Z");

function deps(over: Partial<AgentDeps> & Pick<AgentDeps, "bedrock" | "ledger">): AgentDeps {
  return { tools, config: cfg, modelId: "test-model", systemPrompt: "Test", timeZone: "Europe/Berlin", now, ...over };
}

test("tool call and answer: counters are correct after settling", async () => {
  const { ledger, store } = memoryLedger();
  const bedrock = fakeBedrock([toolCall("t1", { query: "serverless" }), finalAnswer("Yes, 4 copies in stock.")]);

  const result = await runAgent(job, deps({ bedrock: bedrock.client, ledger }));

  assert.equal(result.status, "completed");
  assert.match(result.answer, /4 copies/);
  assert.deepEqual(result.toolsUsed, ["searchBooks"]);
  // (900 + 80*5) + (1100 + 60*5) = 1300 + 1400
  assert.equal(result.weightedTokens, 2_700);
  assert.equal(store.get("JOB#a-1"), 2_700);
  assert.equal(store.get("USER#customer-7#DAY#2026-10-31"), 2_700);
});

test("endless loop: budget ends the job with a partial result", async () => {
  const { ledger, store } = memoryLedger();
  const replies = Array.from({ length: 20 }, (_, i) =>
    toolCall(`t${i}`, { query: "worm" }, { inputTokens: 1_000 + i * 600, outputTokens: 100 }),
  );
  const bedrock = fakeBedrock(replies);

  const result = await runAgent(job, deps({ bedrock: bedrock.client, ledger }));

  assert.equal(result.status, "job_budget");
  assert.match(result.answer, /catalog/);
  assert.ok(result.modelCalls < 8, "budget kicks in before maxIterations");
  assert.ok((store.get("JOB#a-1") ?? 0) <= cfg.jobLimit);
  // maxTokens shrinks as the budget left shrinks
  const maxTokens = bedrock.seen.map((c) => c.input.inferenceConfig?.maxTokens ?? 0);
  assert.ok(maxTokens.at(-1)! < maxTokens[0]!);
});

test("daily budget exhausted: not a single model call", async () => {
  const { ledger } = memoryLedger({ "USER#customer-7#DAY#2026-10-31": 39_900 });
  const bedrock = fakeBedrock([finalAnswer("should never arrive")]);

  const result = await runAgent(job, deps({ bedrock: bedrock.client, ledger }));

  assert.equal(result.status, "daily_budget");
  assert.equal(bedrock.seen.length, 0);
});

test("Bedrock error: reservation is released", async () => {
  const { ledger, store } = memoryLedger();
  const bedrock = fakeBedrock([]);

  await assert.rejects(runAgent(job, deps({ bedrock: bedrock.client, ledger })));

  assert.equal(store.get("JOB#a-1"), 0);
  assert.equal(store.get("USER#customer-7#DAY#2026-10-31"), 0);
});
Enter fullscreen mode Exit fullscreen mode

I test the DynamoDB part separately. There, I care about whether the right expressions end up in the UpdateItemCommand and whether a failed condition comes back cleanly as “not ok”.

// test/ledger.test.ts
import { test } from "node:test";
import assert from "node:assert/strict";
import { ConditionalCheckFailedException, UpdateItemCommand } from "@aws-sdk/client-dynamodb";
import { createLedger } from "../src/ledger.ts";

test("reserve uses ADD with a condition in one UpdateItem", async () => {
  const sent: UpdateItemCommand[] = [];
  const db = {
    async send(cmd: UpdateItemCommand) {
      sent.push(cmd);
      return { Attributes: { committed: { N: "1500" } } };
    },
  };
  const ledger = createLedger(db as never, "budgets");

  const res = await ledger.reserve("JOB#a-1", 1_500, 8_000, 123);

  assert.deepEqual(res, { ok: true, committed: 1_500 });
  const input = sent[0]!.input;
  assert.match(input.UpdateExpression ?? "", /^ADD committed :amount/);
  assert.equal(input.ConditionExpression, "attribute_not_exists(committed) OR committed <= :maxBefore");
  assert.deepEqual(input.ExpressionAttributeValues?.[":maxBefore"], { N: "6500" });
});

test("reserve reports a rejected reservation instead of throwing", async () => {
  const db = {
    async send() {
      throw new ConditionalCheckFailedException({
        message: "The conditional request failed",
        $metadata: {},
        Item: { committed: { N: "7900" } },
      });
    },
  };
  const ledger = createLedger(db as never, "budgets");

  const res = await ledger.reserve("JOB#a-1", 1_500, 8_000, 123);

  assert.deepEqual(res, { ok: false, committed: 7_900 });
});

test("settle always counts, including negative corrections", async () => {
  const sent: UpdateItemCommand[] = [];
  const db = {
    async send(cmd: UpdateItemCommand) {
      sent.push(cmd);
      return { Attributes: { committed: { N: "1200" } } };
    },
  };
  const ledger = createLedger(db as never, "budgets");

  const committed = await ledger.settle("JOB#a-1", -300, {
    inputTokens: 900, outputTokens: 60, cacheReadInputTokens: 0, cacheWriteInputTokens: 0,
  });

  assert.equal(committed, 1_200);
  assert.equal(sent[0]!.input.ConditionExpression, undefined);
  assert.deepEqual(sent[0]!.input.ExpressionAttributeValues?.[":delta"], { N: "-300" });
});
Enter fullscreen mode Exit fullscreen mode 
// test/metrics.test.ts
import { test } from "node:test";
import assert from "node:assert/strict";
import { emfLine } from "../src/metrics.ts";

test("EMF line contains metadata and values at the top level", () => {
  const line = JSON.parse(
    emfLine("Seitenwind/Agent", "support-agent", { weightedTokens: 2700, modelCalls: 2, budgetExhausted: 0 }, { jobId: "a-1" }, 1_700_000_000_000),
  );
  assert.equal(line._aws.Timestamp, 1_700_000_000_000);
  assert.deepEqual(line._aws.CloudWatchMetrics[0].Dimensions, [["Service"]]);
  assert.equal(line.WeightedTokens, 2700);
  assert.equal(line.Service, "support-agent");
  assert.equal(line.jobId, "a-1");
});
Enter fullscreen mode Exit fullscreen mode 
npx tsc --noEmit
npm test
Enter fullscreen mode Exit fullscreen mode

Expected result: tsc prints nothing, and node:test shows pass 17 and fail 0 at the end. The endless loop test is the most telling one. There, maxTokens drops from 1,000 to 1,000 to 535, and then the agent stops before the fourth call with 6,300 of 8,000 weighted tokens.

Step 8: Deploy and try it out 

export MODEL_ID="<model-or-inference-profile-id>"
export ALARM_EMAIL="you@example.com"
npx tsc --noEmit && serverless deploy
serverless invoke -f agent -d '{"jobId":"a-1","userId":"customer-7","question":"Is Cooking with Serverless available?"}'
Enter fullscreen mode Exit fullscreen mode

The response is a JSON object with status, answer, toolsUsed and the usage. Afterwards, DynamoDB holds two items, one for the job and one for the day. Also confirm the SNS email, otherwise the alarms never arrive. If you want to see the stop in action, set JOB_TOKEN_LIMIT very low for a test and invoke the same job again.

Common mistakes

  • A condition on settlement: If you also add a condition when settling, you lose usage exactly when things get tight.
  • Counting only totalTokens: Cache tokens sit in separate fields and have their own prices. Weight them separately.
  • User ID as a dimension: That creates one metric per person. Use properties and analyze them with Logs Insights.
  • Lambda timeout too short: Several model calls take time. If the timeout does not fit the loop, the function dies before the clean stop.
  • Repeated Lambda invocations: With asynchronous invocations, Lambda can start a failed job again. Same jobId then means same counter, which is actually what you want here.

Cost and security

Every model call costs two reservations and two settlements, so four writes to DynamoDB. Compared with a Bedrock call, that is usually negligible, but check it with your own numbers. If the function crashes between reservation and settlement, the reservation stays. The budget is then slightly too strict, never too generous. For a cost limit, I find that behavior right.

The IAM permissions in the example are broad for getting started. In production, restrict bedrock:InvokeModel to exactly your model and inference profile. Also, the user ID must come from the authentication in front of the Lambda function. It must never come from the prompt. Otherwise, someone simply books themselves a fresh daily budget.

How does this fit with prompt caching, approvals and evaluation?

The budget is an emergency brake and does not lower the cost per call. For that, prompt caching with Bedrock is the better tool. With agents in particular, the system prompt and tool list repeat in every round. With caching, those tokens show up as cacheReadInputTokens, which is exactly why separate weights pay off. A well-cached agent gets more rounds out of the same budget.

For expensive or risky actions, a budget is not enough. If the assistant is supposed to trigger a refund, for example, a person belongs in between. I show how that looks with Step Functions in the article on human in the loop for AI agents. By the way, hitting the budget is a good trigger for exactly that handover.

And how do you know whether a budget is too tight? For that, you need test cases with expected results. When you evaluate AI agents, add token usage per test case as a metric of its own. Then you can see whether a new prompt improves the quality or only the bill.

My take

I consider a token budget for AI agents mandatory in production. It is the counterpart to the Lambda timeout: nobody would run a function without a timeout, yet many agents run without any cost limit. The reserve-and-settle pattern is hardly more code than a simple counter. In return, it holds up with parallel jobs too.

I disagree with the idea of setting the budget as generously as possible “so the agent does not get in the way”. An agent that regularly runs into its budget usually has a different problem: unclear tool descriptions, oversized tool results or a job that is really two jobs. The budget makes these problems visible. That is why, for me, the alarm belongs to the setup just as much as the limit.

To be honest, the approach has limits too. The input estimate is rough, and a single large call can easily push past the limit. CountTokens makes this more precise, but it adds an extra call per round. Also, the budget only protects you from an expensive result. It does not save you from a bad one. Still, for me it is the first safeguard I add, even before any fine-tuning of the prompt.

FAQ

Can I calculate the cost of a Bedrock agent before the call?

Only partly. You cap the output with maxTokens. You can estimate the input or, for supported models, count it with the free CountTokens API. The price per token comes from the Bedrock pricing page for your model and service tier.

Is maxTokens alone enough as a cost brake?

No. maxTokens only limits the length of a single reply. The number of rounds and the growing input stay untouched. Only a budget per job and per user limits the total.

Do cache tokens count against the budget?

In my approach, yes, but weighted. The Converse API reports cacheReadInputTokens and cacheWriteInputTokens separately. You give each field a weight that matches the price ratio of your model.

Why DynamoDB and not a counter in Lambda memory?

Under load, Lambda starts several execution environments that know nothing about each other. A counter in memory only sees its own instance. A user's daily budget therefore needs a shared place with atomic updates.

Your next step

Use your existing agent as the starting point for your first token budget for AI agents. For one week, only log the usage per job with the EMF line from step 5. After that, set the job budget slightly above the highest normal value and switch on the alarm. I am curious: how do you limit the cost of your agents today, and which limit has already kicked in for you? Tell me in the comments.

Sources

This post was created with AI assistance.

👉 Next step: Curious what your own Lambda functions cost? Try my Lambda cost calculator – it runs right in your browser.

Originally published on sascha-ms.de. More articles on serverless, AI in everyday development and leading tech teams: sascha-ms.de/en

Top comments (0)