Token Budget Management for Multi-Agent Systems
As large language model (LLM) applications grow more sophisticated, multi-agent architectures have become a popular way to decompose complex tasks into specialized roles. However, every agent invocation consumes tokens — for prompts, context, tool calls, and responses — and those tokens translate directly into latency and cost. Without explicit budget management, a multi-agent system can spiral out of control: agents loop indefinitely, context windows bloat with redundant history, and a single runaway task can burn through an entire API quota in minutes.
Token budget management is the discipline of allocating, tracking, and enforcing token consumption across all agents in a system so that work completes within predictable cost and latency bounds. This tutorial explains what token budgeting means in a multi-agent context, why it matters, how to implement it, and the best practices that separate production-grade systems from fragile prototypes.
What Is Token Budget Management?
At its core, token budget management is the process of assigning a finite token allowance to a unit of work and ensuring that allowance is respected. In a multi-agent system, that unit of work may be a single agent, a sub-task delegated to a group of agents, or an entire user request routed through an orchestrator. The budget covers both input tokens (prompts, retrieved context, conversation history) and output tokens (model completions, tool call arguments).
A well-designed budget system has three components:
- Allocation — deciding how many tokens each agent or task may consume before work begins.
- Accounting — measuring tokens consumed in real time as agents call the model.
- Enforcement — intervening when an agent approaches or exceeds its budget, either by terminating the task, forcing a summary, or escalating to a cheaper model.
Without all three, you either overspend (no enforcement) or incorrectly blame agents for costs they did not incur (bad accounting).
Why It Matters
Multi-agent systems are uniquely vulnerable to token runaway because of three structural properties:
- Compositional blow-up. When an orchestrator delegates to N sub-agents, total cost scales with N. If each sub-agent also calls tools that return large payloads, the multiplier compounds.
- Conversation accumulation. Agents that maintain state across turns carry growing history. Without compaction, the input token cost of each turn grows linearly with the conversation length.
- Retry loops. Agents that fail and retry — especially with self-reflection patterns — can consume several times the expected token count for a single task.
Beyond cost, budgets also protect latency and fairness. A single long-running request should not starve other users of a shared API quota, and a misbehaving agent should not be allowed to retry forever. Budgets turn these implicit risks into explicit, observable constraints.
How to Use It: A Practical Implementation
The following example shows a minimal but complete token budget manager suitable for a multi-agent Python application. It supports per-agent budgets, hierarchical delegation (a parent agent's budget is shared with its children), and enforcement via a BudgetExceeded exception.
from dataclasses import dataclass, field
from typing import Optional
class BudgetExceeded(Exception):
"""Raised when an agent attempts to consume more tokens than allowed."""
pass
@dataclass
class TokenBudget:
"""A token budget that can be shared across agents and sub-tasks."""
total: int
spent: int = 0
parent: Optional["TokenBudget"] = None
@property
def remaining(self) -> int:
return max(0, self.total - self.spent)
def consume(self, n: int) -> None:
"""Record consumption of n tokens, propagating to the parent budget."""
if n <= 0:
return
if self.spent + n > self.total:
raise BudgetExceeded(
f"Budget exceeded: attempted {n}, "
f"spent {self.spent}, total {self.total}"
)
self.spent += n
if self.parent is not None:
# Propagate to parent so orchestrators see child consumption.
self.parent.spent += n
def fork(self, allocation: int) -> "TokenBudget":
"""Create a child budget with its own limit, linked to this parent."""
if allocation > self.remaining:
raise BudgetExceeded(
f"Cannot fork {allocation} tokens, only {self.remaining} remain"
)
return TokenBudget(total=allocation, parent=self)
@dataclass
class Agent:
"""A minimal agent that consumes tokens from a shared budget."""
name: str
model: str
budget: TokenBudget
def call(self, prompt: str, estimated_output: int = 256) -> str:
# In a real system, use the tokenizer of `self.model` to count exactly.
input_tokens = len(prompt.split()) # rough approximation
self.budget.consume(input_tokens + estimated_output)
# ... invoke LLM here ...
return f"[{self.name}] response to: {prompt[:40]}..."
# Example: an orchestrator with a top-level budget delegates to two workers.
root_budget = TokenBudget(total=10_000)
orchestrator = Agent(name="orchestrator", model="gpt-4o", budget=root_budget)
# Each worker gets a slice of the orchestrator's budget.
worker_a = Agent(name="worker-a", model="gpt-4o-mini",
budget=root_budget.fork(3_000))
worker_b = Agent(name="worker-b", model="gpt-4o-mini",
budget=root_budget.fork(3_000))
worker_a.call("Summarize the quarterly report.", estimated_output=500)
worker_b.call("Extract action items from the meeting notes.",
estimated_output=500)
print(f"Root budget remaining: {root_budget.remaining}")
The key design choice is the parent link. When a worker consumes tokens, the cost is also charged against the orchestrator's budget. This means an orchestrator can fork sub-budgets for its children and still observe total consumption in one place, without each child needing to know about the global limit.
Integrating with Real Tokenizers
The example above uses a naive word-count approximation. In production, you should count tokens with the exact tokenizer for the model you are calling. Most providers expose this through their SDKs.
import tiktoken
from openai import OpenAI
client = OpenAI()
enc = tiktoken.encoding_for_model("gpt-4o")
def count_tokens(text: str) -> int:
return len(enc.encode(text))
def chat_with_budget(agent: Agent, messages: list[dict]) -> str:
input_text = "".join(m["content"] for m in messages)
input_tokens = count_tokens(input_text)
# Reserve a conservative output allowance before calling the model.
max_output = 512
agent.budget.consume(input_tokens + max_output)
response = client.chat.completions.create(
model=agent.model,
messages=messages,
max_tokens=max_output,
)
# Refund the difference between reserved and actual output tokens.
actual_output = response.usage.completion_tokens
agent.budget.spent -= (max_output - actual_output)
if agent.budget.parent is not None:
agent.budget.parent.spent -= (max_output - actual_output)
return response.choices[0].message.content
Notice the refund step. Because we cannot know in advance exactly how many output tokens the model will produce, we reserve a conservative allowance and then reconcile against the real usage reported by the API. This keeps accounting accurate without risking a budget violation mid-request.
Enforcement Strategies
Raising an exception is the simplest enforcement mechanism, but it is rarely the best user experience. A production system should degrade gracefully. Common strategies include:
- Summarize and continue. When an agent's context approaches its budget, replace the full history with a compressed summary, freeing tokens for further work.
- Downgrade the model. Switch from a premium model to a cheaper, faster one when the remaining budget is low, accepting lower quality to stay within limits.
- Return partial results. Let the agent emit whatever it has produced so far, rather than discarding the work entirely.
- Escalate to the user. For interactive systems, ask the user whether to extend the budget rather than silently failing.
The following snippet shows a guard that applies the summarize-and-continue strategy when an agent's context grows too large.
def maybe_compact(agent: Agent, messages: list[dict],
summarizer: Agent, threshold: float = 0.8) -> list[dict]:
"""Compact conversation history when input tokens exceed threshold."""
input_tokens = count_tokens("".join(m["content"] for m in messages))
if input_tokens < threshold * agent.budget.total:
return messages # plenty of room
# Keep the system prompt and the most recent user turn; summarize the rest.
system = [m for m in messages if m["role"] == "system"]
recent = messages[-1:]
history = messages[len(system):-1]
summary_prompt = "Summarize the following conversation concisely:\n" + \
"\n".join(m["content"] for m in history)
summary = summarizer.call(summary_prompt, estimated_output=300)
return system + [{"role": "system", "content": f"Prior context: {summary}"}] + recent
Best Practices
- Allocate budgets per request, not per agent. A long-lived agent that serves many requests should not accumulate a single growing budget. Create a fresh budget for each user request and fork sub-budgets as needed.
- Use hierarchical budgets for delegation. When an orchestrator spawns workers, fork the parent budget so that total consumption is visible at the top level and no single worker can starve the others.
- Count input and output separately. Output tokens are typically more expensive than input tokens on modern APIs. Track them independently so your cost projections are accurate.
- Reserve before you call, reconcile after. Reserve a conservative output allowance to avoid mid-request violations, then refund the difference using the provider's reported usage.
- Log every consumption event. Structured logs of who consumed what, when, and for which task are essential for debugging runaway costs and for building accurate cost models.
- Set hard ceilings and soft warnings. Emit a warning at 80% of budget and enforce a hard stop at 100%. This gives observability tooling a chance to alert before failure.
- Test budget enforcement under load. Retry loops and tool-call fan-out are the most common causes of budget overruns. Write integration tests that exercise these paths and assert the budget is respected.
- Cache aggressively. Identical sub-agent calls across requests should be memoized. A cache hit consumes zero new tokens and is the cheapest form of budget management available.
Conclusion
Token budget management is not an optional optimization for multi-agent systems — it is a foundational reliability concern. By treating tokens as a first-class, finite resource with explicit allocation, accurate accounting, and graceful enforcement, you turn unpredictable cost and latency into observable, controllable engineering parameters. The patterns shown here — hierarchical budgets, pre-call reservation with post-call reconciliation, context compaction, and model downgrade — compose into a system that can scale to many agents and many users without surprises. Start with a simple TokenBudget class, instrument every model call, and refine your enforcement strategies as you learn where your real-world workloads actually spend their tokens.