From f2323b67b1239b92cfda524eea321724d13138b4 Mon Sep 17 00:00:00 2001
From: Suhaib Mujahid <smujahid@mozilla.com>
Date: Tue, 17 Mar 2026 16:58:54 -0400
Subject: [PATCH] Add TodoListMiddleware for code review planning

Integrate TodoListMiddleware into the code review agent to enable tracking of investigation tasks during reviews.
---
 bugbug/tools/code_review/agent.py   |  9 +++++++++
 bugbug/tools/code_review/prompts.py | 21 +++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/bugbug/tools/code_review/agent.py b/bugbug/tools/code_review/agent.py
index df65409275..6656df3688 100644
--- a/bugbug/tools/code_review/agent.py
+++ b/bugbug/tools/code_review/agent.py
@@ -12,6 +12,7 @@
 from typing import Optional
 
 from langchain.agents import create_agent
+from langchain.agents.middleware import TodoListMiddleware
 from langchain.agents.structured_output import ProviderStrategy
 from langchain.chat_models import BaseChatModel, init_chat_model
 from langchain.messages import HumanMessage
@@ -32,6 +33,8 @@
     expand_context,
 )
 from bugbug.tools.code_review.prompts import (
+    CODE_REVIEW_TODO_PROMPT,
+    CODE_REVIEW_TODO_TOOL_DESCRIPTION,
     FIRST_MESSAGE_TEMPLATE,
     STATIC_COMMENT_EXAMPLES,
     SYSTEM_PROMPT_TEMPLATE,
@@ -103,6 +106,12 @@ def __init__(
                 target_software=self.target_software,
             ),
             response_format=ProviderStrategy(AgentResponse),
+            middleware=[
+                TodoListMiddleware(
+                    system_prompt=CODE_REVIEW_TODO_PROMPT,
+                    tool_description=CODE_REVIEW_TODO_TOOL_DESCRIPTION,
+                )
+            ],
         )
 
         self.review_comments_db = review_comments_db
diff --git a/bugbug/tools/code_review/prompts.py b/bugbug/tools/code_review/prompts.py
index 8ee20da13a..639b36abc9 100644
--- a/bugbug/tools/code_review/prompts.py
+++ b/bugbug/tools/code_review/prompts.py
@@ -187,6 +187,27 @@
     },
 ]
 
+CODE_REVIEW_TODO_PROMPT = """
+## Review Planning with `write_todos`
+
+Use the `write_todos` tool to track investigation tasks as you review.
+
+- After your initial scan, create todos for any concerns that need deeper investigation
+  (e.g., "Verify that removed error handler is covered elsewhere", "Check callers of
+  renamed function for breakage")
+- As the review progresses, add new todos when you discover additional concerns
+- Remove or complete todos that turn out to be non-issues after verification
+- For small or straightforward patches, skip todos entirely — just review directly
+"""
+
+CODE_REVIEW_TODO_TOOL_DESCRIPTION = (
+    "Track investigation tasks during code review. Add items for concerns that need "
+    "tool-based verification (expand_context, find_function_definition). Evolve the "
+    "list as you go — add new items when you discover concerns, remove irrelevant ones. "
+    "Do not use this as a file checklist."
+)
+
+
 TEMPLATE_PATCH_FROM_HUNK = """diff --git a/{filename} b/{filename}
 --- a/{filename}
 +++ b/{filename}