Repository navigation
[Feature] Per-agent token/cost budget guardrail (enforcement, not just reporting) #1065
Description
Search before asking
- I searched in the issues and found nothing similar.
Description
#858 tracks embedding models not consistently reporting token usage metrics, and chat models already have a working token-accounting path that lands in the metrics layer. That's visibility. There's no enforcement layer on top of it.
What problem does it solve: a Flink Agents job is a long-running streaming application, not a one-shot script. A misbehaving action (a bad prompt loop, an unexpectedly chatty tool-call cycle, a runaway retry) can burn LLM spend continuously for as long as the job stays up, with the only signal being a metrics dashboard someone has to be watching. There's currently no built-in way to say "stop calling this model for this agent/action once N tokens (or $ cost) have been spent in a window" the way, for example, Flink itself lets you bound resource usage with slot/memory limits.
What I have in mind: a budget guardrail configured per agent (or per action/chat-model-setup), expressed as a token or cost ceiling over a sliding or fixed window, enforced on the same durable, metered chat/embedding path the token-accounting and #1062's MODEL_ROUTER work already run through — so it's a policy layered on infrastructure that already exists, not a new call path. When the ceiling is hit, the configured behavior could range from a raised error (fail the action, let normal error-handling strategy take over, same as a routing/judge failure already does in #1062) to a soft warning event for observability-only setups.
This is a runtime safety feature, distinct from #1062's routing/fallback-among-candidates concern — routing decides which model answers a call; this decides whether the call is still allowed to happen given spend so far.
Are you willing to submit a PR?
- I'm willing to submit a PR!
Activity
Ashfaqbs commentedon Sep 4, 2026
Proposed design
Traced where the enforcement hook naturally fits with the current code:
ChatModelAction.recordChatTokenMetrics(), feedingBaseChatModelSetup.recordTokenMetrics()— but that's Flink-metric counters for dashboards only, not state that can gate the next call (and metrics aren't reliably queryable/restorable for decision logic across checkpoints).ChatModelInvoker.chatWithRetries(), on the durable/keyed execution path (ctx.durableExecute). That's the natural enforcement point.Sketch:
AgentExecutionOptions, matching the existingConfigOption<>pattern (MAX_RETRIES,CHAT_ASYNC, etc.):chat.budget.max-tokens(or cost) — ceiling per window, disabled by default so it's zero-impact for existing userschat.budget.window-sec— sliding/fixed window sizechat.budget.on-exceeded—FAIL(raise, letting the existingErrorHandlingStrategyhandle it the same way any other attempt failure already does) orWARN(emit an observability event only)ValueState-backed total (tokens spent + window start) checkpoints/restores correctly, where a metric counter would not.chatWithRetries, beforedurableExecute: reset if the window has expired; if spend would exceed budget, either throw (folds into the existingChatAttemptFailed→ErrorHandlingStrategypath, reusing machinery instead of adding a new failure mode) or emit a warning and proceed.recordChatTokenMetricscall, right after a successful attempt.This is meant to be a thin policy layer on infrastructure that already exists — no new call path, reuses the existing retry/error-handling machinery, opt-in by default.
Before I put together a PR: this touches checkpointed-state semantics (window reset behavior across job restarts, what "per agent" scopes to under keyed parallelism/rescaling). Wanted to get a read on that from a maintainer before writing code, rather than build it cold and risk a rewrite. Does the sketch above look right, or would you scope the state/window semantics differently?
wenjin272 commentedon Sep 10, 2026
Thanks for raising this and for putting together the initial proposal. I agree that budget enforcement is valuable, especially for continuously running streaming jobs where an unnoticed loop or retry could cause significant token consumption.
I have a few initial thoughts on the design. We may need token limits at different scopes, such as per model request, per ReAct loop, per time window, and per subtask. A job-wide per-agent budget may be difficult to enforce correctly because the agent runs in a distributed environment. The current proposal also needs to clarify its error-handling behavior and persistence and recovery semantics.
Following the community’s current practice, I suggest opening a GitHub Discussion with a detailed design covering these semantics and trade-offs before moving to implementation. Please note that the community is currently focused on developing the planned features for the 0.4 release, so we may not have enough bandwidth to discuss the proposal in depth before the release.