Refactoring an Agent Loop Into a State Machine

The standard agent is a while loop around a model call and a tool dispatcher. It is the right thing to write first and the wrong thing to run in production, because every property you eventually need — a bounded step count, resumability, an audit trail, a human-approval pause — has to be bolted onto a construct with no place to put them. Converting it to an explicit state machine takes an afternoon and gives you all four.

What the loop actually gets wrong

The loop is not wrong because loops are bad. It is wrong because the agent’s state is spread across local variables that vanish the moment the process does.

• Termination is implicit. The loop ends when the model stops asking for tools. That is a property of the model’s output, not of your code, so there is no place to enforce a budget other than a counter somebody remembers to increment. Unbounded agents are this failure.
• It cannot be resumed. If the process restarts mid-run, everything is lost, including the tool calls that already had side effects. That is the difference between an agent that can wait for a human and one that cannot.
• The audit trail is the log. Reconstructing what the agent did means parsing prose written for humans, and it is invariably missing the one field you need.
• Repetition is invisible. A model that calls the same tool with the same arguments five times is doing so because nothing in the loop can see that it already did. Repetition loops are the most common runaway-cost mode.

Choosing the states

Five states cover almost every tool-using agent. Resist adding more until a real requirement forces one: the value here comes from the set being small enough to reason about exhaustively.

The loop, before

# agent_loop.py — the version everyone writes first
def run(question, tools, client, max_steps=20):
    messages = [{"role": "user", "content": question}]
    for _ in range(max_steps):
        reply = client.chat(messages=messages, tools=schemas(tools))
        messages.append(reply)
        if not reply.get("tool_calls"):
            return reply["content"]
        for call in reply["tool_calls"]:
            result = tools[call["name"]](**call["arguments"])
            messages.append({
                "role": "tool",
                "tool_call_id": call["id"],
                "content": json.dumps(result),
            })
    raise RuntimeError("step limit")

Note what is unrepresentable here. There is no way to stop between the tool call and the next model call, no way to record that a tool succeeded before the process died, no way to refuse a repeated call, and no way to express “wait for a human” other than blocking a thread.

The state machine, after

The conversion is mechanical: the loop body becomes a single step() function that takes a state and returns the next one, and the driver becomes a loop over step() that persists after each transition. Nothing about the model calls changes.

# agent_fsm.py — Python 3.11, standard library only
import json, hashlib
from dataclasses import dataclass, field, asdict

PLANNING, ACTING, AWAITING_APPROVAL, DONE, FAILED = (
    "PLANNING", "ACTING", "AWAITING_APPROVAL", "DONE", "FAILED")

@dataclass
class Run:
    id: str
    state: str = PLANNING
    messages: list = field(default_factory=list)
    pending: dict | None = None        # the tool call ACTING will execute
    answer: str | None = None
    error: str | None = None
    steps: int = 0
    cost_micros: int = 0
    seen: list = field(default_factory=list)   # fingerprints of executed calls

MAX_STEPS = 20
MAX_COST_MICROS = 2_000_000            # a hard budget, in millionths of a dollar
NEEDS_APPROVAL = {"send_email", "refund", "delete_record"}

def fingerprint(call):
    payload = call["name"] + json.dumps(call["arguments"], sort_keys=True)
    return hashlib.sha256(payload.encode()).hexdigest()[:16]

def step(run, client, tools, schemas):
    """Advance exactly one transition. Never continues on its own:
    the driver decides whether to take another step."""
    if run.steps >= MAX_STEPS:
        return _fail(run, "step budget exhausted")
    if run.cost_micros >= MAX_COST_MICROS:
        return _fail(run, "cost budget exhausted")

    if run.state == PLANNING:
        reply, usage = client.chat(messages=run.messages, tools=schemas)
        run.cost_micros += usage["cost_micros"]
        run.steps += 1
        run.messages.append(reply)
        calls = reply.get("tool_calls") or []
        if not calls:
            run.answer = reply["content"]
            run.state = DONE
            return run
        call = calls[0]                      # one at a time, deliberately
        if fingerprint(call) in run.seen:
            return _fail(run, f"repeated tool call: {call['name']}")
        run.pending = call
        run.state = AWAITING_APPROVAL if call["name"] in NEEDS_APPROVAL else ACTING
        return run

    if run.state == ACTING:
        call = run.pending
        try:
            result = tools[call["name"]](**call["arguments"])
        except Exception as exc:
            result = {"error": type(exc).__name__, "message": str(exc)}
        run.seen.append(fingerprint(call))
        run.messages.append({
            "role": "tool",
            "tool_call_id": call["id"],
            "content": json.dumps(result)[:8000],
        })
        run.pending = None
        run.state = PLANNING
        return run

    return run          # AWAITING_APPROVAL, DONE and FAILED are inert

def _fail(run, reason):
    run.error = reason
    run.state = FAILED
    return run

def approve(run):
    if run.state != AWAITING_APPROVAL:
        raise ValueError(f"cannot approve from {run.state}")
    run.state = ACTING
    return run

def drive(run, client, tools, schemas, save):
    """The whole driver. Persist after every transition, so a crash
    resumes from the last durable state instead of from the start."""
    while run.state in (PLANNING, ACTING):
        run = step(run, client, tools, schemas)
        save(asdict(run))
    return run

The save callback is the point of the exercise. Persisting a serialisable object after every transition is what makes the run resumable, auditable and pausable, and none of those are possible while the state lives in a stack frame.

Where the state actually lives

The refactor is only worth doing if the persistence is real, and there are three decisions to make about it that the code above deliberately leaves open.

One row per run, or one row per transition

Storing the current state as a single row that is overwritten is simpler and loses the history. Storing one immutable row per transition — an event log, with the current state derived by replaying or cached alongside — costs more storage and gives you the audit trail for nothing. For anything with side effects, take the second: the question you will actually be asked is “what did it do, in what order, and what did each step cost”, and only the log answers it.

What the transcript costs to store

The message array grows with every step, so writing the whole run object on every transition is quadratic in the number of steps. At twenty steps this is irrelevant; at two hundred it is not. Two fixes, and the first is usually enough: store the transcript once and append only the delta per transition, or truncate tool results before they enter the transcript at all. The code above caps tool output at 8,000 characters for exactly that reason, and a cap on tool results is worth having regardless of storage — an unbounded tool response is an unbounded prompt on the next step, which is an unbounded bill.

The resumption contract

A crash can land in exactly two places, and the state machine makes both answerable:

• Between transitions. The persisted state is consistent and resumption is trivial: load it and call the driver. This is the case the design is for.
• Inside ACTING, after the tool ran and before the result was written. The state on disk still says ACTING with a pending call, so a naive resume re-executes it. Whether that is safe is a property of the tool, not of the state machine. Either make the tool idempotent with a key derived from the run id and the call fingerprint, or record an intent row before execution and reconcile on resume. Deciding which, per tool, is the real work of the refactor.

Write the answer down per tool, in the tool registry, next to its schema. A tool whose safety on retry is undocumented will be retried by somebody eventually.

What the conversion buys

The wider argument for shaping agents this way — and for when a plain loop is still correct — is in the agent loop and stopping conditions.

What it does not fix

Three things the refactor is regularly claimed to solve and does not.

• It does not make the agent choose better. Every decision is still the model’s. If the agent picks the wrong tool, a state machine gives you a clearer record of it picking the wrong tool. Tool descriptions and prompt design remain the levers.
• It does not make side effects idempotent. Resumption re-runs whatever was in flight when the process died unless the tool itself is safe to repeat. The state machine gives you the place to put an idempotency key; it does not supply one. See idempotency.
• It does not bound cost by itself. The budget check above is only as good as the cost figure fed into it. If your client does not return a per-call cost, the budget is a step count wearing a currency symbol.