"""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_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_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 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 assert after_model() is None base_hook.assert_not_called() @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