LLM Cost Control: Token Budgets and Rate Limits for AI Agents

A bug in an agent loop can spend a surprising amount of money before anyone notices. LLM cost control sets limits that your code enforces, so a mistake stops early and cheaply. In this lesson you will wrap the model client with budgets, add a daily spending cap and limit how often each person can start a run.

Advertisement

📖 What You’ll Learn

  • Where AI agent costs come from
  • How to set per-run limits
  • How to add a daily spending cap
  • How to rate limit each user

💡 Where LLM Cost Control Begins

Explain Like I’m 5: Pocket Money

Imagine: With a set amount for the week, you can spend freely until it runs out. Then you stop.

What it really means: A budget is a limit that stops the spending.

In testing: The run ends safely when it reaches its token limit.

Cost driverControl
Too many calls in a loopA call limit per run
Very long inputsA token limit, and shorter evidence
Too many runsA rate limit for each user
Many users on one dayA daily spending cap
An expensive model for easy workUsing a cheaper model where it is good enough

🔍 Count What Was Used

Explain Like I’m 5: Counting Your Steps

Imagine: A step counter does not guess, and it shows exactly how far you walked.

What it really means: Use the token counts that the provider returns.

In testing: The adapter adds each call’s real usage to a meter.

Create src/llm/usage.ts.

export class UsageMeter {
  calls = 0;
  inputTokens = 0;
  outputTokens = 0;

  add(inputTokens: number, outputTokens: number): void {
    this.inputTokens += inputTokens;
    this.outputTokens += outputTokens;
  }

  get totalTokens(): number {
    return this.inputTokens + this.outputTokens;
  }

  reset(): void {
    this.calls = 0;
    this.inputTokens = 0;
    this.outputTokens = 0;
  }
}

// Safe as one shared object because the worker handles one run at a time.
export const usageMeter = new UsageMeter();

Now update src/llm/provider-adapter.ts. Add this import.

import { usageMeter } from './usage.js';

Then add this line straight after the response arrives, before the stop_reason check.

usageMeter.add(response.usage.input_tokens, response.usage.output_tokens);

Other providers report usage in their own way, so adapt that line to yours.


🧩 Enforce the Budget

Create src/llm/guarded-llm.ts. It wraps any model client and adds limits and logging, so your agents do not change.

import type * as z from 'zod';
import { log } from '../observability/logger.js';
import type { Metrics } from '../observability/metrics.js';
import type { LlmClient, LlmPrompt } from './llm-client.js';
import type { UsageMeter } from './usage.js';

export class BudgetExceededError extends Error {}

export type BudgetLimits = {
  maxCallsPerRun: number;
  maxTokensPerRun: number;
  maxDailyCostUsd: number;
  inputCostPerMillionUsd: number;
  outputCostPerMillionUsd: number;
};

export interface SpendLedger {
  spentToday(): number;
  add(usd: number): void;
}

export class GuardedLlmClient implements LlmClient {
  readonly providerName: string;

  constructor(
    private readonly inner: LlmClient,
    private readonly meter: UsageMeter,
    private readonly limits: BudgetLimits,
    private readonly ledger: SpendLedger,
    private readonly metrics: Metrics,
    private readonly runId: string,
  ) {
    this.providerName = inner.providerName;
  }

  private checkBudget(): void {
    if (this.meter.calls >= this.limits.maxCallsPerRun) {
      throw new BudgetExceededError('Reached the model call limit for this run');
    }
    if (this.meter.totalTokens >= this.limits.maxTokensPerRun) {
      throw new BudgetExceededError('Reached the token limit for this run');
    }
    if (this.ledger.spentToday() >= this.limits.maxDailyCostUsd) {
      throw new BudgetExceededError('Reached the daily model spending limit');
    }
  }

  async generateStructuredOutput<T>(prompt: LlmPrompt, schema: z.ZodType<T>): Promise<T> {
    this.checkBudget();
    this.meter.calls += 1;

    const started = Date.now();
    const inputBefore = this.meter.inputTokens;
    const outputBefore = this.meter.outputTokens;
    let outcome = 'ok';
    try {
      return await this.inner.generateStructuredOutput(prompt, schema);
    } catch (error) {
      outcome = 'error';
      throw error;
    } finally {
      const inputTokens = this.meter.inputTokens - inputBefore;
      const outputTokens = this.meter.outputTokens - outputBefore;
      const costUsd =
        (inputTokens * this.limits.inputCostPerMillionUsd +
          outputTokens * this.limits.outputCostPerMillionUsd) /
        1_000_000;
      this.ledger.add(costUsd);
      this.metrics.increment('qa_agent_llm_calls_total', { outcome });
      this.metrics.increment('qa_agent_llm_tokens_total', {}, inputTokens + outputTokens);
      // Prompt text is deliberately left out, because it can contain application data.
      log('info', 'llm.call', {
        runId: this.runId,
        provider: this.providerName,
        prompt: prompt.name,
        promptVersion: prompt.version,
        durationMs: Date.now() - started,
        inputTokens,
        outputTokens,
        costUsd,
        outcome,
      });
    }
  }
}

When a limit is reached, the error stops the workflow, the AI agent orchestrator records the reason and a report is still produced.

Set the prices from your provider’s current price list in your configuration. They change, so do not hard-code them.

A cap protects you, but check your provider too

Set a spending limit in your model provider’s account as well. If your code has a bug, that second limit still applies.


🛠️ Add Rate Limits and the Daily Ledger

Explain Like I’m 5: One Turn Each

Imagine: On the slide, everyone gets a turn before anyone goes again.

What it really means: A rate limit stops one person using up everything.

In testing: Each person can start ten runs an hour.

Create src/service/limits.ts.

import type { SpendLedger } from '../llm/guarded-llm.js';
import type { ServiceStore } from './store.js';

export class FixedWindowLimiter {
  private readonly windows = new Map<string, { startedAt: number; count: number }>();

  constructor(
    private readonly limit: number,
    private readonly windowMs: number,
    private readonly now: () => number = Date.now,
  ) {}

  allow(key: string): boolean {
    const time = this.now();
    const current = this.windows.get(key);
    if (!current || time - current.startedAt >= this.windowMs) {
      this.windows.set(key, { startedAt: time, count: 1 });
      return true;
    }
    if (current.count >= this.limit) {
      return false;
    }
    current.count += 1;
    return true;
  }
}

export class StoredSpendLedger implements SpendLedger {
  constructor(private readonly store: ServiceStore) {}

  private key(): string {
    return `spend:${new Date().toISOString().slice(0, 10)}`;
  }

  spentToday(): number {
    return Number(this.store.getSetting(this.key()) ?? 0);
  }

  add(usd: number): void {
    this.store.setSetting(this.key(), String(this.spentToday() + usd));
  }
}

Create tests/unit/budget.test.ts.

import assert from 'node:assert/strict';
import { test } from 'node:test';
import * as z from 'zod';
import { BudgetExceededError, GuardedLlmClient, type SpendLedger } from '../../src/llm/guarded-llm.js';
import type { LlmClient } from '../../src/llm/llm-client.js';
import { UsageMeter } from '../../src/llm/usage.js';
import { Metrics } from '../../src/observability/metrics.js';
import { FixedWindowLimiter } from '../../src/service/limits.js';

const prompt = { name: 'p', version: '1', task: 't', untrustedData: {} };
const limits = {
  maxCallsPerRun: 2,
  maxTokensPerRun: 1000,
  maxDailyCostUsd: 1,
  inputCostPerMillionUsd: 1000,
  outputCostPerMillionUsd: 2000,
};

function setup(spent = 0) {
  const meter = new UsageMeter();
  let total = spent;
  const ledger: SpendLedger = { spentToday: () => total, add: (usd) => (total += usd) };
  const inner: LlmClient = {
    providerName: 'fake',
    generateStructuredOutput: async <T>(_p: unknown, schema: z.ZodType<T>) => {
      meter.add(100, 50);
      return schema.parse('ok');
    },
  };
  const client = new GuardedLlmClient(inner, meter, limits, ledger, new Metrics(), 'run-1');
  return { client, spent: () => total };
}

test('stops at the call limit', async () => {
  const { client } = setup();
  await client.generateStructuredOutput(prompt, z.string());
  await client.generateStructuredOutput(prompt, z.string());
  await assert.rejects(client.generateStructuredOutput(prompt, z.string()), BudgetExceededError);
});

test('records the cost of each call', async () => {
  const { client, spent } = setup();
  await client.generateStructuredOutput(prompt, z.string());
  assert.equal(spent(), 0.2);
});

test('stops at the daily spending limit', async () => {
  const { client } = setup(1);
  await assert.rejects(client.generateStructuredOutput(prompt, z.string()), /daily/);
});

test('rate limits each user separately, then resets', () => {
  let now = 0;
  const limiter = new FixedWindowLimiter(1, 1000, () => now);
  assert.equal(limiter.allow('asha'), true);
  assert.equal(limiter.allow('asha'), false);
  assert.equal(limiter.allow('ben'), true);
  now = 1000;
  assert.equal(limiter.allow('asha'), true);
});

This rate limiter lives in memory, so it resets on a restart. That is fine for one server. With several servers, keep the counters in a shared store.


📌 Spend Less

TechniqueEffect
Send only the evidence that is neededFewer input tokens
Use a cheaper model for simple stepsLower cost per call
Use your provider’s prompt cachingLower cost for repeated context
Avoid repeating unchanged workFewer calls
Stop as soon as a stage failsNo wasted later stages

Always check quality using your evaluation pack after any change, because a cheaper setup that fails more often may cost more overall.


🌍 Real-Life Example: The Weekend Bill

An agent retries a bad response in a loop. Nobody is watching over the weekend, and the bill arrives on Monday. A call limit per run would have stopped it after a few requests. A daily cap would have been the second line of defence.


🧪 Try It Yourself

  1. Create the files, and run the unit tests.
  2. Set MAX_LLM_CALLS_PER_RUN=1, and run the mock workflow.
  3. Read the stop reason in the report.
  4. Set your provider’s real prices in the configuration.
  5. Calculate the cost of one run.

📋 Ready-to-Copy Prompt

You are an engineer reducing the cost of an LLM-based QA agent.

Per-run data:
[paste calls, tokens and cost by prompt with no sensitive content]

Find the biggest cost drivers. Suggest changes in order of savings, and
say how each one could reduce quality.

Do not invent prices. Tell me to confirm current pricing with my provider.

⚠️ Common Mistakes

  • Relying on people to watch the bill
  • Having no limit on a single run
  • Hard-coding prices that later change
  • Treating a budget stop as a crash
  • Cutting cost without checking quality

🎯 Key Takeaways

  • Limits in code stop mistakes early.
  • Use real token counts from the provider.
  • Combine per-run, per-user and daily limits.
  • Re-evaluate quality after any cost change.

🎤 Interview Questions

1. How do you control LLM costs in an agent?

Model answer: I enforce limits on calls and tokens per run, rate limit users and apply a daily cap. I log usage and cost for each prompt.

2. Why not rely on the provider’s spending limit alone?

Model answer: It is a useful backstop, but it protects the whole account. Per-run limits stop one bad run quickly and explain why.

3. What should happen when a budget is reached?

Model answer: The workflow should stop safely, record the reason and produce a report. It should not fail silently or continue.

4. How would you reduce token usage?

Model answer: I would send only relevant evidence, use cheaper models where evaluation supports it, use caching and avoid repeated work.

5. Why might a cheaper model cost more?

Model answer: If it fails more often, retries and human review can outweigh the savings. I compare the total cost per accepted result.


📝 Assignment

Report the cost for each prompt across five runs. Find the most expensive one, reduce its input and compare quality and cost before and after.


❓ FAQ

What is a sensible daily limit?

Measure a typical run, multiply by your expected number of runs and add a margin.

Does the mock use real tokens?

No. It reports zero, so the token and cost limits need a real provider or a fake client, as in the unit tests.


🚀 Next Steps

Well done on finishing this lesson. Complete the assignment, then continue with AI Agent Production Readiness: Rollout and Incident Runbook.