From 6f85f38d47cf71ca7d884b6385aa9a8f5a999d19 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Tue, 27 Aug 2024 14:06:54 -0700
Subject: [PATCH] feat: add option to disable shell command suggestions

---
 aider/coders/base_coder.py |  3 +++
 tests/basic/test_coder.py  | 33 ++++++++++++++++++++++++++-
 tests/basic/test_main.py   | 46 --------------------------------------
 3 files changed, 35 insertions(+), 47 deletions(-)

diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py
index 93ad9a4e4..d641505c8 100755
--- a/aider/coders/base_coder.py
+++ b/aider/coders/base_coder.py
@@ -92,6 +92,7 @@ class Coder:
     add_cache_headers = False
     cache_warming_thread = None
     num_cache_warming_pings = 0
+    suggest_shell_commands = True
 
     @classmethod
     def create(
@@ -262,6 +263,8 @@ class Coder:
         self.rejected_urls = set()
         self.abs_root_path_cache = {}
 
+        self.suggest_shell_commands = suggest_shell_commands
+
         self.num_cache_warming_pings = num_cache_warming_pings
 
         if not fnames:
diff --git a/tests/basic/test_coder.py b/tests/basic/test_coder.py
index 79d5650c2..89c2b7c50 100644
--- a/tests/basic/test_coder.py
+++ b/tests/basic/test_coder.py
@@ -776,7 +776,38 @@ This command will print 'Hello, World!' to the console."""
             self.assertEqual(coder.shell_commands[0].strip(), 'echo "Hello, World!"')
 
             # Check if handle_shell_commands was called with the correct argument
-            coder.handle_shell_commands.assert_called_once_with('echo "Hello, World!"', ANY)
+            coder.handle_shell_commands.assert_called_once()
+
+    def test_no_suggest_shell_commands(self):
+        with GitTemporaryDirectory():
+            io = InputOutput(yes=True)
+            coder = Coder.create(self.GPT35, "diff", io=io, suggest_shell_commands=False)
+
+            def mock_send(*args, **kwargs):
+                coder.partial_response_content = """Here's a shell command to run:
+
+```bash
+echo "Hello, World!"
+```
+
+This command will print 'Hello, World!' to the console."""
+                coder.partial_response_function_call = dict()
+                return []
+
+            coder.send = mock_send
+
+            # Mock the handle_shell_commands method to check if it's called
+            coder.handle_shell_commands = MagicMock()
+
+            # Run the coder with a message
+            coder.run(with_message="Suggest a shell command")
+
+            # Check if the shell command was added to the list
+            self.assertEqual(len(coder.shell_commands), 1)
+            self.assertEqual(coder.shell_commands[0].strip(), 'echo "Hello, World!"')
+
+            # Check if handle_shell_commands was called with the correct argument
+            coder.handle_shell_commands.assert_not_called()
 
 
 if __name__ == "__main__":
diff --git a/tests/basic/test_main.py b/tests/basic/test_main.py
index 3d742f72a..ab7ffddb7 100644
--- a/tests/basic/test_main.py
+++ b/tests/basic/test_main.py
@@ -586,20 +586,6 @@ class TestMain(TestCase):
             self.assertIsInstance(coder, Coder)
             self.assertEqual(coder.repo_map.map_mul_no_files, 5)
 
-    def test_apply_shell_commands(self):
-        with GitTemporaryDirectory():
-            shell_md = Path("shell.md")
-            shell_md.write_text("```bash\ntouch file.txt\n```")
-
-            main(
-                ["--apply", "shell.md", "--yes"],
-                input=DummyInput(),
-                output=DummyOutput(),
-            )
-
-            # shell commands require explicit approval, not just --yes
-            self.assertFalse(Path("file.txt").exists())
-
     def test_suggest_shell_commands_default(self):
         with GitTemporaryDirectory():
             coder = main(
@@ -629,35 +615,3 @@ class TestMain(TestCase):
                 return_coder=True,
             )
             self.assertTrue(coder.suggest_shell_commands)
-
-    def test_apply_shell_commands_with_no_suggest(self):
-        with GitTemporaryDirectory():
-            shell_md = Path("shell.md")
-            shell_md.write_text("```bash\ntouch no_suggest_file.txt\n```")
-
-            with patch(
-                "aider.coders.base_coder.Coder.handle_shell_commands"
-            ) as mock_handle_shell_commands:
-                main(
-                    ["--apply", "shell.md", "--no-git"],
-                    input=DummyInput(),
-                    output=DummyOutput(),
-                )
-
-                # Make sure handle_shell_commands IS called when --no-suggest-shell-commands is not used
-                mock_handle_shell_commands.assert_called_once()
-
-            with patch(
-                "aider.coders.base_coder.Coder.handle_shell_commands"
-            ) as mock_handle_shell_commands:
-                main(
-                    ["--apply", "shell.md", "--no-suggest-shell-commands", "--no-git"],
-                    input=DummyInput(),
-                    output=DummyOutput(),
-                )
-
-                # Make sure handle_shell_commands is not called when --no-suggest-shell-commands is used
-                mock_handle_shell_commands.assert_not_called()
-
-            # Check that the file was not created in either case
-            self.assertFalse(Path("no_suggest_file.txt").exists())