Testing That an Agent Respects Its Maximum Step Budget

A step limit is a piece of your own control flow, which means it is fully testable with no model in the loop and no non-determinism anywhere. It is also the guard most likely to be written once, never exercised, and quietly broken by a refactor that moves the counter inside the retry.

The budget is yours, not the model’s

Nothing in a model prevents it from asking for another tool call. The model produces a response whose stop reason indicates it wants to use a tool; your loop decides whether to give it one. The bound is therefore a property of code you own, and the correct test has a mocked transport rather than a real one. A test that tries to provoke a real model into looping is testing something you cannot control, will be slow, will cost money, and will not fail when the guard breaks.

It is worth being precise about what the guard is for. It is not a quality mechanism — an agent that hits the ceiling has usually failed at its task. It is a blast-radius mechanism: it bounds spend, bounds latency, and bounds the number of side effects a confused loop can cause before a human sees it. That framing decides the assertions, because it means “stopped eventually” is not good enough. It has to stop at a number you can multiply by a price.

The runaway mock

The fixture is a transport that always returns the same tool-use response. It never terminates on its own, so if the loop has no working bound the test hangs rather than passing — which is the correct failure and is why this fixture is better than one that returns a finite sequence.

import { describe, expect, it, vi } from "vitest";
import { runAgent } from "../src/agent";

const alwaysAsksForATool = () => ({
  stop_reason: "tool_use",
  content: [{ type: "tool_use", id: "tu", name: "search", input: { q: "anything" } }],
});

describe("step budget", () => {
  it("halts at exactly maxSteps model calls", async () => {
    const transport = vi.fn().mockImplementation(async () => alwaysAsksForATool());
    const search = vi.fn().mockResolvedValue({ hits: [] });

    const result = await runAgent({
      transport,
      tools: { search },
      maxSteps: 5,
      input: "find everything",
    });

    expect(transport).toHaveBeenCalledTimes(5);
    expect(search).toHaveBeenCalledTimes(4);
    expect(result.stopReason).toBe("step_budget");
    expect(result.text).toBeTruthy();
  });
});

Give the test an explicit timeout well under your suite default. If the bound is broken, you want a five-second failure that says the loop did not terminate, not a job that runs until the CI runner is killed and reports nothing useful.

The three assertions

• An exact count, not an upper bound. toHaveBeenCalledTimes(5) rather than a check that the count is at most five. An off-by-one in the safe direction is still a bug: it means the number in your config is not the number in your bill, and every capacity estimate built on it is wrong.
• The last tool did not execute. This is the one people miss. A loop that makes its final model call, receives a tool request, executes the tool and then notices it is out of budget has performed a side effect after the stop condition. If that tool posts a message or charges a card, the bound did not do its job. Assert the tool mock’s call count is one lower than the model call count.
• The result is typed and usable. Hitting the ceiling must not throw, and must not return an empty string. It should return a discriminated result carrying a stopReason the caller can branch on, plus whatever partial work exists. A loop that throws on budget exhaustion turns a routine condition into an incident, and the handler upstream inevitably swallows it.

Decide what a step is, once

“Five steps” is ambiguous between five model calls, five tool executions, and five complete request-and-execute cycles. Any of the three is defensible; having two of them in one codebase is not, and it happens as soon as a second person adds a budget to a nested loop.

Pick model calls. It is the unit that maps to cost and latency, it is countable at exactly one place in the code, and it makes the tool count a derived quantity rather than a second counter that can drift. Write the definition in a comment next to the config value, and let the test above be the documentation of the off-by-one: five model calls, four tool executions, because the fifth call is the one that gets cut off.

Where steps leak: retries and sub-agents

Two refactors reliably break a working budget, and both deserve their own test.

A retry must not consume a step. When a request fails with a 429 or a 5xx and your client retries it, that is the same step being attempted twice, not two steps. But if the retry lives outside the transport — in the loop rather than in the client — the counter increments on each attempt, and an agent running through a period of provider overload silently gets a much smaller effective budget. The test is a transport that fails twice and then returns a tool-use response, with an assertion that the successful path still gets its full count of steps. The inverse bug exists too: a retry wrapper placed around the whole loop restarts the agent from scratch with a fresh budget, multiplying the ceiling by the retry count.

A sub-agent must draw from the same budget. If a parent with a budget of ten can spawn workers that each have a budget of ten, the real ceiling is a hundred, and nobody wrote that number down. Pass a budget object that decrements rather than an integer that is copied, and assert it: run a parent that spawns two children, give the whole thing a budget of six, and check the total transport call count across all three is six. The same argument applies to a circuit breaker and to any other bound that is shared rather than per-component.

Two smaller leaks are worth a test each once the big ones are covered. A tool that itself calls a model — a summariser invoked to compress a large tool result before it goes back into context — spends a model call that the loop counter never sees, so the transport call count and the step count diverge. Count at the transport, not in the loop, and this stops being possible. And a loop that repairs invalid structured output by asking again is making a model call for every repair attempt; whether those count against the budget is a decision, but it has to be a decision, because a validator that rejects persistently can consume the entire budget without the agent making any progress at all. Assert whichever answer you chose with a mocked model that returns malformed output every time.