deer-flow/backend/tests/test_todo_middleware.py
Evan d8e085630a
fix(agents): keep todo reminders in Goal turns after an earlier provider stop (#8004)
* fix(agents): keep todo reminders in Goal turns after an earlier provider stop

TodoMiddleware skipped completion reminders whenever runtime.context held
safety_capped or model_length_capped. The worker keeps that key across the
run's hidden Goal continuation turns so the run record reports the first
cap, and an unmet Goal continues after these provider stops (it stands down
only on loop and token caps). After one capped response, every continuation
turn therefore lost its todo reminders and could stop with unfinished items.

Loop and token stops stay run-level and are still read from the context.
Provider safety and length stops now apply to the response they ended,
identified by the safety_termination / model_length_termination marker that
those guards stamp on it. The check still runs before the parent hook, so a
capped response cannot be resumed by the parallel-write retry.

* refactor(agents): import provider-stop markers from the guards that stamp them

TodoMiddleware matched the safety and length markers by string literals that
duplicated the guards' own. Each guard now exports its marker and stop reason
as module constants, used at its stamp and check sites and imported by
TodoMiddleware, so a rename in a guard cannot silently disable suppression.

Comment that completion_reminder_hard_stop_reasons is the inventory surface
and only its run-level entries are read from runtime.context, and run the
worker-level Goal continuation test for the safety guard too, which
overwrites stop_reason instead of keeping the first one.

---------

Co-authored-by: Totoro <msy626836554@gmail.com>
2026-11-08 19:19:56 +08:00

1285 lines
54 KiB
Python

"""Tests for TodoMiddleware context-loss detection."""
import ast
import asyncio
from pathlib import Path
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
import pytest
from langchain.agents import create_agent
from langchain.agents.middleware import TodoListMiddleware
from langchain.agents.middleware.types import ModelRequest
from langchain.tools import tool
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import PrivateAttr
from deerflow.agents.middlewares.model_length_finish_reason_middleware import (
ModelLengthFinishReasonMiddleware,
)
from deerflow.agents.middlewares.safety_finish_reason_middleware import SafetyFinishReasonMiddleware
from deerflow.agents.middlewares.todo_middleware import (
TODO_REMINDER_MESSAGE_NAME,
TodoMiddleware,
_format_todos,
_has_tool_call_intent_or_error,
_reminder_in_messages,
_todos_in_messages,
)
from deerflow.agents.thread_state import ThreadState
def _ai_with_write_todos():
return AIMessage(content="", tool_calls=[{"name": "write_todos", "id": "tc_1", "args": {}}])
def _reminder_msg():
return HumanMessage(name=TODO_REMINDER_MESSAGE_NAME, content="reminder")
class _CapturingFakeMessagesListChatModel(FakeMessagesListChatModel):
_seen_messages: list[list[Any]] = PrivateAttr(default_factory=list)
@property
def seen_messages(self) -> list[list[Any]]:
return self._seen_messages
def bind_tools(self, tools, *, tool_choice=None, **kwargs):
return self
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
self._seen_messages.append(list(messages))
return super()._generate(
messages,
stop=stop,
run_manager=run_manager,
**kwargs,
)
def _make_runtime():
runtime = MagicMock()
runtime.context = {"thread_id": "test-thread", "run_id": "test-run"}
return runtime
def _make_runtime_for(thread_id: str, run_id: str):
runtime = _make_runtime()
runtime.context = {"thread_id": thread_id, "run_id": run_id}
return runtime
def _make_model_request(messages: list[Any], *, runtime=None) -> ModelRequest:
return ModelRequest(
model=object(),
messages=list(messages),
state={"messages": list(messages)},
runtime=runtime,
)
def _sample_todos():
return [
{"status": "completed", "content": "Set up project"},
{"status": "in_progress", "content": "Write tests"},
{"status": "pending", "content": "Deploy"},
]
class TestTodosInMessages:
def test_true_when_write_todos_present(self):
msgs = [HumanMessage(content="hi"), _ai_with_write_todos()]
assert _todos_in_messages(msgs) is True
def test_false_when_no_write_todos(self):
msgs = [
HumanMessage(content="hi"),
AIMessage(content="hello", tool_calls=[{"name": "bash", "id": "tc_1", "args": {}}]),
]
assert _todos_in_messages(msgs) is False
def test_false_for_empty_list(self):
assert _todos_in_messages([]) is False
def test_false_for_ai_without_tool_calls(self):
msgs = [AIMessage(content="hello")]
assert _todos_in_messages(msgs) is False
class TestReminderInMessages:
def test_true_when_reminder_present(self):
msgs = [HumanMessage(content="hi"), _reminder_msg()]
assert _reminder_in_messages(msgs) is True
def test_false_when_no_reminder(self):
msgs = [HumanMessage(content="hi"), AIMessage(content="hello")]
assert _reminder_in_messages(msgs) is False
def test_false_for_empty_list(self):
assert _reminder_in_messages([]) is False
def test_false_for_human_without_name(self):
msgs = [HumanMessage(content="todo_reminder")]
assert _reminder_in_messages(msgs) is False
class TestFormatTodos:
def test_formats_multiple_items(self):
todos = _sample_todos()
result = _format_todos(todos)
assert "- [completed] Set up project" in result
assert "- [in_progress] Write tests" in result
assert "- [pending] Deploy" in result
def test_empty_list(self):
assert _format_todos([]) == ""
def test_missing_fields_use_defaults(self):
todos = [{"content": "No status"}, {"status": "done"}]
result = _format_todos(todos)
assert "- [pending] No status" in result
assert "- [done] " in result
class TestBeforeModel:
def test_returns_none_when_no_todos(self):
mw = TodoMiddleware()
state = {"messages": [HumanMessage(content="hi")], "todos": []}
assert mw.before_model(state, _make_runtime()) is None
def test_returns_none_when_todos_is_none(self):
mw = TodoMiddleware()
state = {"messages": [HumanMessage(content="hi")], "todos": None}
assert mw.before_model(state, _make_runtime()) is None
def test_returns_none_when_write_todos_still_visible(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_with_write_todos()],
"todos": _sample_todos(),
}
assert mw.before_model(state, _make_runtime()) is None
def test_returns_none_when_reminder_already_present(self):
mw = TodoMiddleware()
state = {
"messages": [HumanMessage(content="hi"), _reminder_msg()],
"todos": _sample_todos(),
}
assert mw.before_model(state, _make_runtime()) is None
def test_injects_reminder_when_todos_exist_but_truncated(self):
mw = TodoMiddleware()
state = {
"messages": [HumanMessage(content="hi"), AIMessage(content="sure")],
"todos": _sample_todos(),
}
result = mw.before_model(state, _make_runtime())
assert result is not None
msgs = result["messages"]
assert len(msgs) == 1
assert isinstance(msgs[0], HumanMessage)
assert msgs[0].name == TODO_REMINDER_MESSAGE_NAME
def test_reminder_contains_formatted_todos(self):
mw = TodoMiddleware()
state = {
"messages": [HumanMessage(content="hi")],
"todos": _sample_todos(),
}
result = mw.before_model(state, _make_runtime())
content = result["messages"][0].content
assert "Set up project" in content
assert "Write tests" in content
assert "Deploy" in content
assert "system_reminder" in content
class TestAbeforeModel:
def test_delegates_to_sync(self):
mw = TodoMiddleware()
state = {
"messages": [HumanMessage(content="hi")],
"todos": _sample_todos(),
}
result = asyncio.run(mw.abefore_model(state, _make_runtime()))
assert result is not None
assert result["messages"][0].name == TODO_REMINDER_MESSAGE_NAME
def _todo_completion_reminders(messages):
reminders = []
for message in messages:
if isinstance(message, HumanMessage) and message.name == "todo_completion_reminder":
reminders.append(message)
return reminders
def _ai_no_tool_calls():
return AIMessage(content="I'm done!")
def _ai_stopped_by(stop_reason):
"""The response a guard leaves behind: provider stops also mark the message."""
marker = {"safety_capped": "safety_termination", "model_length_capped": "model_length_termination"}.get(stop_reason)
return AIMessage(content="I'm done!", additional_kwargs={marker: {"suppressed_tool_call_count": 0}} if marker else {})
def _ai_with_invalid_tool_calls():
return AIMessage(
content="",
tool_calls=[],
invalid_tool_calls=[
{
"type": "invalid_tool_call",
"id": "write_file:36",
"name": "write_file",
"args": "{invalid",
"error": "Failed to parse tool arguments",
}
],
)
def _ai_with_raw_provider_tool_calls():
return AIMessage(
content="",
tool_calls=[],
invalid_tool_calls=[],
additional_kwargs={
"tool_calls": [
{
"id": "raw-tool-call",
"type": "function",
"function": {"name": "write_file", "arguments": '{"path":"report.md"}'},
}
]
},
)
def _ai_with_legacy_function_call():
return AIMessage(
content="",
additional_kwargs={"function_call": {"name": "write_file", "arguments": '{"path":"report.md"}'}},
)
def _ai_with_tool_finish_reason():
return AIMessage(content="", response_metadata={"finish_reason": "tool_calls"})
def _incomplete_todos():
return [
{"status": "completed", "content": "Step 1"},
{"status": "in_progress", "content": "Step 2"},
{"status": "pending", "content": "Step 3"},
]
def _all_completed_todos():
return [
{"status": "completed", "content": "Step 1"},
{"status": "completed", "content": "Step 2"},
]
class TestToolCallIntentOrError:
def test_false_for_plain_final_answer(self):
assert _has_tool_call_intent_or_error(_ai_no_tool_calls()) is False
def test_true_for_structured_tool_calls(self):
assert _has_tool_call_intent_or_error(_ai_with_write_todos()) is True
def test_true_for_invalid_tool_calls(self):
assert _has_tool_call_intent_or_error(_ai_with_invalid_tool_calls()) is True
def test_true_for_raw_provider_tool_calls(self):
assert _has_tool_call_intent_or_error(_ai_with_raw_provider_tool_calls()) is True
def test_true_for_legacy_function_call(self):
assert _has_tool_call_intent_or_error(_ai_with_legacy_function_call()) is True
def test_true_for_tool_finish_reason(self):
assert _has_tool_call_intent_or_error(_ai_with_tool_finish_reason()) is True
def test_langchain_ai_message_tool_fields_are_explicitly_handled(self):
# Sentinel for LangChain compatibility: if future AIMessage versions add
# new top-level tool/function-call fields, this test should fail. When
# it does, update `_has_tool_call_intent_or_error()` so the completion
# reminder guard explicitly decides whether each new field means "not a
# clean final answer"; the helper has a matching comment pointing back
# to this sentinel.
tool_related_fields = {name for name in AIMessage.model_fields if "tool" in name.lower() or ("function" in name.lower() and "call" in name.lower())}
assert tool_related_fields <= {"tool_calls", "invalid_tool_calls"}
class TestAfterModel:
def test_all_middleware_stop_reason_writes_have_an_explicit_todo_policy(self):
middleware_dir = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/agents/middlewares"
stamped_reasons = set()
for source in sorted(middleware_dir.glob("*.py")):
tree = ast.parse(source.read_text(encoding="utf-8"))
constants = {target.id: node.value for node in tree.body if isinstance(node, ast.Assign) for target in node.targets if isinstance(target, ast.Name) and isinstance(node.value, ast.Constant)}
for node in ast.walk(tree):
if not isinstance(node, (ast.Assign, ast.AnnAssign)):
continue
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
for target in targets:
if not isinstance(target, ast.Subscript) or not isinstance(target.slice, ast.Constant) or target.slice.value != "stop_reason":
continue
# This copies a child result into a ledger entry; it does
# not stamp the current run's runtime.context.
if source.name == "delegation_ledger.py" and isinstance(target.value, ast.Name) and target.value.id == "entry":
continue
value = constants.get(node.value.id) if isinstance(node.value, ast.Name) else node.value
assert isinstance(value, ast.Constant) and isinstance(value.value, str), f"Classify stop_reason writer at {source.name}:{node.lineno}"
stamped_reasons.add(value.value)
# The subagent budget only removes task calls; the lead may still
# complete its work using other tools. Every other writer is terminal.
nonterminal_reasons = {"subagent_limit_capped"}
assert nonterminal_reasons <= stamped_reasons
policy = TodoMiddleware().release_policy_parameters()
assert set(policy["completion_reminder_hard_stop_reasons"]) == stamped_reasons - nonterminal_reasons
@pytest.mark.parametrize("marker", ["safety_termination", "model_length_termination"])
@pytest.mark.parametrize("context", [None, {}, {"stop_reason": "custom_stop"}])
def test_termination_marker_suppresses_reminders_without_recognized_context(self, marker, context):
mw = TodoMiddleware()
runtime = _make_runtime()
runtime.context = context
state = {"messages": [AIMessage(content="partial response", additional_kwargs={marker: {"suppressed_tool_call_count": 0}})], "todos": _incomplete_todos()}
assert mw.after_model(state, runtime) is None
assert mw._pending_completion_reminders == {}
@pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("stop_reason", ["loop_capped", "token_capped", "safety_capped", "model_length_capped"])
def test_recognized_stop_reason_suppresses_completion_reminders(self, invoke_async, stop_reason):
mw = TodoMiddleware()
runtime = _make_runtime()
runtime.context["stop_reason"] = stop_reason
state = {"messages": [_ai_stopped_by(stop_reason)], "todos": _incomplete_todos()}
result = asyncio.run(mw.aafter_model(state, runtime)) if invoke_async else mw.after_model(state, runtime)
assert result is None
assert mw._pending_completion_reminders == {}
assert mw._completion_reminder_counts == {}
assert state["todos"] == _incomplete_todos()
@pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("stop_reason", ["loop_capped", "token_capped", "safety_capped", "model_length_capped"])
def test_hard_stop_takes_precedence_over_base_jump(self, monkeypatch, invoke_async, stop_reason):
base_hook = MagicMock(return_value={"jump_to": "model"})
monkeypatch.setattr(TodoListMiddleware, "after_model", base_hook)
mw = TodoMiddleware()
runtime = _make_runtime()
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
# Preserve normal parent behavior, but never let it override a hard stop.
def after_model():
return asyncio.run(mw.aafter_model(state, runtime)) if invoke_async else mw.after_model(state, runtime)
assert after_model() == {"jump_to": "model"}
base_hook.assert_called_once()
base_hook.reset_mock()
runtime.context["stop_reason"] = stop_reason
state["messages"] = [_ai_stopped_by(stop_reason)]
assert after_model() is None
base_hook.assert_not_called()
@pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("stop_reason", ["safety_capped", "model_length_capped"])
def test_earlier_provider_stop_in_context_keeps_completion_reminders(self, invoke_async, stop_reason):
# Goal continuation turns share the run's context, which keeps the first
# cap for the run record. A complete later response is not capped.
mw = TodoMiddleware()
runtime = _make_runtime()
runtime.context["stop_reason"] = stop_reason
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
result = asyncio.run(mw.aafter_model(state, runtime)) if invoke_async else mw.after_model(state, runtime)
assert result == {"jump_to": "model"}
assert mw._completion_reminder_count_for_runtime(runtime) == 1
@pytest.mark.parametrize("stop_reason", [None, "custom_stop", "subagent_limit_capped", [], {}])
def test_unrecognized_stop_reason_keeps_completion_reminders(self, stop_reason):
mw = TodoMiddleware()
runtime = _make_runtime()
runtime.context["stop_reason"] = stop_reason
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
assert mw.after_model(state, runtime) == {"jump_to": "model"}
def test_returns_none_when_agent_still_using_tools(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_with_write_todos()],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_returns_none_when_no_todos(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_no_tool_calls()],
"todos": [],
}
assert mw.after_model(state, _make_runtime()) is None
def test_returns_none_when_todos_is_none(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_no_tool_calls()],
"todos": None,
}
assert mw.after_model(state, _make_runtime()) is None
def test_returns_none_when_all_completed(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_no_tool_calls()],
"todos": _all_completed_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_returns_none_when_no_messages(self):
mw = TodoMiddleware()
state = {
"messages": [],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_queues_reminder_and_jumps_to_model_when_incomplete(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [HumanMessage(content="hi"), _ai_no_tool_calls()],
"todos": _incomplete_todos(),
}
result = mw.after_model(state, runtime)
assert result is not None
assert result["jump_to"] == "model"
assert "messages" not in result
request = _make_model_request(state["messages"], runtime=runtime)
seen: list[ModelRequest] = []
def handler(model_request: ModelRequest):
seen.append(model_request)
return "response"
assert mw.wrap_model_call(request, handler) == "response"
assert len(seen) == 1
sent = seen[0]
assert sent.system_message is not None
assert "write_todos" in sent.system_message.text
reminder = sent.messages[-1]
assert isinstance(reminder, HumanMessage)
assert reminder.name == "todo_completion_reminder"
assert reminder.additional_kwargs["hide_from_ui"] is True
assert "Step 2" in reminder.content
assert "Step 3" in reminder.content
def test_reminder_lists_only_incomplete_items(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [_ai_no_tool_calls()],
"todos": _incomplete_todos(),
}
result = mw.after_model(state, runtime)
assert result is not None
request = _make_model_request(state["messages"], runtime=runtime)
seen: list[ModelRequest] = []
def handler(model_request: ModelRequest):
seen.append(model_request)
return "response"
mw.wrap_model_call(request, handler)
content = seen[0].messages[-1].content
assert "Step 1" not in content # completed — should not appear
assert "Step 2" in content
assert "Step 3" in content
def test_allows_exit_after_max_reminders(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [
_ai_no_tool_calls(),
],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, runtime) is not None
assert mw.after_model(state, runtime) is not None
assert mw.after_model(state, runtime) is None
def test_still_sends_reminder_before_cap(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [
_ai_no_tool_calls(),
],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, runtime) is not None
result = mw.after_model(state, runtime)
assert result is not None
assert result["jump_to"] == "model"
def test_does_not_trigger_for_invalid_tool_calls(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_with_invalid_tool_calls()],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_does_not_trigger_for_raw_provider_tool_calls(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_with_raw_provider_tool_calls()],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_does_not_trigger_for_legacy_function_call(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_with_legacy_function_call()],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_does_not_trigger_for_tool_finish_reason(self):
mw = TodoMiddleware()
state = {
"messages": [_ai_with_tool_finish_reason()],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_does_not_reengage_when_model_length_capped_marker_present(self):
mw = TodoMiddleware()
state = {
"messages": [
AIMessage(
content="nit",
tool_calls=[],
additional_kwargs={
"model_length_termination": {
"detector": "openai_compatible_length",
"suppressed_tool_call_count": 1,
"suppressed_tool_call_names": ["write_file"],
}
},
)
],
"todos": _incomplete_todos(),
}
assert mw.after_model(state, _make_runtime()) is None
def test_pure_text_length_cap_without_marker_still_reengages(self):
mw = TodoMiddleware()
state = {
"messages": [
AIMessage(
content="partial answer",
response_metadata={"finish_reason": "length"},
)
],
"todos": _incomplete_todos(),
}
result = mw.after_model(state, _make_runtime())
assert result is not None
assert result["jump_to"] == "model"
class TestAafterModel:
def test_delegates_to_sync(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [_ai_no_tool_calls()],
"todos": _incomplete_todos(),
}
result = asyncio.run(mw.aafter_model(state, runtime))
assert result is not None
assert result["jump_to"] == "model"
assert "messages" not in result
class TestWrapModelCall:
def test_no_pending_reminder_still_injects_todo_system_prompt(self):
"""Regression for bytedance/deer-flow#4714: the base class system-prompt
injection must survive the override, otherwise the model never learns
about `write_todos` and no todo list is ever produced."""
mw = TodoMiddleware()
runtime = _make_runtime()
request = _make_model_request([HumanMessage(content="hi")], runtime=runtime)
seen: list[ModelRequest] = []
def handler(model_request: ModelRequest):
seen.append(model_request)
return "response"
assert mw.wrap_model_call(request, handler) == "response"
assert len(seen) == 1
sent = seen[0]
assert sent.system_message is not None
assert "write_todos" in sent.system_message.text
# No pending reminder — messages must not be augmented.
assert sent.messages == [HumanMessage(content="hi")]
def test_pending_reminder_is_injected_once(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [_ai_no_tool_calls()],
"todos": _incomplete_todos(),
}
mw.after_model(state, runtime)
request = _make_model_request(state["messages"], runtime=runtime)
seen: list[ModelRequest] = []
def handler(model_request: ModelRequest):
seen.append(model_request)
return "response"
assert mw.wrap_model_call(request, handler) == "response"
assert len(seen) == 1
sent = seen[0]
assert sent.system_message is not None
assert "write_todos" in sent.system_message.text
injected = sent.messages[-1]
assert isinstance(injected, HumanMessage)
assert injected.name == "todo_completion_reminder"
# Second call: the reminder was drained, so only the system prompt is
# injected and the original messages pass through untouched.
seen.clear()
assert mw.wrap_model_call(request, handler) == "response"
assert len(seen) == 1
assert seen[0].messages == state["messages"]
def test_pending_reminder_survives_a_failed_model_call(self):
"""A call that raises is retried by LLMErrorHandlingMiddleware through this wrap; the reminder must still be sent."""
mw = TodoMiddleware()
runtime = _make_runtime()
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
mw.after_model(state, runtime)
request = _make_model_request(state["messages"], runtime=runtime)
sent: list[list[Any]] = []
def flaky_handler(model_request: ModelRequest):
sent.append(model_request.messages)
if len(sent) == 1:
raise RuntimeError("503 Service Unavailable")
return "response"
with pytest.raises(RuntimeError):
mw.wrap_model_call(request, flaky_handler)
assert mw.wrap_model_call(request, flaky_handler) == "response"
assert [len(_todo_completion_reminders(messages)) for messages in sent] == [1, 1]
# Restoring is not a new reminder: the retry must not spend the cap.
assert mw._completion_reminder_count_for_runtime(runtime) == 1
assert mw._drain_completion_reminders(runtime) == []
def test_failed_model_call_does_not_restore_a_reminder_for_a_cleared_run(self):
"""If the run's reminder state was dropped while the call was in flight, the failure must not resurrect it."""
mw = TodoMiddleware()
runtime = _make_runtime()
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
mw.after_model(state, runtime)
request = _make_model_request(state["messages"], runtime=runtime)
def failing_handler(model_request: ModelRequest):
mw.after_agent({}, runtime)
raise RuntimeError("503 Service Unavailable")
with pytest.raises(RuntimeError):
mw.wrap_model_call(request, failing_handler)
assert mw._pending_completion_reminders == {}
assert mw._completion_reminder_counts == {}
@pytest.fixture
def production_chain_config(tmp_path, monkeypatch):
from deerflow.config.app_config import AppConfig, pop_current_app_config, push_current_app_config
from deerflow.sandbox.sandbox_provider import shutdown_sandbox_provider
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
# A missing ambient config must not be masked by a developer's checkout.
monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(tmp_path / "missing-config.yaml"))
monkeypatch.setenv("LANGSMITH_TRACING", "false")
monkeypatch.setenv("LANGCHAIN_TRACING_V2", "false")
monkeypatch.setattr("deerflow.config.paths._paths", None)
monkeypatch.setattr("deerflow.sandbox.sandbox_provider._default_sandbox_provider", None)
monkeypatch.setattr("deerflow.sandbox.local.local_sandbox_provider._singleton", None)
app_config = AppConfig.model_validate(
{
"models": [{"name": "offline", "use": "langchain_openai:ChatOpenAI", "model": "offline"}],
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider", "allow_host_bash": False},
"skills": {"path": str(tmp_path / "skills")},
"memory": {"enabled": False},
"title": {"enabled": False},
"summarization": {"enabled": False},
}
)
# Sandbox lifecycle hooks resolve get_app_config() at invocation time,
# independently of the explicit config passed to build_middlewares().
# The runtime binding propagates through both graph and worker-thread calls.
push_current_app_config(app_config)
try:
yield app_config
finally:
try:
shutdown_sandbox_provider()
finally:
pop_current_app_config()
class TestTodoMiddlewareAgentGraphIntegration:
@pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("response_kind", ["tool", "blank", "partial_text"])
def test_production_chain_stops_after_safety_response(self, production_chain_config, invoke_async, response_kind):
from deerflow.agents.lead_agent.agent import build_middlewares
executed = []
@tool
def record_action(marker: str) -> str:
"""Record a harmless action in memory."""
executed.append(marker)
return marker
filtered = AIMessage(
id="filtered-response",
content="Partial report." if response_kind == "partial_text" else "",
response_metadata={"finish_reason": "content_filter"},
tool_calls=[{"name": "record_action", "args": {"marker": "filtered"}, "id": "filtered-call"}] if response_kind == "tool" else [],
)
model = _CapturingFakeMessagesListChatModel(
responses=[
filtered,
AIMessage(content="", tool_calls=[{"name": "record_action", "args": {"marker": "after-stop"}, "id": "later-call"}]),
AIMessage(content="final"),
]
)
middlewares = build_middlewares({"configurable": {"is_plan_mode": True}}, model_name="offline", app_config=production_chain_config, memory_enabled=False, available_skills=set())
graph = create_agent(model=model, tools=[record_action], middleware=middlewares, state_schema=ThreadState)
todos = _incomplete_todos()
context = {"thread_id": "production-safety-thread", "run_id": "production-safety-run"}
state = {"messages": [("user", "Finish the report.")], "todos": todos}
result = asyncio.run(graph.ainvoke(state, context=context)) if invoke_async else graph.invoke(state, context=context)
assert len(model.seen_messages) == 1
assert executed == []
assert context["stop_reason"] == "safety_capped"
final = result["messages"][-1]
assert final.additional_kwargs["safety_termination"]
assert final.content
assert not final.tool_calls
assert result["todos"] == todos
if response_kind == "partial_text":
assert final.content == filtered.content
todo_mw = next(mw for mw in middlewares if isinstance(mw, TodoMiddleware))
assert todo_mw._pending_completion_reminders == {}
assert todo_mw._completion_reminder_counts == {}
@pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("guard_kind", ["loop", "token", "safety", "model_length"])
@pytest.mark.parametrize("pending", [True, False], ids=["unfinished", "completed"])
def test_run_hard_stop_does_not_reengage_todos(self, invoke_async, guard_kind, pending):
from deerflow.agents.middlewares.loop_detection_middleware import LoopDetectionMiddleware
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
from deerflow.config.token_budget_config import TokenBudgetConfig
executed = []
@tool
def record_action(marker: str) -> str:
"""Record a harmless action in memory."""
executed.append(marker)
return marker
def action(marker, call_id, tokens=10):
return AIMessage(
content="",
tool_calls=[{"name": "record_action", "args": {"marker": marker}, "id": call_id}],
usage_metadata={"input_tokens": tokens, "output_tokens": 0, "total_tokens": tokens},
)
todos = [{"content": "Record the requested actions", "status": "pending" if pending else "completed"}]
blocked_action = action("repeated", "blocked", tokens=1000 if guard_kind == "token" else 10)
if guard_kind in ("safety", "model_length"):
blocked_action.response_metadata = {"finish_reason": "content_filter" if guard_kind == "safety" else "length"}
model = _CapturingFakeMessagesListChatModel(
responses=[
AIMessage(content="", tool_calls=[{"name": "write_todos", "args": {"todos": todos}, "id": "todos-start"}]),
action("repeated", "first"),
blocked_action,
# A different action makes unintended post-stop execution observable.
action("after-stop", "after-stop"),
AIMessage(content="final answer"),
AIMessage(content="final answer"),
AIMessage(content="final answer"),
]
)
todo_mw = TodoMiddleware()
guards = {
"loop": LoopDetectionMiddleware(warn_threshold=2, hard_limit=2),
"token": TokenBudgetMiddleware(TokenBudgetConfig(enabled=True, max_tokens=1000)),
"safety": SafetyFinishReasonMiddleware(),
"model_length": ModelLengthFinishReasonMiddleware(),
}
guard = guards[guard_kind]
# Production registers Todo before the guards; after_model runs in reverse.
graph = create_agent(model=model, tools=[record_action], middleware=[todo_mw, guard], state_schema=ThreadState, checkpointer=InMemorySaver())
config = {"configurable": {"thread_id": "hard-stop-thread"}, "recursion_limit": 100}
context = {"thread_id": "hard-stop-thread", "run_id": "hard-stop-run"}
state = {"messages": [("user", "Record actions and track todos.")]}
def invoke(state, context):
return asyncio.run(graph.ainvoke(state, context=context, config=config)) if invoke_async else graph.invoke(state, context=context, config=config)
result = invoke(state, context)
assert context["stop_reason"] == f"{guard_kind}_capped"
assert executed == ["repeated"]
assert len(model.seen_messages) == 3
# Pin the real producers to Todo's declared policy, without importing
# runtime Goal policy (which currently covers only loop/token).
assert context["stop_reason"] in todo_mw.release_policy_parameters()["completion_reminder_hard_stop_reasons"]
assert result["todos"] == todos
assert not result["messages"][-1].tool_calls
assert not any(_todo_completion_reminders(messages) for messages in model.seen_messages)
assert todo_mw._pending_completion_reminders == {}
snapshot = graph.get_state(config)
assert snapshot.next == ()
assert snapshot.values["todos"] == todos
if pending:
# Restore the same thread from its checkpoint without resupplying Todo.
# The old cap must not disable the new run or recount historical usage.
next_context = {"thread_id": "hard-stop-thread", "run_id": "next-run"}
next_result = invoke({"messages": [("user", "Continue the unfinished task.")]}, next_context)
assert "stop_reason" not in next_context
assert executed == ["repeated", "after-stop"]
assert any(_todo_completion_reminders(messages) for messages in model.seen_messages[3:])
assert next_result["todos"] == todos
def test_reuses_thread_state_todos_schema_in_real_agent_graph(self):
mw = TodoMiddleware()
model = _CapturingFakeMessagesListChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
{
"name": "write_todos",
"id": "todos-1",
"args": {
"todos": [
{"content": "Step 1", "status": "pending"},
]
},
}
],
),
AIMessage(content="final"),
],
)
graph = create_agent(
model=model,
tools=[],
middleware=[mw],
state_schema=ThreadState,
)
result = graph.invoke(
{"messages": [("user", "create a todo")]},
context={"thread_id": "schema-thread", "run_id": "schema-run"},
)
assert result["todos"] == [{"content": "Step 1", "status": "pending"}]
# Regression for bytedance/deer-flow#4714: the model request must carry
# the `write_todos` system prompt (injected by the base class), otherwise
# the model never produces a todo list.
first_model_call = model.seen_messages[0]
assert isinstance(first_model_call[0], SystemMessage)
assert "write_todos" in first_model_call[0].text
def test_completion_reminder_is_transient_in_real_agent_graph(self):
mw = TodoMiddleware()
model = _CapturingFakeMessagesListChatModel(
responses=[
AIMessage(
content="",
tool_calls=[
{
"name": "write_todos",
"id": "todos-1",
"args": {
"todos": [
{"content": "Step 1", "status": "completed"},
{"content": "Step 2", "status": "pending"},
]
},
}
],
),
AIMessage(content="premature final 1"),
AIMessage(content="premature final 2"),
AIMessage(content="premature final 3"),
],
)
graph = create_agent(model=model, tools=[], middleware=[mw])
result = graph.invoke(
{"messages": [("user", "finish all todos")]},
context={"thread_id": "integration-thread", "run_id": "integration-run"},
)
assert len(model.seen_messages) == 4
reminders_by_call = [_todo_completion_reminders(messages) for messages in model.seen_messages]
assert reminders_by_call[0] == []
assert reminders_by_call[1] == []
assert len(reminders_by_call[2]) == 1
assert len(reminders_by_call[3]) == 1
assert "Step 1" not in reminders_by_call[2][0].content
assert "Step 2" in reminders_by_call[2][0].content
persisted_reminders = _todo_completion_reminders(result["messages"])
assert persisted_reminders == []
assert result["messages"][-1].content == "premature final 3"
assert result["todos"] == [
{"content": "Step 1", "status": "completed"},
{"content": "Step 2", "status": "pending"},
]
assert mw._pending_completion_reminders == {}
assert mw._completion_reminder_counts == {}
def test_completion_reminder_survives_a_retried_model_call_in_real_agent_graph(self):
"""LLMErrorHandlingMiddleware retries a failed call by running the inner wraps again; the retry must still carry the reminder."""
from deerflow.agents.middlewares.llm_error_handling_middleware import LLMErrorHandlingMiddleware
from deerflow.config.app_config import AppConfig, LlmCallConfig
from deerflow.config.sandbox_config import SandboxConfig
class ProviderUnavailable(Exception):
def __init__(self) -> None:
super().__init__("503 Service Unavailable")
self.status_code = 503
self.response = SimpleNamespace(status_code=503, headers={})
class FailsOnceOnReminder(_CapturingFakeMessagesListChatModel):
_failed: bool = PrivateAttr(default=False)
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
if not self._failed and _todo_completion_reminders(messages):
self._failed = True
self._seen_messages.append(list(messages))
raise ProviderUnavailable()
return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
model = FailsOnceOnReminder(
responses=[
AIMessage(
content="",
tool_calls=[
{
"name": "write_todos",
"id": "todos-1",
"args": {"todos": [{"content": "Step 1", "status": "pending"}]},
}
],
),
AIMessage(content="premature final 1"),
AIMessage(content="premature final 2"),
AIMessage(content="premature final 3"),
],
)
app_config = AppConfig(
sandbox=SandboxConfig(use="test"),
llm_call=LlmCallConfig(retry_max_attempts=3, retry_base_delay_ms=0, retry_cap_delay_ms=0),
)
graph = create_agent(
model=model,
tools=[],
middleware=[LLMErrorHandlingMiddleware(app_config=app_config), TodoMiddleware()],
state_schema=ThreadState,
)
result = graph.invoke(
{"messages": [("user", "finish all todos")]},
context={"thread_id": "retry-thread", "run_id": "retry-run"},
)
# write_todos, the premature exit, the failed attempt and its retry, then
# the second reminder: the retry did not spend the two-reminder cap.
assert [len(_todo_completion_reminders(messages)) for messages in model.seen_messages] == [0, 0, 1, 1, 1]
assert result["messages"][-1].content == "premature final 3"
@pytest.mark.parametrize("invoke_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize(
"content,response_metadata,additional_kwargs",
[
("partial answer", {"finish_reason": "length"}, {}),
("partial answer", {"stop_reason": "max_tokens"}, {}),
([{"type": "text", "text": "partial answer"}], {"finish_reason": "MAX_TOKENS"}, {}),
("partial answer", {"status": "incomplete", "incomplete_details": {"reason": "max_output_tokens"}}, {}),
("partial answer", {}, {"finish_reason": "length"}),
],
ids=["openai", "anthropic", "gemini-blocks", "responses", "additional-kwargs"],
)
def test_text_only_length_cap_does_not_reengage_todos(self, invoke_async, content, response_metadata, additional_kwargs):
todo_mw = TodoMiddleware()
capped_message = AIMessage(content=content, response_metadata=response_metadata, additional_kwargs=additional_kwargs)
model = _CapturingFakeMessagesListChatModel(responses=[capped_message])
graph = create_agent(
model=model,
tools=[],
middleware=[todo_mw, ModelLengthFinishReasonMiddleware()],
state_schema=ThreadState,
)
todos = [{"content": "Finish the report", "status": "in_progress"}]
state = {"messages": [("user", "write the report")], "todos": todos}
context = {"thread_id": "text-cap-thread", "run_id": "text-cap-run"}
result = asyncio.run(graph.ainvoke(state, context=context)) if invoke_async else graph.invoke(state, context=context)
assert len(model.seen_messages) == 1
final_ai = result["messages"][-1]
assert final_ai.content == content
assert final_ai.response_metadata == response_metadata
assert final_ai.additional_kwargs["model_length_termination"]["suppressed_tool_call_count"] == 0
assert all(final_ai.additional_kwargs[key] == value for key, value in additional_kwargs.items())
assert "model_length_termination" not in capped_message.additional_kwargs
assert result["todos"] == todos
assert todo_mw._pending_completion_reminders == {}
assert not any(_todo_completion_reminders(messages) for messages in model.seen_messages)
def test_length_capped_write_file_does_not_reengage_todos(self):
"""Reproduces the incident (thread b1723286): model emits a write_file
call with finish_reason=length, ModelLength suppresses it, and
TodoMiddleware must NOT re-engage via jump_to=model."""
todo_mw = TodoMiddleware()
model = _CapturingFakeMessagesListChatModel(
responses=[
AIMessage(
content="nit",
tool_calls=[{"name": "write_file", "id": "call_1", "args": {"path": "/mnt/user-data/outputs/report.md", "content": "# truncated report\n| ext4 | jbd2"}}],
response_metadata={"finish_reason": "length", "model_name": "deepseek-v4-pro"},
),
],
)
graph = create_agent(
model=model,
tools=[],
middleware=[todo_mw, ModelLengthFinishReasonMiddleware()],
state_schema=ThreadState,
)
result = graph.invoke(
{"messages": [("user", "write the report")], "todos": [{"content": "Phase 9: Synthesis and final report writing", "status": "in_progress"}]},
context={"thread_id": "cap-thread", "run_id": "cap-run"},
)
assert len(model.seen_messages) == 1
reminders_by_call = [_todo_completion_reminders(messages) for messages in model.seen_messages]
assert all(len(r) == 0 for r in reminders_by_call)
final_ai = result["messages"][-1]
assert final_ai.additional_kwargs.get("model_length_termination")
assert "output limit" in str(final_ai.content)
assert "nit" in str(final_ai.content)
assert result["todos"] == [{"content": "Phase 9: Synthesis and final report writing", "status": "in_progress"}]
assert todo_mw._pending_completion_reminders == {}
class TestRunScopedReminderCleanup:
def test_before_agent_clears_stale_count_without_pending_reminder(self):
mw = TodoMiddleware()
stale_runtime = _make_runtime()
stale_runtime.context = {"thread_id": "test-thread", "run_id": "stale-run"}
current_runtime = _make_runtime()
current_runtime.context = {"thread_id": "test-thread", "run_id": "current-run"}
other_thread_runtime = _make_runtime()
other_thread_runtime.context = {"thread_id": "other-thread", "run_id": "stale-run"}
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
assert mw.after_model(state, stale_runtime) is not None
assert mw.after_model(state, other_thread_runtime) is not None
# Simulate a model call that drained the pending message, followed by an
# abnormal run end where after_agent did not clear the reminder count.
assert mw._drain_completion_reminders(stale_runtime)
assert mw._completion_reminder_count_for_runtime(stale_runtime) == 1
mw.before_agent({}, current_runtime)
assert mw._completion_reminder_count_for_runtime(stale_runtime) == 0
assert mw._completion_reminder_count_for_runtime(other_thread_runtime) == 1
def test_size_guard_prunes_oldest_count_only_reminder_state(self):
mw = TodoMiddleware()
mw._MAX_COMPLETION_REMINDER_KEYS = 2
first_runtime = _make_runtime_for("thread-a", "run-a")
second_runtime = _make_runtime_for("thread-b", "run-b")
third_runtime = _make_runtime_for("thread-c", "run-c")
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
assert mw.after_model(state, first_runtime) is not None
# Simulate the normal model request path: pending reminder is consumed,
# but the run count remains until after_agent() or stale cleanup.
assert mw._drain_completion_reminders(first_runtime)
assert mw._completion_reminder_count_for_runtime(first_runtime) == 1
assert mw.after_model(state, second_runtime) is not None
assert mw.after_model(state, third_runtime) is not None
assert mw._completion_reminder_count_for_runtime(first_runtime) == 0
assert mw._completion_reminder_count_for_runtime(second_runtime) == 1
assert mw._completion_reminder_count_for_runtime(third_runtime) == 1
assert ("thread-a", "run-a") not in mw._completion_reminder_touch_order
def test_size_guard_prunes_pending_and_count_state_together(self):
mw = TodoMiddleware()
mw._MAX_COMPLETION_REMINDER_KEYS = 1
stale_runtime = _make_runtime_for("thread-a", "run-a")
current_runtime = _make_runtime_for("thread-b", "run-b")
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
assert mw.after_model(state, stale_runtime) is not None
assert mw.after_model(state, current_runtime) is not None
assert mw._drain_completion_reminders(stale_runtime) == []
assert mw._completion_reminder_count_for_runtime(stale_runtime) == 0
assert mw._completion_reminder_count_for_runtime(current_runtime) == 1
class TestAwrapModelCall:
def test_async_no_pending_reminder_still_injects_todo_system_prompt(self):
"""Mirror of the sync no-reminder test: the async path must also keep the
base class system-prompt injection while leaving messages untouched when
there are no pending completion reminders."""
mw = TodoMiddleware()
runtime = _make_runtime()
request = _make_model_request([HumanMessage(content="hi")], runtime=runtime)
seen: list[ModelRequest] = []
async def handler(model_request: ModelRequest):
seen.append(model_request)
return "response"
result = asyncio.run(mw.awrap_model_call(request, handler))
assert result == "response"
assert len(seen) == 1
sent = seen[0]
assert sent.system_message is not None
assert "write_todos" in sent.system_message.text
assert sent.messages == [HumanMessage(content="hi")]
def test_async_pending_reminder_is_injected(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {
"messages": [_ai_no_tool_calls()],
"todos": _incomplete_todos(),
}
mw.after_model(state, runtime)
request = _make_model_request(state["messages"], runtime=runtime)
seen: list[ModelRequest] = []
async def handler(model_request: ModelRequest):
seen.append(model_request)
return "response"
result = asyncio.run(mw.awrap_model_call(request, handler))
assert result == "response"
assert len(seen) == 1
sent = seen[0]
assert sent.system_message is not None
assert "write_todos" in sent.system_message.text
injected = sent.messages[-1]
assert isinstance(injected, HumanMessage)
assert injected.name == "todo_completion_reminder"
assert injected.additional_kwargs["hide_from_ui"] is True
def test_async_pending_reminder_survives_a_failed_model_call(self):
mw = TodoMiddleware()
runtime = _make_runtime()
state = {"messages": [_ai_no_tool_calls()], "todos": _incomplete_todos()}
mw.after_model(state, runtime)
request = _make_model_request(state["messages"], runtime=runtime)
sent: list[list[Any]] = []
async def flaky_handler(model_request: ModelRequest):
sent.append(model_request.messages)
if len(sent) == 1:
raise RuntimeError("503 Service Unavailable")
return "response"
with pytest.raises(RuntimeError):
asyncio.run(mw.awrap_model_call(request, flaky_handler))
assert asyncio.run(mw.awrap_model_call(request, flaky_handler)) == "response"
assert [len(_todo_completion_reminders(messages)) for messages in sent] == [1, 1]
assert mw._completion_reminder_count_for_runtime(runtime) == 1
class TestMalformedStoredTodos:
"""State is a boundary: a stored value need not match the write_todos shape.
``todos`` is a reducer channel that checkpoint writes and direct integrations
can Overwrite directly, so the middleware must not assume every entry is a
mapping. One unknown shape used to raise on every model call of the thread.
"""
def test_before_model_ignores_entries_it_cannot_render(self):
mw = TodoMiddleware()
state = {"messages": [HumanMessage(content="hi")], "todos": ["oops", {"content": "Keep me"}]}
result = mw.before_model(state, _make_runtime())
content = result["messages"][0].content
assert "Keep me" in content
assert "oops" not in content
def test_after_model_keeps_guarding_for_the_usable_entries(self):
mw = TodoMiddleware()
state = {"messages": [AIMessage(content="final")], "todos": ["oops", {"content": "Keep me", "status": "pending"}]}
assert mw.after_model(state, _make_runtime()) == {"jump_to": "model"}
def test_after_model_ignores_malformed_siblings_of_completed_entries(self):
mw = TodoMiddleware()
state = {"messages": [AIMessage(content="final")], "todos": [{"content": "Done", "status": "completed"}, "junk"]}
assert mw.after_model(state, _make_runtime()) is None
@pytest.mark.parametrize("todos", [{"content": "x"}, "oops", 5])
def test_todos_that_are_not_a_list_allow_a_clean_exit(self, todos):
mw = TodoMiddleware()
assert mw.after_model({"messages": [AIMessage(content="final")], "todos": todos}, _make_runtime()) is None
assert mw.before_model({"messages": [HumanMessage(content="hi")], "todos": todos}, _make_runtime()) is None