github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-15 00:19:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(agents): keep token budget signals for runs without a run_id (#5436)

Browse Source 
* fix(agents): keep token budget signals for runs without a run_id

#5410 moved every invocation without a non-empty string run_id onto
str(id(runtime)). Two things break on that key:

- SubagentExecutor passes the parent's run_id, None when the parent run has
  none (LangGraph Server, direct create_deerflow_agent callers), and reads the
  stop reason back with that None. The hard stop stored it under the id string,
  so a token-capped subagent reported a clean completion to the lead.
- LangGraph gives each graph node its own Runtime wrapper, so the key changed
  between after_model and the next model call: the budget warning was never
  delivered, and each after_model counted every AIMessage in the thread.

Key those invocations by Runtime.control, as LoopDetectionMiddleware does,
release it in after_agent, and store the stop reason under the context run_id
as given.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* fix(agents): keep an active invocation's budget key when the anchor map is full

The fallback anchor map was FIFO, so with 1000 run_id-less invocations on a
shared instance an active one could lose its anchor mid-run and restart with a
fresh budget. Move the anchor to the end on every lookup, as loop detection
does, and note why execution_info.run_id is not consulted.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
alanhuangyoo 2026-09-15 06:56:38 +08:00 committed by GitHub
parent 6469833886
commit 99902b7791
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 127 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
backend/packages/harness/deerflow/agents/middlewares/token_budget_middleware.py
View File

@ -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)

80
backend/tests/test_token_budget_middleware.py
View File

@ -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