diff --git a/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py index 15125b74f..fea19d424 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/todo_middleware.py @@ -337,7 +337,11 @@ class TodoMiddleware(TodoListMiddleware): request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse], ) -> ModelCallResult: - return handler(self._augment_request(request)) + # The base class appends the `write_todos` system prompt to the request; + # without calling it the model is never told about the todo list feature. + # Augment with pending completion reminders on the request that already + # carries the injected system prompt. + return super().wrap_model_call(request, lambda req: handler(self._augment_request(req))) @override async def awrap_model_call( @@ -345,7 +349,11 @@ class TodoMiddleware(TodoListMiddleware): request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]], ) -> ModelCallResult: - return await handler(self._augment_request(request)) + # See wrap_model_call: preserve the base class system-prompt injection. + async def augmented_handler(req: ModelRequest) -> ModelResponse: + return await handler(self._augment_request(req)) + + return await super().awrap_model_call(request, augmented_handler) @override def after_agent(self, state: ThreadState, runtime: Runtime) -> dict[str, Any] | None: diff --git a/backend/tests/test_todo_middleware.py b/backend/tests/test_todo_middleware.py index 4e7b2c61c..8a7b52c55 100644 --- a/backend/tests/test_todo_middleware.py +++ b/backend/tests/test_todo_middleware.py @@ -2,11 +2,12 @@ import asyncio from typing import Any -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock from langchain.agents import create_agent +from langchain.agents.middleware.types import ModelRequest from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel -from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage from pydantic import PrivateAttr from deerflow.agents.middlewares.todo_middleware import ( @@ -59,6 +60,15 @@ def _make_runtime_for(thread_id: str, run_id: str): 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"}, @@ -342,21 +352,24 @@ class TestAfterModel: assert result["jump_to"] == "model" assert "messages" not in result - request = MagicMock() - request.runtime = runtime - request.messages = state["messages"] - request.override.return_value = "patched-request" - handler = MagicMock(return_value="response") + 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" - request.override.assert_called_once() - reminder = request.override.call_args.kwargs["messages"][-1] + 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 - handler.assert_called_once_with("patched-request") def test_reminder_lists_only_incomplete_items(self): mw = TodoMiddleware() @@ -368,12 +381,15 @@ class TestAfterModel: result = mw.after_model(state, runtime) assert result is not None - request = MagicMock() - request.runtime = runtime - request.messages = state["messages"] - request.override.return_value = "patched-request" - mw.wrap_model_call(request, MagicMock(return_value="response")) - content = request.override.call_args.kwargs["messages"][-1].content + 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 @@ -453,16 +469,26 @@ class TestAafterModel: class TestWrapModelCall: - def test_no_pending_reminder_passthrough(self): + 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() - request = MagicMock() - request.runtime = _make_runtime() - request.messages = [HumanMessage(content="hi")] - handler = MagicMock(return_value="response") + 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" - request.override.assert_not_called() - handler.assert_called_once_with(request) + 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() @@ -473,22 +499,28 @@ class TestWrapModelCall: } mw.after_model(state, runtime) - request = MagicMock() - request.runtime = runtime - request.messages = state["messages"] - request.override.return_value = "patched-request" - handler = MagicMock(return_value="response") + 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" - injected_messages = request.override.call_args.kwargs["messages"] - assert injected_messages[-1].name == "todo_completion_reminder" + 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" - request.override.reset_mock() - handler.reset_mock() - handler.return_value = "second-response" - assert mw.wrap_model_call(request, handler) == "second-response" - request.override.assert_not_called() - handler.assert_called_once_with(request) + # 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"] class TestTodoMiddlewareAgentGraphIntegration: @@ -528,6 +560,13 @@ class TestTodoMiddlewareAgentGraphIntegration: 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( @@ -642,6 +681,27 @@ class TestRunScopedReminderCleanup: 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() @@ -651,14 +711,20 @@ class TestAwrapModelCall: } mw.after_model(state, runtime) - request = MagicMock() - request.runtime = runtime - request.messages = state["messages"] - request.override.return_value = "patched-request" - handler = AsyncMock(return_value="response") + 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" - injected_messages = request.override.call_args.kwargs["messages"] - assert injected_messages[-1].name == "todo_completion_reminder" - handler.assert_awaited_once_with("patched-request") + 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