diff --git a/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py index 88d46690b..0c867de81 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py @@ -21,15 +21,18 @@ Run scope: and those continuations share one budget; a later user run gets a new ``run_id`` and a fresh budget. Only the per-message ``seen`` map is dropped (``before_agent`` rebuilds it). Invocations without a non-empty string - ``run_id`` use runtime-local identity and clear their usage/warning state - in ``after_agent``. + ``run_id`` are keyed by LangGraph's run-scoped ``Runtime.control`` object + (each graph node gets its own ``Runtime`` wrapper, but they share it) and + clear their usage/warning state in ``after_agent``. Stop-reason surfacing (#3875 Phase 2): The hard stop does NOT raise — it strips tool_calls so the agent loop terminates naturally and produces a final answer. To let the caller (e.g. the subagent executor) distinguish a budget-capped completion from a clean one, the run that triggered the hard stop is recorded in ``_stop_reason`` - and exposed via :meth:`consume_stop_reason`. That dict is intentionally NOT + and exposed via :meth:`consume_stop_reason`. It is keyed by the context + ``run_id`` exactly as given, ``None`` included, because that is what the + executor passes back. That dict is intentionally NOT cleared by ``after_agent``/``_clear_run_state`` so the executor can read it after the run returns; the bounded dict prevents unbounded growth on abandoned runs, and each subagent run builds a fresh middleware instance so @@ -40,6 +43,7 @@ from __future__ import annotations import logging import threading +import uuid from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import Any, override @@ -84,7 +88,10 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]): # Stop reason set when the hard-stop fires. NOT cleared by # ``_clear_run_state``/``after_agent`` so the executor can consume it # after the run returns; bounded so abandoned runs cannot leak. - self._stop_reason: BoundedDict[str, str] = BoundedDict(1000) + self._stop_reason: BoundedDict[str | None, str] = BoundedDict(1000) + # id(Runtime.control) -> (control, generated key) for invocations + # without a context run_id; released in ``after_agent``. + self._fallback_run_ids: BoundedDict[int, tuple[object, str]] = BoundedDict(1000) def release_policy_parameters(self) -> dict[str, object]: return {"config": self._config.model_dump(mode="python")} @@ -100,6 +107,7 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]): self._seen_messages.clear() self._cumulative_usage.clear() self._stop_reason.clear() + self._fallback_run_ids.clear() def consume_stop_reason(self, run_id: str | None) -> str | None: """Pop and return the stop reason the hard-stop set for this run. @@ -120,10 +128,43 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]): run_id = ctx.get("run_id") if isinstance(ctx, dict) else None return run_id if isinstance(run_id, str) and run_id else None - @classmethod - def _get_run_id(cls, runtime: Runtime) -> str: - # Fallback to runtime object ID to prevent collisions across embedded client runs - return cls._context_run_id(runtime) or str(id(runtime)) + def _get_run_id(self, runtime: Runtime) -> str: + run_id = self._context_run_id(runtime) + if run_id is not None: + return run_id + # Same anchor as LoopDetectionMiddleware: ``id(runtime)`` changes from + # one graph node to the next, ``Runtime.control`` does not. The key is a + # generated token rather than the address, which can be reused once the + # object is collected. Unlike loop detection, ``execution_info.run_id`` + # is skipped on purpose: without a context run_id the budget is per + # invocation, not per RunnableConfig run. + control = getattr(runtime, "control", None) + anchor = control if control is not None else runtime + with self._lock: + entry = self._fallback_run_ids.get(id(anchor)) + if entry is None or entry[0] is not anchor: + entry = (anchor, f"__invocation__:{uuid.uuid4().hex}") + self._fallback_run_ids[id(anchor)] = entry + # Least recently used goes first, so a full map never evicts an active invocation. + self._fallback_run_ids.move_to_end(id(anchor)) + return entry[1] + + def _release_fallback_run_id(self, runtime: Runtime) -> None: + control = getattr(runtime, "control", None) + anchor = control if control is not None else runtime + with self._lock: + entry = self._fallback_run_ids.get(id(anchor)) + if entry is not None and entry[0] is anchor: + del self._fallback_run_ids[id(anchor)] + + @staticmethod + def _stop_reason_key(runtime: Runtime, run_id: str) -> str | None: + # SubagentExecutor consumes the stop reason with its raw run_id, which + # is None when the parent run has none. + ctx = getattr(runtime, "context", None) + if isinstance(ctx, dict) and "run_id" in ctx and (ctx["run_id"] is None or isinstance(ctx["run_id"], str)): + return ctx["run_id"] + return run_id def _clear_run_state(self, run_id: str) -> None: with self._lock: @@ -171,6 +212,7 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]): self._seen_messages.pop(run_id, None) return self._clear_run_state(run_id) + self._release_fallback_run_id(runtime) @override async def aafter_agent(self, state: AgentState, runtime: Runtime) -> None: @@ -274,7 +316,7 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]): # ``stop_reason=token_capped`` to the lead after the run # returns (the hard stop itself does not raise). See # ``consume_stop_reason``. - self._stop_reason[run_id] = "token_capped" + self._stop_reason[self._stop_reason_key(runtime, run_id)] = "token_capped" # Also write to runtime.context so the lead worker can read it # without needing a reference to this middleware instance (#4176). ctx = getattr(runtime, "context", None) diff --git a/backend/tests/test_token_budget_middleware.py b/backend/tests/test_token_budget_middleware.py index dbba48917..b302c4d16 100644 --- a/backend/tests/test_token_budget_middleware.py +++ b/backend/tests/test_token_budget_middleware.py @@ -1,3 +1,4 @@ +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -102,8 +103,8 @@ class TestTokenBudgetLifecycle: mw.after_model(state, runtime) # Missing identities are invocation-local, never shared under None or "". - key = str(id(runtime)) - assert mw._get_run_id(runtime) == key + key = mw._get_run_id(runtime) + assert key.startswith("__invocation__:") assert mw._cumulative_usage[key].total == 850 assert mw._warned[key] assert mw._pending_warnings[key] @@ -121,8 +122,20 @@ class TestTokenBudgetLifecycle: follow_up = _make_state_with_usage(total=200) follow_up["messages"][0].id = "next-msg" assert mw.after_model(follow_up, runtime) is None - assert mw._cumulative_usage[key].total == 200 - assert not mw._warned.get(key) + next_key = mw._get_run_id(runtime) + assert next_key != key + assert mw._cumulative_usage[next_key].total == 200 + assert not mw._warned.get(next_key) + + def test_active_invocation_keeps_its_key_when_the_anchor_map_is_full(self): + mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000)) + mw._fallback_run_ids.maxsize = 3 + active = SimpleNamespace(context={}, control=object()) + key = mw._get_run_id(active) + + for _ in range(5): + mw._get_run_id(SimpleNamespace(context={}, control=object())) + assert mw._get_run_id(active) == key @pytest.mark.asyncio async def test_valid_run_id_preserves_usage_warnings_and_stop_reason(self): @@ -226,6 +239,17 @@ class TestTokenBudgetHardStop: # A run that never hit the cap has no stop reason. assert mw.consume_stop_reason("uncapped-run") is None + def test_stop_reason_round_trips_an_explicit_none_run_id(self): + """A subagent whose parent run has no run_id runs with ``run_id=None``; + ``SubagentExecutor`` reads the reason back with that same ``None``.""" + mw = TokenBudgetMiddleware.from_config(TokenBudgetConfig(max_tokens=1000, enabled=True)) + runtime = _make_runtime(run_id=None) + tool_calls = [{"name": "bash", "args": {"command": "ls"}, "id": "call_1"}] + assert mw._apply(_make_state_with_usage(total=1500, tool_calls=tool_calls), runtime) is not None + + assert mw.consume_stop_reason(None) == "token_capped" + assert mw.consume_stop_reason(None) is None + def test_below_threshold_does_not_stamp_stop_reason(self): """A run that only crosses the warn threshold (not the hard stop) keeps running and must not stamp ``token_capped`` — the run is not capped.""" @@ -270,6 +294,16 @@ class _ToolCallingFakeModel(FakeMessagesListChatModel): return self +class _RecordingToolCallingFakeModel(_ToolCallingFakeModel): + def __init__(self, **kwargs): + super().__init__(**kwargs) + object.__setattr__(self, "requests", []) + + def _generate(self, messages, stop=None, run_manager=None, **kwargs): + self.requests.append(list(messages)) + return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs) + + class TestTokenBudgetAgentGraph: def test_goal_continuation_shares_the_run_budget(self): """A hidden goal continuation re-enters the graph under the same run_id; it must not get a fresh budget.""" @@ -315,3 +349,41 @@ class TestTokenBudgetAgentGraph: # A later user run still starts with a fresh budget. graph.invoke({"messages": [HumanMessage("next question")]}, config=config, context={"thread_id": "goal-thread", "run_id": "run-2"}) assert executed == ["a", "b", "d"] + + @pytest.mark.parametrize("context", [{"thread_id": "no-run-id"}, {"thread_id": "no-run-id", "run_id": None}]) + def test_invocation_without_run_id_keeps_one_budget_across_graph_nodes(self, context): + """LangGraph hands each node its own Runtime, so an invocation without a run_id can't be keyed by id(runtime).""" + executed: list[str] = [] + + @as_tool + def bash(command: str) -> str: + """Run a fake shell command.""" + executed.append(command) + return "ok" + + def call(command: str, tokens: int) -> AIMessage: + return AIMessage( + content="", + id=f"ai-{command}", + tool_calls=[{"name": "bash", "id": f"call-{command}", "args": {"command": command}}], + usage_metadata={"input_tokens": tokens, "output_tokens": 0, "total_tokens": tokens}, + ) + + def answer(text: str) -> AIMessage: + return AIMessage(content=text, id=f"ai-{text}", usage_metadata={"input_tokens": 500, "output_tokens": 0, "total_tokens": 500}) + + model = _RecordingToolCallingFakeModel(responses=[call("a", 4000), call("b", 4500), answer("first answer"), call("c", 4000), call("d", 4500), answer("second answer")]) + mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000, warn_threshold=0.8)) + graph = create_agent(model=model, tools=[bash], middleware=[mw], checkpointer=InMemorySaver()) + config = {"configurable": {"thread_id": "no-run-id"}} + + # 8.5k of 10k after "b": the warning reaches the next model request. + graph.invoke({"messages": [HumanMessage("research")]}, config=config, context=dict(context)) + assert [getattr(message, "name", None) for message in model.requests[2]][-1] == "budget_warning" + + # The next invocation has its own 10k; the first one's 9k doesn't count. + result = graph.invoke({"messages": [HumanMessage("next question")]}, config=config, context=dict(context)) + assert executed == ["a", "b", "c", "d"] + assert result["messages"][-1].content == "second answer" + for values in (mw._cumulative_usage, mw._warned, mw._pending_warnings, mw._seen_messages, mw._fallback_run_ids): + assert not values