Token Budgets per Agent Run: Allocating, Tracking, and Enforcing

Agents consume tokens at every step — reasoning, tool calls, retries, and context accumulation. Without budgets, a single runaway agent loop can cost hundreds of dollars. Token budgets give you control: allocate per-run limits, track usage in real-time, and enforce hard stops when limits are hit.

Why Token Budgets Matter

An agent that loops through tool calls can burn tokens exponentially:

User request: 500 tokens
Agent reasoning: 1,000 tokens
Tool call 1 + result: 2,000 tokens
Agent reasoning: 1,500 tokens
Tool call 2 + result: 2,500 tokens
... (continues for 20 iterations)
Total: 50,000+ tokens

At $10/1M tokens, that’s $0.50 per request. At scale, it’s thousands per day.

Architecture

Agent Run Request
    ↓
Budget Allocator (assign budget based on task type)
    ↓
Token Counter (track usage in real-time)
    ↓
Agent Loop
    ├── LLM Call (deduct tokens)
    ├── Tool Call (deduct tokens)
    └── Check remaining budget
    ↓
Budget Enforcer (halt if exceeded)

Implementation

Token Budget Manager

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import asyncio

class BudgetStatus(Enum):
    OK = "ok"
    WARNING = "warning"  # 80% used
    CRITICAL = "critical"  # 95% used
    EXCEEDED = "exceeded"

@dataclass
class TokenBudget:
    run_id: str
    allocated: int  # Total budget for this run
    used: int = 0
    warning_threshold: float = 0.8
    critical_threshold: float = 0.95
    created_at: datetime = field(default_factory=datetime.now)
    
    @property
    def remaining(self) -> int:
        return max(0, self.allocated - self.used)
    
    @property
    def usage_percent(self) -> float:
        return self.used / self.allocated if self.allocated > 0 else 0
    
    @property
    def status(self) -> BudgetStatus:
        if self.used >= self.allocated:
            return BudgetStatus.EXCEEDED
        elif self.usage_percent >= self.critical_threshold:
            return BudgetStatus.CRITICAL
        elif self.usage_percent >= self.warning_threshold:
            return BudgetStatus.WARNING
        return BudgetStatus.OK

class TokenBudgetManager:
    def __init__(self):
        self.budgets: dict[str, TokenBudget] = {}
        self.default_budgets = {
            "simple_query": 5000,
            "complex_task": 50000,
            "research": 100000,
            "unlimited": float('inf'),
        }
    
    def allocate(
        self,
        run_id: str,
        task_type: str = "complex_task",
        custom_budget: int = None
    ) -> TokenBudget:
        """Allocate a token budget for an agent run."""
        budget_amount = custom_budget or self.default_budgets.get(task_type, 50000)
        
        budget = TokenBudget(
            run_id=run_id,
            allocated=budget_amount
        )
        self.budgets[run_id] = budget
        
        return budget
    
    def track_usage(
        self,
        run_id: str,
        input_tokens: int,
        output_tokens: int,
        operation: str = "llm_call"
    ) -> TokenBudget:
        """Track token usage for a run."""
        budget = self.budgets.get(run_id)
        
        if not budget:
            raise ValueError(f"No budget found for run {run_id}")
        
        total_tokens = input_tokens + output_tokens
        budget.used += total_tokens
        
        # Log usage
        self._log_usage(run_id, input_tokens, output_tokens, operation)
        
        return budget
    
    def check_budget(self, run_id: str) -> tuple[bool, str]:
        """Check if the run is within budget."""
        budget = self.budgets.get(run_id)
        
        if not budget:
            return False, "No budget found"
        
        status = budget.status
        
        if status == BudgetStatus.EXCEEDED:
            return False, f"Budget exceeded: {budget.used}/{budget.allocated} tokens"
        elif status == BudgetStatus.CRITICAL:
            return True, f"Warning: {budget.usage_percent:.1%} of budget used"
        elif status == BudgetStatus.WARNING:
            return True, f"Notice: {budget.usage_percent:.1%} of budget used"
        
        return True, "Within budget"
    
    def _log_usage(self, run_id, input_tokens, output_tokens, operation):
        # Store in database for analysis
        pass

Budget-Aware Agent

class BudgetAwareAgent:
    def __init__(self, llm, budget_manager: TokenBudgetManager):
        self.llm = llm
        self.budget_manager = budget_manager
    
    async def execute(
        self,
        task: str,
        task_type: str = "complex_task",
        budget: int = None
    ) -> dict:
        run_id = str(uuid.uuid4())
        
        # Allocate budget
        token_budget = self.budget_manager.allocate(
            run_id=run_id,
            task_type=task_type,
            custom_budget=budget
        )
        
        messages = [{"role": "user", "content": task}]
        
        while True:
            # Check budget before each LLM call
            within_budget, message = self.budget_manager.check_budget(run_id)
            
            if not within_budget:
                return {
                    "status": "budget_exceeded",
                    "error": message,
                    "tokens_used": token_budget.used,
                    "budget": token_budget.allocated,
                }
            
            # Make LLM call
            response = await self.llm.ainvoke(messages)
            
            # Track usage
            usage = response.usage_metadata
            self.budget_manager.track_usage(
                run_id=run_id,
                input_tokens=usage["input_tokens"],
                output_tokens=usage["output_tokens"],
                operation="llm_call"
            )
            
            # Process response
            if response.tool_calls:
                # Execute tools and track their token usage
                for tool_call in response.tool_calls:
                    tool_result = await self._execute_tool(tool_call)
                    
                    # Tools may use tokens (e.g., embedding calls)
                    if "token_usage" in tool_result:
                        self.budget_manager.track_usage(
                            run_id=run_id,
                            input_tokens=tool_result["token_usage"].get("input", 0),
                            output_tokens=tool_result["token_usage"].get("output", 0),
                            operation=f"tool:{tool_call['name']}"
                        )
                    
                    messages.append({"role": "tool", "content": str(tool_result)})
                
                messages.append(response.message)
            else:
                return {
                    "status": "completed",
                    "result": response.content,
                    "tokens_used": token_budget.used,
                    "budget": token_budget.allocated,
                }

Budget Allocation Strategies

class BudgetAllocator:
    def __init__(self):
        self.task_profiles = {
            "simple_query": {"base": 3000, "multiplier": 1.0},
            "complex_task": {"base": 20000, "multiplier": 1.5},
            "research": {"base": 50000, "multiplier": 2.0},
            "code_generation": {"base": 30000, "multiplier": 1.8},
        }
    
    def allocate(
        self,
        task_type: str,
        complexity: float = 1.0,
        max_budget: int = 100000
    ) -> int:
        """Allocate budget based on task type and complexity."""
        profile = self.task_profiles.get(task_type, {"base": 20000, "multiplier": 1.0})
        
        budget = int(profile["base"] * complexity * profile["multiplier"])
        
        return min(budget, max_budget)
    
    def dynamic_adjustment(
        self,
        run_id: str,
        current_budget: TokenBudget,
        performance_metrics: dict
    ) -> int:
        """Dynamically adjust budget based on performance."""
        # If agent is making progress, allow more tokens
        if performance_metrics.get("steps_completed", 0) > 0:
            progress_rate = (
                performance_metrics["steps_completed"] /
                performance_metrics.get("estimated_steps", 10)
            )
            
            if progress_rate > 0.5:
                # Agent is progressing well, allow 20% more
                return int(current_budget.allocated * 1.2)
        
        # If agent is stuck, reduce budget
        if performance_metrics.get("repeated_failures", 0) > 2:
            return int(current_budget.allocated * 0.8)
        
        return current_budget.allocated

Cost Tracking Dashboard

from fastapi import FastAPI
from datetime import datetime, timedelta

app = FastAPI()

@app.get("/budgets/{run_id}")
async def get_budget_status(run_id: str):
    budget = budget_manager.budgets.get(run_id)
    
    if not budget:
        return {"error": "Run not found"}
    
    return {
        "run_id": run_id,
        "allocated": budget.allocated,
        "used": budget.used,
        "remaining": budget.remaining,
        "usage_percent": f"{budget.usage_percent:.1%}",
        "status": budget.status.value,
        "cost_estimate": self._estimate_cost(budget.used),
    }

@app.get("/costs/summary")
async def get_cost_summary(timeframe: str = "24h"):
    """Get cost summary for the timeframe."""
    since = datetime.now() - timedelta(hours=int(timeframe.replace("h", "")))
    
    usage = await get_usage_since(since)
    
    total_tokens = sum(u["tokens"] for u in usage)
    total_cost = sum(u["cost"] for u in usage)
    
    by_task_type = {}
    for u in usage:
        task_type = u.get("task_type", "unknown")
        if task_type not in by_task_type:
            by_task_type[task_type] = {"tokens": 0, "cost": 0}
        by_task_type[task_type]["tokens"] += u["tokens"]
        by_task_type[task_type]["cost"] += u["cost"]
    
    return {
        "timeframe": timeframe,
        "total_tokens": total_tokens,
        "total_cost": f"${total_cost:.2f}",
        "by_task_type": by_task_type,
        "avg_tokens_per_run": total_tokens / max(len(usage), 1),
    }

def _estimate_cost(self, tokens: int) -> str:
    """Estimate cost based on token count."""
    # GPT-4o pricing
    cost_per_1k = 0.005  # $5/1M tokens
    cost = (tokens / 1000) * cost_per_1k
    return f"${cost:.4f}"

Budget Policies

class BudgetPolicy:
    def __init__(self):
        self.policies = {
            "default": {"daily_limit": 1000000, "per_run_limit": 100000},
            "premium": {"daily_limit": 5000000, "per_run_limit": 500000},
            "free": {"daily_limit": 100000, "per_run_limit": 10000},
        }
    
    def check_policy(self, user_tier: str, run_id: str) -> tuple[bool, str]:
        policy = self.policies.get(user_tier, self.policies["default"])
        
        # Check per-run limit
        budget = budget_manager.budgets.get(run_id)
        if budget and budget.allocated > policy["per_run_limit"]:
            return False, f"Exceeds {user_tier} per-run limit of {policy['per_run_limit']} tokens"
        
        # Check daily limit
        daily_usage = get_daily_usage(user_tier)
        if daily_usage >= policy["daily_limit"]:
            return False, f"Exceeds {user_tier} daily limit of {policy['daily_limit']} tokens"
        
        return True, "Within policy limits"

Best Practices

• Set per-task budgets — Different tasks need different limits
• Monitor in real-time — Track usage as it happens
• Implement warnings — Alert before hitting hard limits
• Log everything — For cost analysis and optimization
• Use tiered policies — Different limits for different user levels
• Auto-adjust — Increase budgets for progressing agents, decrease for stuck ones

Conclusion

Token budgets transform unbounded agent costs into predictable expenses. By allocating budgets per run, tracking usage in real-time, and enforcing hard limits, you maintain control over your AI spend. Start with conservative budgets, adjust based on actual usage patterns, and implement tiered policies for different user levels.