github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-19 19:16:17 +00:00
Code Issues Packages Projects Releases Wiki Activity
deer-flow/backend/tests/test_token_budget_middleware.py
alanhuangyoo 0f2813f6b9
fix(agents): keep queued guard warnings when a model call is retried (#5433) 
* fix(agents): keep queued guard warnings when a model call is retried

LoopDetectionMiddleware, TokenBudgetMiddleware and ToolProgressMiddleware
pop their queued warning/hint before calling the model. When the call
raises, LLMErrorHandlingMiddleware (outside them) retries by running
their wrap_model_call again, and by then the queue is empty, so the
retried request goes out without the warning. Loop detection and the
token budget have already marked it as sent, so it is never queued
again, and a loop runs on to the hard stop unwarned.

Put the drained items back in front of the queue when the handler
raises, then re-raise.

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

* fix(agents): trim restored loop warnings from the tail and drop a dead helper

_restore_pending_warnings put the restored warnings at the front and then
trimmed the front, so if the cap ever fired it would drop exactly what it
restored. Trim the tail, as tool progress does. _augment_request had no
callers after the wrap_model_call change. Add the sync twin of the tool
progress retry test.

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

* refactor(agents): drop tool progress's unused _augment_request

Its only remaining reference was a test name; the dedup that test checks lives
in _inject_hints.

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

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
2026-09-16 15:52:17 +08:00

466 lines
22 KiB
Python
Raw Blame History

from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from langchain.agents import create_agent
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langchain_core.tools import tool as as_tool
from langgraph.checkpoint.memory import InMemorySaver
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
from deerflow.config.token_budget_config import TokenBudgetConfig
from deerflow.models.claude_provider import ClaudeChatModel
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 TestTokenBudgetLifecycle:
@pytest.mark.parametrize("context", [None, {}, {"run_id": None}, {"run_id": ""}, {"run_id": 0}, {"run_id": []}])
@pytest.mark.parametrize("async_hooks", [False, True])
@pytest.mark.asyncio
async def test_missing_or_invalid_run_id_clears_invocation_state(self, context, async_hooks):
mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
runtime = _make_runtime()
runtime.context = context
state = _make_state_with_usage(total=850)
if async_hooks:
await mw.abefore_agent({"messages": []}, runtime)
await mw.aafter_model(state, runtime)
else:
mw.before_agent({"messages": []}, runtime)
mw.after_model(state, runtime)
# Missing identities are invocation-local, never shared under None or "".
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]
assert mw._seen_messages[key]
if async_hooks:
await mw.aafter_agent(state, runtime)
else:
mw.after_agent(state, runtime)
for values in (mw._cumulative_usage, mw._warned, mw._pending_warnings, mw._seen_messages):
assert key not in values
# Reusing even the same runtime object starts a fresh invocation budget.
mw.before_agent(state, runtime)
follow_up = _make_state_with_usage(total=200)
follow_up["messages"][0].id = "next-msg"
assert mw.after_model(follow_up, runtime) is None
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):
mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000))
runtime = _make_runtime(run_id="goal-run")
mw.after_model(_make_state_with_usage(total=850), runtime)
state = _make_state_with_usage(total=1100)
assert mw.after_model(state, runtime) is not None
await mw.aafter_agent(state, runtime)
assert "goal-run" not in mw._seen_messages
assert mw._cumulative_usage["goal-run"].total == 1100
assert mw._warned["goal-run"]
assert len(mw._pending_warnings["goal-run"]) == 1
assert mw.consume_stop_reason("goal-run") == "token_capped"
# Continuations may get another Runtime object with the same run identity.
continuation = _make_runtime(run_id="goal-run")
await mw.abefore_agent(state, continuation)
await mw.aafter_model(state, continuation)
assert mw._cumulative_usage["goal-run"].total == 1100
assert len(mw._pending_warnings["goal-run"]) == 1
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_warning_survives_a_failed_model_call(self):
"""A call that raises is retried by LLMErrorHandlingMiddleware through this wrap; the warning must still be sent."""
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)
request = _make_request([AIMessage(content="hi")], runtime)
sent = []
def flaky_handler(req):
sent.append(req.messages)
if len(sent) == 1:
raise RuntimeError("503 Service Unavailable")
return MagicMock()
with pytest.raises(RuntimeError):
mw.wrap_model_call(request, flaky_handler)
mw.wrap_model_call(request, flaky_handler)
assert [any(getattr(message, "name", None) == "budget_warning" for message in messages) for messages in sent] == [True, True]
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
def test_hard_stop_drops_provider_tool_call_content_blocks(self):
# Anthropic keeps tool_use blocks in content; one left behind without a
# tool_result makes every later request on the thread fail with a 400.
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": "toolu_1"}]
content = [
{"type": "text", "text": "Listing"},
{"type": "tool_use", "id": "toolu_1", "name": "bash", "input": {"command": "ls"}},
]
res = mw._apply(_make_state_with_usage(total=105000, tool_calls=tool_calls, content=content), _make_runtime())
stopped = res["messages"][0]
assert [block["type"] for block in stopped.content] == ["text", "text"]
assert "TOKEN BUDGET EXCEEDED" in stopped.content[-1]["text"]
def test_hard_stop_stamps_token_capped_stop_reason_consumed_once(self):
"""#3875 Phase 2: a hard-stop stamps ``token_capped`` on a per-run
accessor the executor reads post-run. It pops on read so a second read
(e.g. a retry over the same executor) does not double-report, and a
non-capped run yields ``None``."""
config = TokenBudgetConfig(max_tokens=100000, hard_stop_threshold=1.0, enabled=True)
mw = TokenBudgetMiddleware.from_config(config)
runtime = _make_runtime(run_id="capped-run")
tool_calls = [{"name": "bash", "args": {"command": "ls"}, "id": "call_1"}]
state = _make_state_with_usage(total=105000, tool_calls=tool_calls, content="partial answer")
mw._apply(state, runtime)
# First read pops the reason.
assert mw.consume_stop_reason("capped-run") == "token_capped"
# Second read is None — the reason is per-run and consumed once.
assert mw.consume_stop_reason("capped-run") is None
# 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."""
config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.7, hard_stop_threshold=1.0, enabled=True)
mw = TokenBudgetMiddleware.from_config(config)
runtime = _make_runtime(run_id="warn-run")
# 80k of 100k -> crosses warn (0.7) but not hard stop (1.0).
state = _make_state_with_usage(total=80000)
mw._apply(state, runtime)
assert mw.consume_stop_reason("warn-run") is None
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
class _ToolCallingFakeModel(FakeMessagesListChatModel):
def bind_tools(self, tools, *, tool_choice=None, **kwargs):
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."""
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 = 4000) -> 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},
)
model = _ToolCallingFakeModel(
responses=[
call("a"),
call("b"),
AIMessage(content="first answer", id="ai-answer-1", usage_metadata={"input_tokens": 1000, "output_tokens": 0, "total_tokens": 1000}),
call("c"),
call("d"),
AIMessage(content="second answer", id="ai-answer-2", usage_metadata={"input_tokens": 1000, "output_tokens": 0, "total_tokens": 1000}),
]
)
mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
graph = create_agent(model=model, tools=[bash], middleware=[mw], checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "goal-thread"}}
# User turn: 9k of 10k.
graph.invoke({"messages": [HumanMessage("research")]}, config=config, context={"thread_id": "goal-thread", "run_id": "run-1"})
assert executed == ["a", "b"]
# Goal continuation in the same run: the next 4k call crosses the cap.
result = graph.invoke({"messages": [HumanMessage("keep going")]}, config=config, context={"thread_id": "goal-thread", "run_id": "run-1"})
assert executed == ["a", "b"]
assert "TOKEN BUDGET EXCEEDED" in result["messages"][-1].content
# 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"]
def test_checkpointed_hard_stop_leaves_next_anthropic_turn_well_formed(self):
"""The stopped message is checkpointed and replayed; its tool_use must not reach the next request unpaired."""
executed: list[str] = []
@as_tool
def bash(command: str) -> str:
"""Run a fake shell command."""
executed.append(command)
return "ok"
over_budget_call = AIMessage(
content=[
{"type": "text", "text": "Listing files."},
{"type": "tool_use", "id": "toolu_ls", "name": "bash", "input": {"command": "ls"}},
],
id="ai-over-budget",
tool_calls=[{"name": "bash", "id": "toolu_ls", "args": {"command": "ls"}}],
usage_metadata={"input_tokens": 12_000, "output_tokens": 0, "total_tokens": 12_000},
)
answer = AIMessage(content="second answer", id="ai-answer", usage_metadata={"input_tokens": 100, "output_tokens": 0, "total_tokens": 100})
model = _RecordingToolCallingFakeModel(responses=[over_budget_call, answer])
mw = TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=10_000))
graph = create_agent(model=model, tools=[bash], middleware=[mw], checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "stopped-thread"}}
graph.invoke({"messages": [HumanMessage("list files")]}, config=config, context={"thread_id": "stopped-thread", "run_id": "run-1"})
graph.invoke({"messages": [HumanMessage("continue")]}, config=config, context={"thread_id": "stopped-thread", "run_id": "run-2"})
assert executed == []
payload = ClaudeChatModel(model="claude-sonnet-4-5", anthropic_api_key="sk-ant-offline")._get_request_payload(model.requests[1])
turns = [turn["content"] if isinstance(turn["content"], list) else [] for turn in payload["messages"]]
for index, blocks in enumerate(turns):
tool_use_ids = {block["id"] for block in blocks if block["type"] == "tool_use"}
following = turns[index + 1] if index + 1 < len(turns) else []
assert tool_use_ids <= {block["tool_use_id"] for block in following if block["type"] == "tool_result"}
@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
Reference in New Issue View Git Blame Copy Permalink