Budget Controls for AI Agents (Steps, Time, $) + Code

from dataclasses import dataclass, field
import time
from typing import Any


TOOL_USD = {
  "search.read": 0.00,
  "http.get": 0.00,
  "browser.run": 0.20,  # placeholder
}


@dataclass(frozen=True)
class BudgetPolicy:
  max_steps: int = 25
  max_seconds: int = 60
  max_tool_calls: int = 12
  max_usd: float = 1.00


@dataclass
class BudgetState:
  started_at: float = field(default_factory=time.time)
  steps: int = 0
  tool_calls: int = 0
  tokens_in: int = 0
  tokens_out: int = 0
  tool_usd: float = 0.0

  def elapsed_s(self) -> float:
      return time.time() - self.started_at


def estimate_model_usd(tokens_in: int, tokens_out: int) -> float:
  # Replace with your real pricing model(s). Approximate is fine for guards.
  return (tokens_in + tokens_out) * 0.000002


class BudgetExceeded(RuntimeError):
  def __init__(self, stop_reason: str, *, state: BudgetState):
      super().__init__(stop_reason)
      self.stop_reason = stop_reason
      self.state = state


class BudgetGuard:
  def __init__(self, policy: BudgetPolicy):
      self.policy = policy
      self.state = BudgetState()

  def total_usd(self) -> float:
      return estimate_model_usd(self.state.tokens_in, self.state.tokens_out) + self.state.tool_usd

  def check(self) -> None:
      if self.state.steps > self.policy.max_steps:
          raise BudgetExceeded("max_steps", state=self.state)
      if self.state.elapsed_s() > self.policy.max_seconds:
          raise BudgetExceeded("max_seconds", state=self.state)
      if self.state.tool_calls > self.policy.max_tool_calls:
          raise BudgetExceeded("max_tool_calls", state=self.state)
      if self.total_usd() > self.policy.max_usd:
          raise BudgetExceeded("max_usd", state=self.state)

  def on_step(self) -> None:
      self.state.steps += 1
      self.check()

  def on_model_call(self, *, tokens_in: int, tokens_out: int) -> None:
      self.state.tokens_in += tokens_in
      self.state.tokens_out += tokens_out
      self.check()

  def on_tool_call(self, *, tool: str) -> None:
      self.state.tool_calls += 1
      self.state.tool_usd += float(TOOL_USD.get(tool, 0.0))
      self.check()


def run_agent(task: str, *, policy: BudgetPolicy) -> dict[str, Any]:
  guard = BudgetGuard(policy)

  try:
      while True:
          guard.on_step()

          # model decides next action (pseudo)
          action, tokens_in, tokens_out = llm_decide(task)  # (pseudo)
          guard.on_model_call(tokens_in=tokens_in, tokens_out=tokens_out)

          if action.kind == "tool":
              guard.on_tool_call(tool=action.name)
              obs = call_tool(action.name, action.args)  # (pseudo)
              task = update_state(task, action, obs)  # (pseudo)
              continue

          return {"status": "ok", "answer": action.final_answer, "usage": guard.state.__dict__}

  except BudgetExceeded as e:
      return {
          "status": "stopped",
          "stop_reason": e.stop_reason,
          "usage": e.state.__dict__,
          "partial": "Stopped by budget. Return partial results + a reason users can understand.",
      }

const TOOL_USD = {
"search.read": 0.0,
"http.get": 0.0,
"browser.run": 0.2, // placeholder
};

export class BudgetExceeded extends Error {
constructor(stopReason, { state }) {
  super(stopReason);
  this.stopReason = stopReason;
  this.state = state;
}
}

export class BudgetGuard {
constructor(policy) {
  this.policy = policy;
  this.state = {
    startedAtMs: Date.now(),
    steps: 0,
    toolCalls: 0,
    tokensIn: 0,
    tokensOut: 0,
    toolUsd: 0,
  };
}

elapsedS() {
  return (Date.now() - this.state.startedAtMs) / 1000;
}

estimateModelUsd(tokensIn, tokensOut) {
  return (tokensIn + tokensOut) * 0.000002;
}

totalUsd() {
  return this.estimateModelUsd(this.state.tokensIn, this.state.tokensOut) + this.state.toolUsd;
}

check() {
  if (this.state.steps > this.policy.maxSteps) throw new BudgetExceeded("max_steps", { state: this.state });
  if (this.elapsedS() > this.policy.maxSeconds) throw new BudgetExceeded("max_seconds", { state: this.state });
  if (this.state.toolCalls > this.policy.maxToolCalls) throw new BudgetExceeded("max_tool_calls", { state: this.state });
  if (this.totalUsd() > this.policy.maxUsd) throw new BudgetExceeded("max_usd", { state: this.state });
}

onStep() {
  this.state.steps += 1;
  this.check();
}

onModelCall({ tokensIn, tokensOut }) {
  this.state.tokensIn += tokensIn;
  this.state.tokensOut += tokensOut;
  this.check();
}

onToolCall({ tool }) {
  this.state.toolCalls += 1;
  this.state.toolUsd += Number(TOOL_USD[tool] || 0);
  this.check();
}
}