From f9086e66d388435acb54567e3f4bc20dc4928e21 Mon Sep 17 00:00:00 2001
From: Paul Gauthier <aider@paulg.org>
Date: Wed, 21 Jun 2023 21:02:07 -0700
Subject: [PATCH] better live output of multi file edits

---
 aider/coders/base_coder.py | 16 ++++++++++------
 aider/coders/func_coder.py | 12 +++++++-----
 aider/diffs.py             | 13 ++++++++++---
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/aider/coders/base_coder.py b/aider/coders/base_coder.py
index 8c5ad1d15..59e6b6a48 100755
--- a/aider/coders/base_coder.py
+++ b/aider/coders/base_coder.py
@@ -527,19 +527,23 @@ class Coder:
                     continue
 
                 if self.pretty:
-                    show_resp = self.modify_incremental_response()
-                    if show_resp:
-                        md = Markdown(
-                            show_resp, style=self.assistant_output_color, code_theme="default"
-                        )
-                        live.update(md)
+                    self.live_incremental_response(live, False)
                 else:
                     sys.stdout.write(text)
                     sys.stdout.flush()
         finally:
+            self.live_incremental_response(live, True)
             if live:
                 live.stop()
 
+    def live_incremental_response(self, live, final):
+        show_resp = self.modify_incremental_response(final)
+        if not show_resp:
+            return
+
+        md = Markdown(show_resp, style=self.assistant_output_color, code_theme="default")
+        live.update(md)
+
     def modify_incremental_response(self):
         return self.partial_response_content
 
diff --git a/aider/coders/func_coder.py b/aider/coders/func_coder.py
index 57997c422..87fa39025 100644
--- a/aider/coders/func_coder.py
+++ b/aider/coders/func_coder.py
@@ -57,7 +57,7 @@ class FunctionCoder(Coder):
         else:
             self.cur_messages += [dict(role="assistant", content=content)]
 
-    def modify_incremental_response(self):
+    def modify_incremental_response(self, final=False):
         args = self.parse_partial_args()
 
         if not args:
@@ -70,7 +70,7 @@ class FunctionCoder(Coder):
         if explanation:
             res += f"{explanation}\n\n"
 
-        for file_upd in files:
+        for i, file_upd in enumerate(files):
             path = file_upd.get("path")
             if not path:
                 continue
@@ -79,11 +79,13 @@ class FunctionCoder(Coder):
                 continue
 
             res += path + ":\n"
-            res += self.live_diffs(path, content)
+
+            this_final = (i < len(files) - 1) or final
+            res += self.live_diffs(path, content, this_final)
 
         return res
 
-    def live_diffs(self, fname, content):
+    def live_diffs(self, fname, content, final):
         lines = content.splitlines(keepends=True)
 
         # ending an existing block
@@ -95,7 +97,7 @@ class FunctionCoder(Coder):
         show_diff = diffs.diff_partial_update(
             orig_lines,
             lines,
-            final=True,
+            final,
         ).splitlines()
 
         return "\n".join(show_diff)
diff --git a/aider/diffs.py b/aider/diffs.py
index 471212fc9..dd6297d1b 100644
--- a/aider/diffs.py
+++ b/aider/diffs.py
@@ -43,12 +43,17 @@ def diff_partial_update(lines_orig, lines_updated, final=False):
     # dump(lines_orig)
     # dump(lines_updated)
 
-    last_non_deleted = find_last_non_deleted(lines_orig, lines_updated)
+    num_orig_lines = len(lines_orig)
+
+    if final:
+        last_non_deleted = num_orig_lines
+    else:
+        last_non_deleted = find_last_non_deleted(lines_orig, lines_updated)
+
     # dump(last_non_deleted)
     if last_non_deleted is None:
         return ""
 
-    num_orig_lines = len(lines_orig)
     pct = last_non_deleted * 100 / num_orig_lines
     bar = create_progress_bar(pct)
     bar = f"! {last_non_deleted:3d} / {num_orig_lines:3d} lines [{bar}] {pct:3.0f}%\n\n"
@@ -63,12 +68,14 @@ def diff_partial_update(lines_orig, lines_updated, final=False):
     diff = list(diff)[2:]
 
     diff = "".join(diff)
+    if not diff.endswith("\n"):
+        diff += "\n"
 
     show = "```diff\n"
     if not final:
         show += bar
 
-    show += diff + "```\n"
+    show += diff + "```\n\n"
 
     # print(diff)