From 78fff5a5e2048c1a3e0f664d2ab7e61355e68aa1 Mon Sep 17 00:00:00 2001 From: Chetan Sharma Date: Sun, 21 Jun 2026 19:55:15 +0530 Subject: [PATCH] =?UTF-8?q?feat(middleware):=20add=20TokenBudgetMiddleware?= =?UTF-8?q?=20for=20per-run=20token=20budget=20e=E2=80=A6=20(#3412)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * eat(middleware): add TokenBudgetMiddleware for per-run token budget enforcement * address copilot comments * resolve feedback --------- Co-authored-by: Willem Jiang --- .../harness/deerflow/agents/factory.py | 12 +- .../harness/deerflow/agents/features.py | 1 + .../deerflow/agents/lead_agent/agent.py | 7 + .../middlewares/token_budget_middleware.py | 286 ++++++++++++++++++ .../harness/deerflow/config/app_config.py | 2 + .../deerflow/config/token_budget_config.py | 21 ++ backend/tests/test_create_deerflow_agent.py | 2 +- backend/tests/test_token_budget_middleware.py | 165 ++++++++++ config.example.yaml | 18 +- 9 files changed, 511 insertions(+), 3 deletions(-) create mode 100644 backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py create mode 100644 backend/packages/harness/deerflow/config/token_budget_config.py create mode 100644 backend/tests/test_token_budget_middleware.py diff --git a/backend/packages/harness/deerflow/agents/factory.py b/backend/packages/harness/deerflow/agents/factory.py index e847023b9..e75b32436 100644 --- a/backend/packages/harness/deerflow/agents/factory.py +++ b/backend/packages/harness/deerflow/agents/factory.py @@ -282,7 +282,17 @@ def _assemble_from_features( chain.append(LoopDetectionMiddleware.from_config(LoopDetectionConfig())) - # --- [13] Clarification (always last among built-ins) --- + # --- [13] TokenBudget --- + if feat.token_budget is not False: + if isinstance(feat.token_budget, AgentMiddleware): + chain.append(feat.token_budget) + else: + from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware + from deerflow.config.token_budget_config import TokenBudgetConfig + + chain.append(TokenBudgetMiddleware.from_config(TokenBudgetConfig())) + + # --- [14] Clarification (always last among built-ins) --- chain.append(ClarificationMiddleware()) extra_tools.append(ask_clarification_tool) diff --git a/backend/packages/harness/deerflow/agents/features.py b/backend/packages/harness/deerflow/agents/features.py index db2bc9580..5aca6cf06 100644 --- a/backend/packages/harness/deerflow/agents/features.py +++ b/backend/packages/harness/deerflow/agents/features.py @@ -32,6 +32,7 @@ class RuntimeFeatures: auto_title: bool | AgentMiddleware = False guardrail: Literal[False] | AgentMiddleware = False loop_detection: bool | AgentMiddleware = True + token_budget: bool | AgentMiddleware = False # --------------------------------------------------------------------------- diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index cc3eee449..586a5a9a5 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -359,6 +359,13 @@ def build_middlewares( if loop_detection_config.enabled: middlewares.append(LoopDetectionMiddleware.from_config(loop_detection_config)) + # TokenBudgetMiddleware - enforce per-run token limits + token_budget_config = resolved_app_config.token_budget + if token_budget_config.enabled: + from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware + + middlewares.append(TokenBudgetMiddleware.from_config(token_budget_config)) + # Inject custom middlewares before ClarificationMiddleware if custom_middlewares: middlewares.extend(custom_middlewares) diff --git a/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py new file mode 100644 index 000000000..a9fe9ae4d --- /dev/null +++ b/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py @@ -0,0 +1,286 @@ +"""Middleware to enforce per-run token budget limits. +Tracks cumulative token usage (input, output, total) across model calls within +a single agent run and enforces configurable soft-warning and hard-stop +thresholds. +Detection strategy: + 1. After each model response, sum the `usage_metadata` of all `AIMessage`s + in the current thread history. This automatically captures tokens from + subagents because `TokenUsageMiddleware` retroactively adds them to the + history. + 2. If the highest fraction (input, output, or total) >= warn_threshold, + queue a warning. + 3. If the highest fraction >= hard_stop_threshold, strip tool_calls. +Warning injection uses the deferred pattern: + - after_model queues the warning (does NOT mutate state). + - wrap_model_call injects it as a HumanMessage at the next model call. +This preserves AIMessage(tool_calls) → ToolMessage pairing. +""" + +from __future__ import annotations + +import logging +import threading +from collections import OrderedDict +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, override + +from langchain.agents import AgentState +from langchain.agents.middleware import AgentMiddleware +from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.runtime import Runtime + +from deerflow.config.token_budget_config import TokenBudgetConfig + +logger = logging.getLogger(__name__) + +_BUDGET_WARNING_MSG = ( + "[TOKEN BUDGET WARNING] You have used {used:,} of your {budget:,} {reason} token budget ({percent:.0f}%). Wrap up your current work and produce a final answer. Avoid starting new tool calls unless absolutely necessary." +) +_BUDGET_EXCEEDED_MSG = "[TOKEN BUDGET EXCEEDED] The {reason} token usage ({used:,}) has exceeded the safety limit ({budget:,}). Producing final answer with results collected so far." + + +@dataclass +class TokenUsage: + input: int = 0 + output: int = 0 + total: int = 0 + + +class BoundedDict(OrderedDict): + """A bounded dictionary to prevent unbounded state growth on abandoned runs.""" + + def __init__(self, maxsize=1000, *args, **kwds): + self.maxsize = maxsize + super().__init__(*args, **kwds) + + def __setitem__(self, key, value): + if key not in self: + if len(self) >= self.maxsize: + self.popitem(last=False) + super().__setitem__(key, value) + + +class TokenBudgetMiddleware(AgentMiddleware[AgentState]): + """Enforce per-run token budget limits.""" + + def __init__(self, config: TokenBudgetConfig) -> None: + super().__init__() + self._config = config + self._lock = threading.Lock() + + # Keyed strictly by run_id (clobber-safe) and bounded (leak-safe) + self._warned: BoundedDict[str, bool] = BoundedDict(1000) + self._pending_warnings: BoundedDict[str, list[str]] = BoundedDict(1000) + self._seen_messages: BoundedDict[str, dict[str, tuple[int, int]]] = BoundedDict(1000) + self._cumulative_usage: BoundedDict[str, TokenUsage] = BoundedDict(1000) + + @classmethod + def from_config(cls, config: TokenBudgetConfig) -> TokenBudgetMiddleware: + return cls(config=config) + + def reset(self) -> None: + with self._lock: + self._warned.clear() + self._pending_warnings.clear() + self._seen_messages.clear() + self._cumulative_usage.clear() + + @staticmethod + def _get_run_id(runtime: Runtime) -> str: + ctx = getattr(runtime, "context", None) + if isinstance(ctx, dict) and "run_id" in ctx: + return ctx["run_id"] + # Fallback to runtime object ID to prevent collisions across embedded client runs + return str(id(runtime)) + + def _clear_run_state(self, run_id: str) -> None: + self._warned.pop(run_id, None) + self._pending_warnings.pop(run_id, None) + self._seen_messages.pop(run_id, None) + self._cumulative_usage.pop(run_id, None) + + @override + def before_agent(self, state: AgentState, runtime: Runtime) -> None: + if not self._config.enabled: + return + + # Mark all old messages from previous runss as 'seen' so they don't count toward THIS run's budget + messages = state.get("messages", []) + if not messages: + return + + run_id = self._get_run_id(runtime) + seen = self._seen_messages.setdefault(run_id, {}) + self._cumulative_usage.setdefault(run_id, TokenUsage()) + + for msg in messages: + if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"): + usage = msg.usage_metadata or {} + input_tokens = usage.get("input_tokens", 0) + output_tokens = usage.get("output_tokens", 0) + seen[msg.id] = (input_tokens, output_tokens) + + @override + async def abefore_agent(self, state: AgentState, runtime: Runtime) -> None: + self.before_agent(state, runtime) + + @override + def after_agent(self, state: AgentState, runtime: Runtime) -> None: + if not self._config.enabled: + return + self._clear_run_state(self._get_run_id(runtime)) + + @override + async def aafter_agent(self, state: AgentState, runtime: Runtime) -> None: + self.after_agent(state, runtime) + + @staticmethod + def _append_text(content: str | list[dict | None] | None, stop_msg: str) -> str | list[dict | str]: + """Append a stop message to an AIMessage.content field.""" + if content is None: + return stop_msg + if isinstance(content, str): + if content: + return f"{content}\n\n{stop_msg}" + return f"\n\n{stop_msg}" + if isinstance(content, list): + new_content = list(content) + new_content.append({"type": "text", "text": f"\n\n{stop_msg}"}) + return new_content + return f"{content}\n\n{stop_msg}" + + def _build_hard_stop_update(self, msg: AIMessage, stop_msg: str) -> dict[str, Any]: + """Build the state update dictionary for a hard stop.""" + updated_content = self._append_text(msg.content, stop_msg) + kwargs = dict(msg.additional_kwargs) if msg.additional_kwargs else {} + if "tool_calls" in kwargs: + del kwargs["tool_calls"] + if "function_call" in kwargs: + del kwargs["function_call"] + + response_metadata = dict(getattr(msg, "response_metadata", {}) or {}) + + if response_metadata.get("finish_reason") == "tool_calls": + response_metadata["finish_reason"] = "stop" + + stopped_msg = msg.model_copy(update={"content": updated_content, "tool_calls": [], "additional_kwargs": kwargs, "response_metadata": response_metadata}) + return {"messages": [stopped_msg]} + + def _apply(self, state: AgentState, runtime: Runtime) -> dict | None: + if not self._config.enabled: + return None + + messages = state.get("messages", []) + if not messages: + return None + + last_msg = messages[-1] + if not isinstance(last_msg, AIMessage): + return None + + run_id = self._get_run_id(runtime) + + seen = self._seen_messages.setdefault(run_id, {}) + usage_accum = self._cumulative_usage.setdefault(run_id, TokenUsage()) + + for msg in messages: + if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"): + usage = msg.usage_metadata or {} + + input_tokens = usage.get("input_tokens", 0) + output_tokens = usage.get("output_tokens", 0) + + # Check what previously recorded for this exact message + prev_input, prev_output = seen.get(msg.id, (0, 0)) + + # Calculate if any new tokens were added (handles retroactive subagent tokens) + diff_input = max(0, input_tokens - prev_input) + diff_output = max(0, output_tokens - prev_output) + + if diff_input > 0 or diff_output > 0: + usage_accum.input += diff_input + usage_accum.output += diff_output + usage_accum.total += diff_input + diff_output + seen[msg.id] = (input_tokens, output_tokens) + + if usage_accum.total <= 0: + return None + + fractions = [("total", usage_accum.total, self._config.max_tokens)] + if self._config.max_input_tokens: + fractions.append(("input", usage_accum.input, self._config.max_input_tokens)) + if self._config.max_output_tokens: + fractions.append(("output", usage_accum.output, self._config.max_output_tokens)) + + highest_fraction = 0.0 + trigger_reason = "" + trigger_used = 0 + trigger_budget = 0 + + for reason, used, limit in fractions: + frac = used / limit + if frac > highest_fraction: + highest_fraction = frac + trigger_reason = reason + trigger_used = used + trigger_budget = limit + + if highest_fraction >= self._config.hard_stop_threshold: + logger.warning("Token budget hard stop triggered for run %s: %s limit exceeded", run_id, trigger_reason) + stop_text = _BUDGET_EXCEEDED_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget) + return self._build_hard_stop_update(last_msg, stop_text) + + if highest_fraction >= self._config.warn_threshold and not self._warned.get(run_id, False): + self._warned[run_id] = True + percent = highest_fraction * 100 + warn_text = _BUDGET_WARNING_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget, percent=percent) + logger.info("Token budget warning triggered for run %s: %s limit at %.1f%%", run_id, trigger_reason, percent) + # queue warning for wrap_model_call + warnings = self._pending_warnings.setdefault(run_id, []) + warnings.append(warn_text) + return None + + return None + + @override + def after_model(self, state: AgentState, runtime: Runtime) -> dict | None: + return self._apply(state, runtime) + + @override + async def aafter_model(self, state: AgentState, runtime: Runtime) -> dict | None: + return self._apply(state, runtime) + + def _drain_pending_warnings(self, runtime: Runtime) -> list[str]: + if not self._config.enabled: + return [] + + run_id = self._get_run_id(runtime) + warnings = self._pending_warnings.pop(run_id, None) + return warnings or [] + + def _inject_warnings(self, request: ModelRequest, warnings: list[str]) -> ModelRequest: + if not warnings: + return request + + merged_text = "\n\n".join(warnings) + warning_msg = HumanMessage(content=merged_text, name="budget_warning") + + messages = getattr(request, "messages", []) + new_messages = list(messages) + [warning_msg] + return request.override(messages=new_messages) + + @override + def wrap_model_call(self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelCallResult: + + warnings = self._drain_pending_warnings(request.runtime) + request = self._inject_warnings(request, warnings) + + return handler(request) + + @override + async def awrap_model_call(self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]]) -> ModelCallResult: + warnings = self._drain_pending_warnings(request.runtime) + request = self._inject_warnings(request, warnings) + return await handler(request) diff --git a/backend/packages/harness/deerflow/config/app_config.py b/backend/packages/harness/deerflow/config/app_config.py index fcfa44fb3..6e796a5f2 100644 --- a/backend/packages/harness/deerflow/config/app_config.py +++ b/backend/packages/harness/deerflow/config/app_config.py @@ -33,6 +33,7 @@ from deerflow.config.subagents_config import SubagentsAppConfig, load_subagents_ from deerflow.config.suggestions_config import SuggestionsConfig from deerflow.config.summarization_config import SummarizationConfig, load_summarization_config_from_dict from deerflow.config.title_config import TitleConfig, load_title_config_from_dict +from deerflow.config.token_budget_config import TokenBudgetConfig from deerflow.config.token_usage_config import TokenUsageConfig from deerflow.config.tool_config import ToolConfig, ToolGroupConfig from deerflow.config.tool_output_config import ToolOutputConfig @@ -98,6 +99,7 @@ class AppConfig(BaseModel): ), ) token_usage: TokenUsageConfig = Field(default_factory=TokenUsageConfig, description="Token usage tracking configuration") + token_budget: TokenBudgetConfig = Field(default_factory=TokenBudgetConfig, description="Token Budget tracking and limits configuration.") models: list[ModelConfig] = Field(default_factory=list, description="Available models") sandbox: SandboxConfig = Field( description=format_field_description( diff --git a/backend/packages/harness/deerflow/config/token_budget_config.py b/backend/packages/harness/deerflow/config/token_budget_config.py new file mode 100644 index 000000000..32dc21bc0 --- /dev/null +++ b/backend/packages/harness/deerflow/config/token_budget_config.py @@ -0,0 +1,21 @@ +"""Config for token budget middleware.""" + +from pydantic import BaseModel, Field, model_validator + + +class TokenBudgetConfig(BaseModel): + """Configuration for per-run token budget enforcement.""" + + enabled: bool = Field(default=False, description="Whether to enable per-run token budget enforcement.") + max_tokens: int = Field(default=200000, ge=1000, description="Maximum total tokens (input + output) allowed per run.") + max_input_tokens: int | None = Field(default=None, ge=1, description="Optional separate limit for input tokens only.") + max_output_tokens: int | None = Field(default=None, ge=1, description="Optional separate limit for output tokens only.") + warn_threshold: float = Field(default=0.8, ge=0.0, le=1.0, description="Fraction of max_tokens at which a soft warning is injected. E.g., 0.8 means warn at 80% of max_tokens") + hard_stop_threshold: float = Field(default=1.0, ge=0.0, le=1.0, description=("Fraction of max_tokens at which tool calls are stripped and the agent is forced to produce a final answer. E.g., 1.0 means stop at 100% of max_tokens.")) + + @model_validator(mode="after") + def validate_thresholds(self) -> "TokenBudgetConfig": + """Ensure hard stop cannot trigger before the warning.""" + if self.hard_stop_threshold < self.warn_threshold: + raise ValueError("hard_stop_threshold must be >= warn_threshold") + return self diff --git a/backend/tests/test_create_deerflow_agent.py b/backend/tests/test_create_deerflow_agent.py index d0fc6c9a8..73ec09a48 100644 --- a/backend/tests/test_create_deerflow_agent.py +++ b/backend/tests/test_create_deerflow_agent.py @@ -671,7 +671,7 @@ def test_loop_detection_custom_middleware(mock_create_agent): mw_types = [type(m).__name__ for m in middleware] # Default LoopDetectionMiddleware must not also appear. assert "LoopDetectionMiddleware" not in mw_types - # Custom replacement still sits immediately before ClarificationMiddleware. + # Custom replacement sits immediately before TokenBudgetMiddleware and ClarificationMiddleware. assert mw_types[-1] == "ClarificationMiddleware" assert mw_types[-2] == "MyLoopDetection" diff --git a/backend/tests/test_token_budget_middleware.py b/backend/tests/test_token_budget_middleware.py new file mode 100644 index 000000000..9d68d978b --- /dev/null +++ b/backend/tests/test_token_budget_middleware.py @@ -0,0 +1,165 @@ +from unittest.mock import MagicMock + +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware +from deerflow.config.token_budget_config import TokenBudgetConfig + + +def _make_runtime(thread_id="test-thread", run_id="test-run"): + runtime = MagicMock() + runtime.context = {"thread_id": thread_id, "run_id": run_id} + return runtime + + +def _make_request(messages, runtime): + request = MagicMock() + request.messages = list(messages) + request.runtime = runtime + + def override_fn(messages=None, **kwags): + new_req = MagicMock() + new_req.messages = messages if messages is not None else request.messages + new_req.runtime = request.runtime + return new_req + + request.override = override_fn + return request + + +def _capture_handler(): + captured: list = [] + + def handler(req): + captured.append(req) + return MagicMock() + + return captured, handler + + +def _make_state_with_usage(total: int, input_tk: int = 0, output_tk: int = 0, tool_calls=None, content=""): + """Build a state dict with a single AIMessage containing usage.""" + if input_tk == 0 and output_tk == 0: + input_tk = total + msg = AIMessage(id="test-msg", content=content, tool_calls=tool_calls or [], usage_metadata={"input_tokens": input_tk, "output_tokens": output_tk, "total_tokens": total}) + return {"messages": [msg]} + + +class TestTokenBudgetTracking: + def test_no_usage_metadata_returns_none(self): + config = TokenBudgetConfig(max_tokens=1000, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + + state = {"messages": [AIMessage(content="hello", tool_calls=[])]} + result = mw._apply(state, _make_runtime()) + assert result is None + + def test_below_threshold_returns_none(self): + config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + + state = _make_state_with_usage(total=50000) + result = mw._apply(state, _make_runtime()) + assert result is None + + def test_warning_threshold_injects_warning_and_returns_none(self): + config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + + # history with multiple AIMessages that add up to 85000 tokens (>80%) + msg1 = AIMessage(id="msg1", content="1", usage_metadata={"total_tokens": 45000, "input_tokens": 45000, "output_tokens": 0}) + msg2 = ToolMessage(content="ok", tool_call_id="call1") + msg3 = AIMessage(id="msg3", content="3", usage_metadata={"total_tokens": 45000, "input_tokens": 45000, "output_tokens": 0}) + + state = {"messages": [msg1, msg2, msg3]} + result = mw._apply(state, _make_runtime()) + + # should queue warning but not mutate state (return None) + assert result is None + assert len(mw._pending_warnings["test-run"]) == 1 + assert "TOKEN BUDGET WARNING" in mw._pending_warnings["test-run"][0] + + +class TestTokenBudgetWarning: + def test_warn_injected_at_next_model_call(self): + config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + runtime = _make_runtime() + + # trigger warning queue + mw._apply(_make_state_with_usage(total=85000), runtime) + + ai_msg = AIMessage(content="", tool_calls=[{"name": "test", "args": {}, "id": "1"}]) + tool_msg = ToolMessage(content="ok", tool_call_id="1") + + request = _make_request([ai_msg, tool_msg], runtime) + + captured, handler = _capture_handler() + mw.wrap_model_call(request, handler) + + sent = captured[0].messages + + assert sent[0] is ai_msg + assert sent[1] is tool_msg + assert isinstance(sent[2], HumanMessage) + assert sent[2].name == "budget_warning" + assert "TOKEN BUDGET WARNING" in sent[2].content + + def test_warn_only_once_per_run(self): + config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + runtime = _make_runtime() + + mw._apply(_make_state_with_usage(total=85000), runtime) + + assert len(mw._pending_warnings["test-run"]) == 1 + + # call 2: still above threshold, but already warning -> no second enqueue + mw._apply(_make_state_with_usage(total=90000), runtime) + assert len(mw._pending_warnings["test-run"]) == 1 + + +class TestTokenBudgetHardStop: + def test_hard_stop_strip_tool_calls(self): + config = TokenBudgetConfig(max_tokens=100000, hard_stop_threshold=1.0, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + + tool_calls = [{"name": "bash", "args": {"command": "ls"}, "id": "call_1"}] + state = _make_state_with_usage(total=105000, tool_calls=tool_calls, content="Thinking") + + res = mw._apply(state, _make_runtime()) + + assert res is not None + msgs = res["messages"] + assert len(msgs) == 1 + + # tool calls must be stripped + assert msgs[0].tool_calls == [] + # content must have the warning appended + assert "Thinking" in msgs[0].content + assert "TOKEN BUDGET EXCEEDED" in msgs[0].content + + +class TestIndependentDimensions: + def test_input_tokens_trigger_limit(self): + config = TokenBudgetConfig(max_tokens=100000, max_input_tokens=10000, warn_threshold=0.8, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + + # total is safe (10k < 100k) but input is over limit (9k >= 8k) + state = _make_state_with_usage(total=10000, input_tk=9000, output_tk=1000) + mw._apply(state, _make_runtime()) + + warnings = mw._pending_warnings["test-run"] + assert len(warnings) == 1 + assert "input token" in warnings[0] + + def test_output_tokens_trigger_limit(self): + config = TokenBudgetConfig(max_tokens=100_000, max_output_tokens=5_000, hard_stop_threshold=1.0, enabled=True) + mw = TokenBudgetMiddleware.from_config(config) + + # Total is safe (10k < 100k) but output is over hard limit (6k >= 5k) + state = _make_state_with_usage(total=10_000, input_tk=4000, output_tk=6000) + result = mw._apply(state, _make_runtime()) + + assert result is not None + assert "output token" in result["messages"][0].content diff --git a/config.example.yaml b/config.example.yaml index d629cbc8c..e4176d5d0 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -15,7 +15,7 @@ # ============================================================================ # Bump this number when the config schema changes. # Run `make config-upgrade` to merge new fields into your local config.yaml. -config_version: 14 +config_version: 15 # ============================================================================ # Logging @@ -32,6 +32,22 @@ log_level: info token_usage: enabled: true +# ============================================================================ +# Token Budget — Per-run token limits +# ============================================================================ +# Prevents runaway API costs by enforcing hard token limits per run. +# When warn_threshold is crossed, the agent receives an in-context warning. +# When hard_stop_threshold is crossed, tool_calls are stripped and the agent +# is forced to produce a final answer immediately. +token_budget: + enabled: false # Set to true to activate budget enforcement + max_tokens: 200000 # Total token limit (input + output) per run + max_input_tokens: null # Optional separate input-only limit + max_output_tokens: null # Optional separate output-only limit + warn_threshold: 0.8 # Warn at 80% of the budget + hard_stop_threshold: 1.0 # Force stop at 100% of the budget + + # ============================================================================ # Models Configuration # ============================================================================