mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-09 12:30:37 +00:00
* 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>
1285 lines
54 KiB
Python
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
|