Budget Explosion (When Agents Burn Money) + Fixes + Code

from dataclasses import dataclass
import time


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


@dataclass
class Usage:
  tool_calls: int = 0
  model_tokens_in: int = 0
  model_tokens_out: int = 0
  estimated_usd: float = 0.0


class BudgetExceeded(RuntimeError):
  pass


def estimate_usd(tokens_in: int, tokens_out: int) -> float:
  # Replace with real pricing for your model(s).
  # This is a placeholder to show the pattern; pricing varies by provider and model.
  return (tokens_in + tokens_out) * 0.000002  # $/token (placeholder)


class BudgetGuard:
  def __init__(self, budget: Budget) -> None:
      self.budget = budget
      self.usage = Usage()
      self.started = time.time()
      self.steps = 0

  def check_step(self) -> None:
      self.steps += 1
      if self.steps > self.budget.max_steps:
          raise BudgetExceeded("step budget exceeded")
      if time.time() - self.started > self.budget.max_seconds:
          raise BudgetExceeded("time budget exceeded")

  def on_tool_call(self) -> None:
      self.usage.tool_calls += 1
      if self.usage.tool_calls > self.budget.max_tool_calls:
          raise BudgetExceeded("tool budget exceeded")

  def on_model_call(self, *, tokens_in: int, tokens_out: int) -> None:
      self.usage.model_tokens_in += tokens_in
      self.usage.model_tokens_out += tokens_out
      self.usage.estimated_usd = estimate_usd(
          self.usage.model_tokens_in, self.usage.model_tokens_out
      )
      if self.usage.estimated_usd > self.budget.max_usd:
          raise BudgetExceeded("cost budget exceeded")


def run(task: str, *, budget: Budget) -> str:
  guard = BudgetGuard(budget)

  while True:
      guard.check_step()

      # model call (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()
          result = call_tool(action.name, action.args)  # (pseudo)
          task = update_state(task, action, result)  # (pseudo)
      else:
          return action.final_answer

export class BudgetExceeded extends Error {}

export class BudgetGuard {
constructor(budget) {
  this.budget = budget;
  this.started = Date.now();
  this.steps = 0;
  this.usage = { toolCalls: 0, tokensIn: 0, tokensOut: 0, estimatedUsd: 0 };
}

estimateUsd(tokensIn, tokensOut) {
  // Replace with real pricing for your provider/model.
  // This is a placeholder to show the pattern.
  return (tokensIn + tokensOut) * 0.000002;
}

checkStep() {
  this.steps += 1;
  const elapsedS = (Date.now() - this.started) / 1000;
  if (this.steps > this.budget.maxSteps) throw new BudgetExceeded("step budget exceeded");
  if (elapsedS > this.budget.maxSeconds) throw new BudgetExceeded("time budget exceeded");
}

onToolCall() {
  this.usage.toolCalls += 1;
  if (this.usage.toolCalls > this.budget.maxToolCalls) throw new BudgetExceeded("tool budget exceeded");
}

onModelCall({ tokensIn, tokensOut }) {
  this.usage.tokensIn += tokensIn;
  this.usage.tokensOut += tokensOut;
  this.usage.estimatedUsd = this.estimateUsd(this.usage.tokensIn, this.usage.tokensOut);
  if (this.usage.estimatedUsd > this.budget.maxUsd) throw new BudgetExceeded("cost budget exceeded");
}
}

export function run(task, { budget }) {
const guard = new BudgetGuard(budget);

while (true) {
  guard.checkStep();

  // model call (pseudo)
  const { action, tokensIn, tokensOut } = llmDecide(task); // (pseudo)
  guard.onModelCall({ tokensIn, tokensOut });

  if (action.kind === "tool") {
    guard.onToolCall();
    const result = callTool(action.name, action.args); // (pseudo)
    task = updateState(task, action, result); // (pseudo)
  } else {
    return action.final_answer;
  }
}
}