refactor: Enhance file mention handling and context coder logic

Author: paul-gauthier

aider/coders/base_coder.py

@@ -1456,7 +1456,8 @@ def send_message(self, inp):
                 return
 
             try:
-                self.reply_completed()
+                if self.reply_completed():
+                    return
             except KeyboardInterrupt:
                 interrupted = True
 
@@ -1599,7 +1600,7 @@ def add_assistant_reply_to_cur_messages(self):
                 )
             ]
 
-    def get_file_mentions(self, content):
+    def get_file_mentions(self, content, ignore_current=False):
         words = set(word for word in content.split())
 
         # drop sentence punctuation from the end
@@ -1609,12 +1610,16 @@ def get_file_mentions(self, content):
         quotes = "\"'`*_"
         words = set(word.strip(quotes) for word in words)
 
-        addable_rel_fnames = self.get_addable_relative_files()
+        if ignore_current:
+            addable_rel_fnames = self.get_all_relative_files()
+            existing_basenames = {}
+        else:
+            addable_rel_fnames = self.get_addable_relative_files()
 
-        # Get basenames of files already in chat or read-only
-        existing_basenames = {os.path.basename(f) for f in self.get_inchat_relative_files()} | {
-            os.path.basename(self.get_rel_fname(f)) for f in self.abs_read_only_fnames
-        }
+            # Get basenames of files already in chat or read-only
+            existing_basenames = {os.path.basename(f) for f in self.get_inchat_relative_files()} | {
+                os.path.basename(self.get_rel_fname(f)) for f in self.abs_read_only_fnames
+            }
 
         mentioned_rel_fnames = set()
         fname_to_rel_fnames = {}

aider/coders/context_coder.py

@@ -1,3 +1,5 @@
+from aider.dump import dump
+
 from .base_coder import Coder
 from .context_prompts import ContextPrompts
 
@@ -7,3 +9,28 @@ class ContextCoder(Coder):
 
     edit_format = "context"
     gpt_prompts = ContextPrompts()
+
+    def reply_completed(self):
+        content = self.partial_response_content
+        if not content or not content.strip():
+            return True
+
+        dump(repr(content))
+        current_rel_fnames = set(self.get_inchat_relative_files())
+        mentioned_rel_fnames = set(self.get_file_mentions(content, ignore_current=True))
+
+        dump(current_rel_fnames)
+        dump(mentioned_rel_fnames)
+        dump(current_rel_fnames == mentioned_rel_fnames)
+
+        if mentioned_rel_fnames != current_rel_fnames:
+            self.abs_fnames = set()
+            for fname in mentioned_rel_fnames:
+                self.add_rel_fname(fname)
+            dump(self.get_inchat_relative_files())
+            self.reflected_message = self.gpt_prompts.try_again
+
+        return True
+
+    def check_for_file_mentions(self, content):
+        pass

aider/coders/context_prompts.py

@@ -43,4 +43,9 @@ class ContextPrompts(CoderPrompts):
 
     system_reminder = """
 NEVER RETURN CODE!
+"""
+
+    try_again = """I have updated the set of files added to the chat.
+Review them to decide if this is the correct set of files or if we need to add more.
+If this is the right set, just return the current list of files.
 """