feat: integrate todo middleware

backend/src/agents/lead_agent/agent.py

@@ -1,5 +1,5 @@
 from langchain.agents import create_agent
-from langchain.agents.middleware import SummarizationMiddleware
+from langchain.agents.middleware import SummarizationMiddleware, TodoListMiddleware
 from langchain_core.runnables import RunnableConfig
 
 from src.agents.lead_agent.prompt import apply_prompt_template
@@ -54,10 +54,134 @@ def _create_summarization_middleware() -> SummarizationMiddleware | None:
     return SummarizationMiddleware(**kwargs)
 
 
+def _create_todo_list_middleware(is_plan_mode: bool) -> TodoListMiddleware | None:
+    """Create and configure the TodoList middleware.
+
+    Args:
+        is_plan_mode: Whether to enable plan mode with TodoList middleware.
+
+    Returns:
+        TodoListMiddleware instance if plan mode is enabled, None otherwise.
+    """
+    if not is_plan_mode:
+        return None
+
+    # Custom prompts matching DeerFlow's style
+    system_prompt = """
+<todo_list_system>
+You have access to the `write_todos` tool to help you manage and track complex multi-step objectives.
+
+**CRITICAL RULES:**
+- Mark todos as completed IMMEDIATELY after finishing each step - do NOT batch completions
+- Keep EXACTLY ONE task as `in_progress` at any time (unless tasks can run in parallel)
+- Update the todo list in REAL-TIME as you work - this gives users visibility into your progress
+- DO NOT use this tool for simple tasks (< 3 steps) - just complete them directly
+
+**When to Use:**
+This tool is designed for complex objectives that require systematic tracking:
+- Complex multi-step tasks requiring 3+ distinct steps
+- Non-trivial tasks needing careful planning and execution
+- User explicitly requests a todo list
+- User provides multiple tasks (numbered or comma-separated list)
+- The plan may need revisions based on intermediate results
+
+**When NOT to Use:**
+- Single, straightforward tasks
+- Trivial tasks (< 3 steps)
+- Purely conversational or informational requests
+- Simple tool calls where the approach is obvious
+
+**Best Practices:**
+- Break down complex tasks into smaller, actionable steps
+- Use clear, descriptive task names
+- Remove tasks that become irrelevant
+- Add new tasks discovered during implementation
+- Don't be afraid to revise the todo list as you learn more
+
+**Task Management:**
+Writing todos takes time and tokens - use it when helpful for managing complex problems, not for simple requests.
+</todo_list_system>
+"""
+
+    tool_description = """Use this tool to create and manage a structured task list for complex work sessions.
+
+**IMPORTANT: Only use this tool for complex tasks (3+ steps). For simple requests, just do the work directly.**
+
+## When to Use
+
+Use this tool in these scenarios:
+1. **Complex multi-step tasks**: When a task requires 3 or more distinct steps or actions
+2. **Non-trivial tasks**: Tasks requiring careful planning or multiple operations
+3. **User explicitly requests todo list**: When the user directly asks you to track tasks
+4. **Multiple tasks**: When users provide a list of things to be done
+5. **Dynamic planning**: When the plan may need updates based on intermediate results
+
+## When NOT to Use
+
+Skip this tool when:
+1. The task is straightforward and takes less than 3 steps
+2. The task is trivial and tracking provides no benefit
+3. The task is purely conversational or informational
+4. It's clear what needs to be done and you can just do it
+
+## How to Use
+
+1. **Starting a task**: Mark it as `in_progress` BEFORE beginning work
+2. **Completing a task**: Mark it as `completed` IMMEDIATELY after finishing
+3. **Updating the list**: Add new tasks, remove irrelevant ones, or update descriptions as needed
+4. **Multiple updates**: You can make several updates at once (e.g., complete one task and start the next)
+
+## Task States
+
+- `pending`: Task not yet started
+- `in_progress`: Currently working on (can have multiple if tasks run in parallel)
+- `completed`: Task finished successfully
+
+## Task Completion Requirements
+
+**CRITICAL: Only mark a task as completed when you have FULLY accomplished it.**
+
+Never mark a task as completed if:
+- There are unresolved issues or errors
+- Work is partial or incomplete
+- You encountered blockers preventing completion
+- You couldn't find necessary resources or dependencies
+- Quality standards haven't been met
+
+If blocked, keep the task as `in_progress` and create a new task describing what needs to be resolved.
+
+## Best Practices
+
+- Create specific, actionable items
+- Break complex tasks into smaller, manageable steps
+- Use clear, descriptive task names
+- Update task status in real-time as you work
+- Mark tasks complete IMMEDIATELY after finishing (don't batch completions)
+- Remove tasks that are no longer relevant
+- **IMPORTANT**: When you write the todo list, mark your first task(s) as `in_progress` immediately
+- **IMPORTANT**: Unless all tasks are completed, always have at least one task `in_progress` to show progress
+
+Being proactive with task management demonstrates thoroughness and ensures all requirements are completed successfully.
+
+**Remember**: If you only need a few tool calls to complete a task and it's clear what to do, it's better to just do the task directly and NOT use this tool at all.
+"""
+
+    return TodoListMiddleware(system_prompt=system_prompt, tool_description=tool_description)
+
+
 # ThreadDataMiddleware must be before SandboxMiddleware to ensure thread_id is available
 # SummarizationMiddleware should be early to reduce context before other processing
+# TodoListMiddleware should be before ClarificationMiddleware to allow todo management
 # ClarificationMiddleware should be last to intercept clarification requests after model calls
-def _build_middlewares():
+def _build_middlewares(config: RunnableConfig):
+    """Build middleware chain based on runtime configuration.
+
+    Args:
+        config: Runtime configuration containing configurable options like is_plan_mode.
+
+    Returns:
+        List of middleware instances.
+    """
     middlewares = [ThreadDataMiddleware(), SandboxMiddleware()]
 
     # Add summarization middleware if enabled
@@ -65,6 +189,12 @@ def _build_middlewares():
     if summarization_middleware is not None:
         middlewares.append(summarization_middleware)
 
+    # Add TodoList middleware if plan mode is enabled
+    is_plan_mode = config.get("configurable", {}).get("is_plan_mode", False)
+    todo_list_middleware = _create_todo_list_middleware(is_plan_mode)
+    if todo_list_middleware is not None:
+        middlewares.append(todo_list_middleware)
+
     middlewares.extend([TitleMiddleware(), ClarificationMiddleware()])
     return middlewares
 
@@ -75,11 +205,12 @@ def make_lead_agent(config: RunnableConfig):
 
     thinking_enabled = config.get("configurable", {}).get("thinking_enabled", True)
     model_name = config.get("configurable", {}).get("model_name") or config.get("configurable", {}).get("model")
-    print(f"thinking_enabled: {thinking_enabled}, model_name: {model_name}")
+    is_plan_mode = config.get("configurable", {}).get("is_plan_mode", False)
+    print(f"thinking_enabled: {thinking_enabled}, model_name: {model_name}, is_plan_mode: {is_plan_mode}")
     return create_agent(
         model=create_chat_model(name=model_name, thinking_enabled=thinking_enabled),
         tools=get_available_tools(),
-        middleware=_build_middlewares(),
+        middleware=_build_middlewares(config),
         system_prompt=apply_prompt_template(),
         state_schema=ThreadState,
     )