← Back to DevBytes

Token Budget Management for Multi-Agent Systems

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:

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:

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:

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

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.

— Ad —

Google AdSense will appear here after approval

← Back to all articles