Skip merged messages when inserting new ones

Antonio Scandurra and Julia Risley created 2 years ago

Co-Authored-By: Julia Risley <[e-mail]>

Change summary

crates/ai/src/assistant.rs | 31 +++++++++++++++++++++++++++++--
1 file changed, 29 insertions(+), 2 deletions(-)

Detailed changes

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(