A proposal to add token-usage budgets to LangChain agents via a new middleware, since the existing limiters only cap call count, not tokens

OSS Product Help LangChain
1

Feature Description

LangChain’s agent middleware suite can cap the number of model calls

(ModelCallLimitMiddleware) and tool calls (ToolCallLimitMiddleware), but nothing bounds

token consumption. I’d like a TokenBudgetMiddleware that accumulates token usage across

model calls (from AIMessage.usage_metadata) and halts the agent — via jump_to=“end” or

a TokenBudgetExceededError — when a configured budget is exceeded. It would support

independent input/output/total limits at both thread and run scope, with sync and async

hooks, mirroring ModelCallLimitMiddleware’s structure and exit_behavior contract.

Use Case

Production agents can consume a very large token volume within a small number of calls

(context grows with tool outputs and retrieved documents), so a call-count cap doesn’t bound

spend or context blowup. SummarizationMiddleware and ContextEditingMiddleware reduce

context reactively but don’t enforce a hard ceiling. Today users hand-roll before_model /

after_model hooks re-reading usage_metadata and managing state manually. A reusable,

provider-agnostic budget guard would cover a common production need.

Proposed Solution

from langchain.agents.middleware import TokenBudgetMiddleware

from langchain.agents import create_agent

budget = TokenBudgetMiddleware(

thread_total_limit=200_000,

run_total_limit=50_000,

exit_behavior=“end”, # or “error”

)

agent = create_agent(“openai:gpt-5.5”, middleware=[budget])
  • after_model reads usage_metadata off the latest AIMessage and accumulates

input/output/total tokens for thread + run scope (defensive: missing usage = zero

contribution, never raises).

  • before_model (with @hook_config(can_jump_to=[“end”])) checks accumulated usage against

any set limit; on breach it injects an explanatory AIMessage and jumps to end, or raises

TokenBudgetExceededError, per exit_behavior.

  • State extends AgentState with private token-count fields, reusing the same

PrivateStateAttr / UntrackedValue channel annotations ModelCallLimitMiddleware uses

for thread- vs run-scoped persistence.

  • Scope is token-only; monetary cost (which needs drifting per-model pricing tables) is left

as a possible follow-up.

2

Hey, @Juyotal! This is an interesting request and addresses a real gap.

If you’re interested in seeing something like this implemented, I would recommend opening a feature request on our GitHub repo.

Just note that our contribution flow is issue-first:

  1. File it at github.com/langchain-ai/langchain/issues as a feature request. Your write-up above is already 90% of the way there.
  2. The team reviews; if approved, they assign an owner.
  3. The PR comes after that.

One heads-up: PRs without an approved, assigned issue get auto-closed, so please hold off on opening one until the issue is reviewed.

In the meantime, you already called out the hand-rolled route. A before_model hook with @hook_config(can_jump_to=["end"]) reading accumulated usage_metadata is precisely the pattern in the custom middleware docs, so you can run a budget guard today while the built-in is under review.

3

Thank you for the reply, @dariel.datoon will create the issue for it right away.

4

Agree this is a real production gap. Call-count limits don’t help much when a single step pulls in large tool outputs or retrieved context.

One complementary angle we’ve been using: reduce the tokens that even reach the model in the first place, especially the ones that are already stale or low-value.

A lot of the context growth in longer agent runs comes from retrieved documents and tool results that are no longer current. We’ve been running a thin OpenAI-compatible proxy (KU-Gateway) that scores those blocks on temporal decay and drops the high-decay ones before they enter the LLM. It doesn’t replace a proper token budget, but it meaningfully lowers the baseline token consumption that the budget then has to police.

Happy to share how we’re measuring the reduction if useful. And +1 on the TokenBudgetMiddleware proposal ,having both a hard ceiling and a way to keep low-value context out seems like the right combination.