From d1dd40590faf9fdfa3d9ed8a0901768c6cc7a249 Mon Sep 17 00:00:00 2001
From: "Paul Gauthier (aider)" <aider@paulg.org>
Date: Sun, 1 Dec 2024 08:25:06 -0800
Subject: [PATCH] feat: use TreeContext to show AI comments in file context

---
 aider/watch.py | 38 +++++++++++++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/aider/watch.py b/aider/watch.py
index 151de4376..5d0f0d997 100644
--- a/aider/watch.py
+++ b/aider/watch.py
@@ -172,17 +172,45 @@ class FileWatcher:
             if line_nums:
                 ai_comments[fname] = comments
 
-        # instead of just joining the comments, use TreeContext to render them in context of their files ai!
-        res = "\n".join(
-            comment for comments in ai_comments.values() if comments for comment in comments
-        )
+        from grep_ast import TreeContext
 
         res = """The "ai" comments below can be found in the code files I've shared with you.
     They contain your instructions.
     Make the requested changes.
     Be sure to remove all these "ai" comments from the code!
 
-    """ + res
+    """
+
+        for fname, comments in ai_comments.items():
+            if not comments:
+                continue
+
+            code = self.io.read_text(fname)
+            if not code:
+                continue
+
+            rel_fname = self.coder.get_rel_fname(fname)
+            res += f"\n{rel_fname}:\n"
+
+            # Convert comment line numbers to line indices (0-based)
+            lois = [ln - 1 for ln, _ in zip(line_nums, comments) if ln > 0]
+
+            context = TreeContext(
+                rel_fname,
+                code,
+                color=False,
+                line_number=False,
+                child_context=False,
+                last_line=False,
+                margin=0,
+                mark_lois=False,
+                loi_pad=0,
+                show_top_of_file_parent_scope=False,
+            )
+            context.lines_of_interest = set()
+            context.add_lines_of_interest(lois)
+            context.add_context()
+            res += context.format()
 
         dump(res)
         return res