How do you manage shared token budgets across multiple AI agents?
A shared token budget is a limit on how many LLM input and output tokens several agents may consume together for one workflow, tenant, or accounting period. Unlike a per-minute rate limit, the budget answers: *How much model usage can this whole job afford?*
Beginner-friendly example: API test generation
Suppose an orchestrator has a 100,000-token budget for generating test suites from API requirements. It assigns:
- Requirements agent: Up to 20,000 tokens.
- Test-generation agent: Up to 55,000 tokens.
- Review agent: Up to 25,000 tokens.
The allocations total 100,000 tokens. Each agent must request a reservation from a shared budget manager before an LLM call, because several branches may run at once.
Define total budget → Reserve estimated usage → Call model → Record actual usage → Release unused reservation → Stop or replan at the limit
Step-by-step numerical example
Suppose the requirements agent reserves 20,000 tokens but consumes 14,000:
- The shared budget begins at 100,000 tokens.
- Reserving 20,000 leaves 80,000 unreserved.
- After the call, the usage record reports 14,000 tokens.
- The manager finalizes 14,000 spent and returns 6,000 unused to the pool.
- With no other usage, 86,000 tokens remain available.
Now suppose two agents simultaneously ask to reserve 60,000 and 40,000 tokens. The total request exceeds the remaining 86,000. An atomic reservation admits only allocations that fit; it must not let each agent independently assume the full remainder is available.
How do you manage budgets safely?
- Choose the budget scope: Set limits per workflow, tenant, project, model, or billing period.
- Track input and output separately: Input context and generated responses may have different costs; model pricing can also differ.
- Reserve before every call: Estimate prompt tokens and cap maximum generated tokens, including a margin for uncertainty.
- Use atomic updates: Store reservations and actual usage in a transactional database or trusted budget service.
- Reconcile provider usage: Settle reservations using actual usage reports when available; handle missing reports conservatively.
- Set priorities: Protect essential validation and security-review stages instead of letting an early exploratory agent spend the whole budget.
- Define exhaustion behavior: Shorten optional work, request approval for more budget, or mark the workflow BUDGET_EXCEEDED rather than silently skipping mandatory checks.
- Monitor spending: Record tokens used per agent, model, task, and run without storing sensitive prompt contents unnecessarily.
Example shared budget record
JSON{ "workflow_id": "api-tests-run-42", "token_limit": 100000, "tokens_spent": 14000, "tokens_reserved": 0, "tokens_available": 86000 }
This simplified record assumes no other in-flight reservations. A real system should also track reservation IDs, expiry, call attempts, and reconciliation status.
What about retries and failed calls?
A failed or timed-out request may still have consumed billable tokens. Do not automatically refund its entire reservation just because the client did not receive a response. Reconcile with provider usage where possible and retain a conservative unresolved charge until the outcome is known.
Budget enforcement should also consider actual currency cost: 10,000 tokens on one model may cost more than 10,000 on another.
LangGraph and Celery implementation
In LangGraph, the supervisor can pass a workflow budget identifier to specialist nodes and consult a shared budget manager before model invocations. Checkpoints preserve workflow progress, but do not replace atomic budget accounting.
With Celery, multiple workers need one shared ledger. Per-worker counters cannot enforce a workflow-wide token ceiling. Use stable request IDs to avoid double-accounting when tasks are retried or messages are redelivered.
Shared token budget vs. distributed rate limit
- Rate limit: Controls request or token *throughput* within a time window.
- Token budget: Controls the *total consumption* allowed for a task, workflow, or account.
A production system often needs both controls.
In short
Manage a shared token budget with explicit allocations, atomic reservations, actual-usage reconciliation, and a clear exhaustion policy. This keeps parallel agents from unexpectedly overspending while protecting required workflow stages.
Learning Sequence
AI Agents Curriculum
Explore Related Questions
Deepen your understanding across related AI Agents topics.