Problem 130 · Practical builds

Tool-Using Agent Runtime with Budgets, Checkpoints and Approval

HardPractical multi-stage build

01Problem

Build the runtime loop for a tool-using agent. A policy (the model) maps history to a tool call or a final answer. The runtime must bound work, survive a restart mid-run, and never run a side-effecting tool without human approval.

Stages — finish each before reading the next

1Agent loop

start(goal) repeatedly asks the policy for the next action, executes tools, and appends results to history until a final answer. Tool exceptions and unknown tools become observations for the model, never crashes.

2Step budget

Stop with status budget_exhausted once max_steps policy calls have been made, including across resumes.

3Checkpoint / resume

checkpoint() serialises all run state to JSON; Agent.resume(blob, policy, tools) rebuilds an agent that continues exactly where the original stopped.

4Human approval

Tools flagged side_effect pause the run with status awaiting_approval and a pending call. decide(approved, note) runs it once with an idempotency key run_id:step, or reports the denial to the model as an observation.

02Key insight

Treat the run as a state machine whose whole state is plain JSON: history, step count, status, and the pending call. Pausing for approval is then just a status plus a checkpoint, and resume is deserialisation. Pass an idempotency key derived from run id and step to side-effecting tools so a crash between executing and checkpointing cannot double-send.

03Solution

solution.pypython86 lines

import json
from dataclasses import dataclass
from typing import Any, Callable


@dataclass(frozen=True)
class Tool:
    name: str
    fn: Callable[[dict, str], Any]   # (args, idempotency_key) -> JSON-serialisable result
    side_effect: bool = False        # True => needs human approval before it runs


Policy = Callable[[list[dict]], dict]  # history -> {"tool": name, "args": {...}} | {"final": text}


class Agent:
    def __init__(self, policy: Policy, tools: list[Tool], max_steps: int = 8, run_id: str = "run-1") -> None:
        self.policy, self.max_steps = policy, max_steps
        self.tools = {t.name: t for t in tools}
        self.state: dict[str, Any] = {"run_id": run_id, "history": [], "steps": 0,
                                      "status": "idle", "pending": None, "answer": None}

    # -- stage 1: the loop ----------------------------------------------------
    def start(self, goal: str) -> str:
        if self.state["status"] != "idle":
            raise RuntimeError("already started")
        self.state["history"].append({"role": "user", "content": goal})
        self.state["status"] = "running"
        return self._loop()

    def _execute(self, name: str, args: dict, key: str) -> None:
        try:
            result = self.tools[name].fn(args, key)
            self.state["history"].append({"role": "tool", "tool": name, "result": result})
        except Exception as e:  # tool errors go back to the model, never crash the run
            self.state["history"].append({"role": "tool", "tool": name, "error": f"{type(e).__name__}: {e}"})

    def _loop(self) -> str:
        s = self.state
        while s["status"] == "running":
            if s["steps"] >= self.max_steps:            # stage 2: step budget
                s["status"] = "budget_exhausted"
                break
            action = self.policy(json.loads(json.dumps(s["history"])))  # model sees a copy
            s["steps"] += 1
            if "final" in action:
                s["answer"], s["status"] = action["final"], "done"
                s["history"].append({"role": "assistant", "final": action["final"]})
                break
            name, args = action["tool"], action.get("args", {})
            s["history"].append({"role": "assistant", "tool": name, "args": args})
            if name not in self.tools:
                s["history"].append({"role": "tool", "tool": name, "error": f"unknown tool {name}"})
                continue
            key = f"{s['run_id']}:{s['steps']}"
            if self.tools[name].side_effect:            # stage 4: pause for a human
                s["pending"], s["status"] = {"tool": name, "args": args, "key": key}, "awaiting_approval"
                break
            self._execute(name, args, key)
        return s["status"]

    # -- stage 3: checkpoint / resume -----------------------------------------
    def checkpoint(self) -> str:
        return json.dumps(self.state, sort_keys=True)

    @classmethod
    def resume(cls, blob: str, policy: Policy, tools: list[Tool], max_steps: int = 8) -> "Agent":
        state = json.loads(blob)
        agent = cls(policy, tools, max_steps, state["run_id"])
        agent.state = state
        return agent

    def run(self) -> str:
        return self._loop()

    # -- stage 4: human approval ----------------------------------------------
    def decide(self, approved: bool, note: str = "") -> str:
        s = self.state
        if s["status"] != "awaiting_approval":
            raise RuntimeError("nothing awaiting approval")
        p, s["pending"], s["status"] = s["pending"], None, "running"
        if approved:
            self._execute(p["tool"], p["args"], p["key"])
        else:
            s["history"].append({"role": "tool", "tool": p["tool"], "error": f"denied by human: {note}"})
        return self._loop()

Try it first.

TimeO(steps) policy calls; checkpoint O(history size)

SpaceO(history size)

tests.py · 12 passing

tests.pypython50 lines

sent: list[tuple[str, str]] = []
def search(args: dict, key: str) -> str:
    if not args.get("q"):
        raise ValueError("empty query")
    return f"3 results for {args['q']}"
def send_email(args: dict, key: str) -> str:
    sent.append((args["to"], key))
    return "sent"
tools = [Tool("search", search), Tool("send_email", send_email, side_effect=True)]

def policy(history: list[dict]) -> dict:
    last = history[-1]
    if last["role"] == "user":
        return {"tool": "search", "args": {"q": "refund policy"}}
    if last.get("tool") == "search" and "result" in last:
        return {"tool": "send_email", "args": {"to": "ops@example.com"}}
    if last.get("tool") == "send_email":
        return {"final": "emailed" if "result" in last else "cancelled"}
    return {"final": "confused"}

# stage 1: plain loop with a pure tool
a = Agent(lambda h: {"final": "hi"} if len(h) > 2 else {"tool": "search", "args": {"q": "x"}}, tools)
assert a.start("go") == "done" and a.state["answer"] == "hi" and a.state["steps"] == 2
# tool errors and unknown tools become observations
b = Agent(lambda h: {"tool": "nope"} if len(h) == 1 else ({"tool": "search", "args": {}} if len(h) == 3 else {"final": h[-1]["error"]}), tools)
assert b.start("go") == "done" and b.state["answer"] == "ValueError: empty query"
assert b.state["history"][2]["error"] == "unknown tool nope"
# stage 2: step budget
c = Agent(lambda h: {"tool": "search", "args": {"q": "again"}}, tools, max_steps=3)
assert c.start("loop forever") == "budget_exhausted" and c.state["steps"] == 3
# stage 3 + 4: pause on side effect, checkpoint, resume in a "new process", approve
d = Agent(policy, tools)
assert d.start("email ops the refund policy") == "awaiting_approval"
assert sent == [] and d.state["pending"]["key"] == "run-1:2"
blob = d.checkpoint()
d2 = Agent.resume(blob, policy, tools)
assert d2.checkpoint() == blob
assert d2.decide(True) == "done" and d2.state["answer"] == "emailed"
assert sent == [("ops@example.com", "run-1:2")]
try:
    d2.decide(True); assert False
except RuntimeError:
    pass
# denial goes back to the model as an observation; nothing is sent
e = Agent.resume(blob, policy, tools)
assert e.decide(False, "not today") == "done" and e.state["answer"] == "cancelled"
assert len(sent) == 1
# a resumed run keeps its step budget
f = Agent.resume(blob, policy, tools, max_steps=2)
assert f.decide(True) == "budget_exhausted" and len(sent) == 2

04Staff-level follow-up

What changes in production? Store checkpoints after every step in a durable store with a version for optimistic concurrency, bound tokens and wall-clock time as well as steps, compact long histories by summarising old turns, and route approvals through a queue with expiry so a stale approval cannot run against a changed world.