Migrated sandbox tools to middleware.

Author: srtab

daiv/automation/agents/plan_and_execute/agent.py

@@ -18,13 +18,12 @@
 from automation.agents import BaseAgent
 from automation.agents.middleware import InjectImagesMiddleware
 from automation.agents.tools.navigation import READ_MAX_LINES
-from automation.agents.tools.sandbox import BASH_TOOL_NAME, format_code_tool
+from automation.agents.tools.sandbox import BASH_TOOL_NAME, FORMAT_CODE_TOOL_NAME, SandboxMiddleware
 from automation.agents.tools.toolkits import (
     FileEditingToolkit,
     FileNavigationToolkit,
     MCPToolkit,
     MergeRequestToolkit,
-    SandboxToolkit,
     WebSearchToolkit,
 )
 from codebase.context import RuntimeCtx
@@ -99,17 +98,7 @@ class ExecutorMiddleware(AgentMiddleware):
     Middleware to select the tools for the executor agent based on the tool calls.
     """
 
-    def __init__(self, *, enable_bash: bool = False, enable_format_code: bool = False):
-        """
-        Initialize the middleware.
-
-        Args:
-            enable_bash (bool): Whether to enable the bash tool.
-            enable_format_code (bool): Whether to enable the format code tool.
-        """
-        super().__init__()
-        self.enable_bash = enable_bash
-        self.enable_format_code = enable_format_code
+    name = "executor_middleware"
 
     async def abefore_agent(self, state: ExecutorState, runtime: Runtime[RuntimeCtx]) -> dict[str, Any] | None:
         """
@@ -138,14 +127,14 @@ async def awrap_model_call(
         Returns:
             ModelCallResult: The result of the model call.
         """
+        tools_names = [tool.name for tool in request.tools]
         request.system_prompt = execute_plan_system.format(
             current_date_time=timezone.now().strftime("%d %B, %Y"),
             repository=request.runtime.context.repo_id,
-            commands_enabled=self.enable_bash,
-            format_code_enabled=self.enable_format_code,
-            tools_names=[tool.name for tool in request.tools],
+            commands_enabled=BASH_TOOL_NAME in tools_names,
+            format_code_enabled=FORMAT_CODE_TOOL_NAME in tools_names,
+            tools_names=tools_names,
         ).content
-
         return await handler(request)
 
 
@@ -195,15 +184,21 @@ async def plan(
         Returns:
             Command[Literal["plan_approval", "__end__"]]: The next step in the workflow.
         """
-        mcp_tools = await MCPToolkit.get_tools()
-        file_navigation_tools = FileNavigationToolkit.get_tools()
-        web_search_tools = WebSearchToolkit.get_tools()
 
-        all_tools: list[BaseTool] = mcp_tools + file_navigation_tools + web_search_tools + [plan_think_tool]
+        all_tools: list[BaseTool] = (
+            (await MCPToolkit.get_tools())
+            + FileNavigationToolkit.get_tools()
+            + WebSearchToolkit.get_tools()
+            + [plan_think_tool]
+        )
 
         if runtime.context.merge_request_id:
             all_tools.extend(MergeRequestToolkit.get_tools())
 
+        conditional_middlewares: list[AgentMiddleware] = []
+        if runtime.context.config.sandbox.enabled:
+            conditional_middlewares.append(SandboxMiddleware(read_only_bash=True))
+
         planner_agent = create_agent(
             model=BaseAgent.get_model(
                 model=settings.PLANNING_MODEL_NAME, max_tokens=8_192, thinking_level=settings.PLANNING_THINKING_LEVEL
@@ -213,7 +208,12 @@ async def plan(
             checkpointer=False,
             context_schema=RuntimeCtx,
             response_format=ToolStrategy(FinalizerOutput),
-            middleware=[plan_system_prompt, InjectImagesMiddleware(), AnthropicPromptCachingMiddleware()],
+            middleware=[
+                plan_system_prompt,
+                InjectImagesMiddleware(),
+                *conditional_middlewares,
+                AnthropicPromptCachingMiddleware(),
+            ],
             name="planner_agent",
         )
 
@@ -274,27 +274,25 @@ async def execute_plan(
         Returns:
             Command[Literal["__end__"]]: The next step in the workflow.
         """
-        all_tools: list[BaseTool] = (
-            FileNavigationToolkit.get_tools() + FileEditingToolkit.get_tools() + [review_code_changes_tool]
-        )
-
+        conditional_middlewares: list[AgentMiddleware] = []
         if runtime.context.config.sandbox.enabled:
-            all_tools += SandboxToolkit.get_tools()
-
-            if not self.skip_format_code and runtime.context.config.sandbox.format_code_enabled:
-                all_tools.append(format_code_tool)
+            conditional_middlewares.append(
+                SandboxMiddleware(
+                    include_format_code=bool(
+                        not self.skip_format_code and runtime.context.config.sandbox.format_code_enabled
+                    )
+                )
+            )
 
         executor_agent = create_agent(
             model=BaseAgent.get_model(model=settings.EXECUTION_MODEL_NAME, max_tokens=8_192),
             state_schema=ExecutorState,
             context_schema=RuntimeCtx,
-            tools=all_tools,
+            tools=(FileNavigationToolkit.get_tools() + FileEditingToolkit.get_tools() + [review_code_changes_tool]),
             store=runtime.store,
             middleware=[
-                ExecutorMiddleware(
-                    enable_bash=runtime.context.config.sandbox.enabled,
-                    enable_format_code=not self.skip_format_code and runtime.context.config.sandbox.format_code_enabled,
-                ),
+                ExecutorMiddleware(),
+                *conditional_middlewares,
                 TodoListMiddleware(),
                 AnthropicPromptCachingMiddleware(),
             ],

daiv/automation/agents/plan_and_execute/prompts.py

@@ -69,7 +69,7 @@
 
 ### Phase 1: Understand (Required)
 1. **Plan your approach** using `think` - outline what you need to investigate
-2. **Gather context** using investigation tools (`ls`, `read`, `grep`, `glob`, `fetch`, `web_search`, etc.)
+2. **Gather context** using investigation tools (`ls`, `read`, `grep`, `glob`, `fetch`, `web_search`,{% if commands_enabled %} `bash`,{% endif %} etc.)
 3. **Update your understanding** with `think` as you learn new information
 
 ### Phase 2: Deliver (Required)

daiv/automation/agents/plan_and_execute/tools.py

@@ -26,7 +26,8 @@ def plan_think_tool(plan: str):
 
     **Usage rules:**
     - Does NOT fetch new information or modify anything, it's just a placeholder to help you track progress.
-    - Update tasks as you learn new information to help you track progress.
+    - Add any new follow-up tasks as you discover them during your investigation.
+    - You can also update future tasks, such as deleting them if they are no longer necessary, or adding new tasks that are necessary. Don't change previously completed tasks.
     - **Important:** It is critical that you mark tasks as completed as soon as you are done with them. Do not batch up multiple tasks before marking them as completed.
 
     **Skip using this tool when:**
@@ -35,8 +36,10 @@ def plan_think_tool(plan: str):
     - The task can be completed in less than 3 trivial steps
     - The task is purely conversational or informational
 
+    Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully.
+
     Args:
-        plan (str): The plan to investigate.
+        plan (str): The plan to investigate in markdown format.
 
     Returns:
         A message indicating that the thought has been registered.