From 0db0a1ccef443fe0ed943c25fec43267a073e5b6 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Tue, 13 Jun 2023 16:06:50 +0200
Subject: [PATCH] Skip merged messages when inserting new ones

Co-Authored-By: Julia Risley <julia@zed.dev>
---
 crates/ai/src/assistant.rs | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/crates/ai/src/assistant.rs b/crates/ai/src/assistant.rs
index 1777253f13..b650a6673f 100644
--- a/crates/ai/src/assistant.rs
+++ b/crates/ai/src/assistant.rs
@@ -661,8 +661,9 @@ impl Assistant {
         {
             let start = self.buffer.update(cx, |buffer, cx| {
                 let offset = self
-                    .messages
-                    .get(prev_message_ix + 1)
+                    .messages[prev_message_ix + 1..]
+                    .iter()
+                    .find(|message| message.start.is_valid(buffer))
                     .map_or(buffer.len(), |message| message.start.to_offset(buffer) - 1);
                 buffer.edit([(offset..offset, "\n")], None, cx);
                 buffer.anchor_before(offset + 1)
@@ -1451,6 +1452,7 @@ mod tests {
             ]
         );
 
+        // Undoing the deletion should also undo the merge.
         buffer.update(cx, |buffer, cx| buffer.undo(cx));
         assert_eq!(
             messages(&assistant, cx),
@@ -1461,6 +1463,31 @@ mod tests {
                 (message_3.id, Role::User, 6..7),
             ]
         );
+
+        // Redoing the deletion should also redo the merge.
+        buffer.update(cx, |buffer, cx| buffer.redo(cx));
+        assert_eq!(
+            messages(&assistant, cx),
+            vec![
+                (message_1.id, Role::User, 0..3),
+                (message_3.id, Role::User, 3..4),
+            ]
+        );
+
+        // Ensure we can still insert after a merged message.
+        let message_5 = assistant.update(cx, |assistant, cx| {
+            assistant
+                .insert_message_after(message_1.id, Role::System, cx)
+                .unwrap()
+        });
+        assert_eq!(
+            messages(&assistant, cx),
+            vec![
+                (message_1.id, Role::User, 0..3),
+                (message_5.id, Role::System, 3..4),
+                (message_3.id, Role::User, 4..5)
+            ]
+        );
     }
 
     fn messages(