mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
1f1192fd80
* fix(agents): keep the todo completion reminder when a model call is retried TodoMiddleware drained its queued completion reminder inside wrap_model_call before calling the handler. LLMErrorHandlingMiddleware sits outside it and retries a failed call by running that wrap again, so the retry went out with an empty queue while the run had already spent one of its two reminders on the lost one. Drain the reminder once per wrap, and when the handler raises put it back in front of the queue without counting it again, then re-raise. A run whose reminder state was cleared while the call was in flight does not get it back. * docs(changelog): reference #6132 in the todo reminder retry entry --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
936 lines
35 KiB
Python
936 lines
35 KiB
Python
"""Tests for TodoMiddleware context-loss detection."""
|
|
|
|
import asyncio
|
|
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.types import ModelRequest
|
|
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
|
from pydantic import PrivateAttr
|
|
|
|
from deerflow.agents.middlewares.model_length_finish_reason_middleware import (
|
|
ModelLengthFinishReasonMiddleware,
|
|
)
|
|
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_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_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 == {}
|
|
|
|
|
|
class TestTodoMiddlewareAgentGraphIntegration:
|
|
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"
|
|
|
|
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
|