Allow redoing edits performed by inline assistant after cancelling it

crates/ai/src/assistant.rs

@@ -442,7 +442,7 @@ impl AssistantPanel {
                     if let Some(transaction_id) = pending_assist.transaction_id {
                         editor.update(cx, |editor, cx| {
                             editor.buffer().update(cx, |buffer, cx| {
-                                buffer.undo_and_forget(transaction_id, cx)
+                                buffer.undo_transaction(transaction_id, cx)
                             });
                         });
                     }

crates/editor/src/multi_buffer.rs

@@ -824,13 +824,15 @@ impl MultiBuffer {
         None
     }
 
-    pub fn undo_and_forget(&mut self, transaction_id: TransactionId, cx: &mut ModelContext<Self>) {
+    pub fn undo_transaction(&mut self, transaction_id: TransactionId, cx: &mut ModelContext<Self>) {
         if let Some(buffer) = self.as_singleton() {
-            buffer.update(cx, |buffer, cx| buffer.undo_and_forget(transaction_id, cx));
+            buffer.update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
-        } else if let Some(transaction) = self.history.forget(transaction_id) {
+        } else if let Some(transaction) = self.history.remove_from_undo(transaction_id) {
-            for (buffer_id, transaction_id) in transaction.buffer_transactions {
+            for (buffer_id, transaction_id) in &transaction.buffer_transactions {
-                if let Some(BufferState { buffer, .. }) = self.buffers.borrow().get(&buffer_id) {
+                if let Some(BufferState { buffer, .. }) = self.buffers.borrow().get(buffer_id) {
-                    buffer.update(cx, |buffer, cx| buffer.undo_and_forget(transaction_id, cx));
+                    buffer.update(cx, |buffer, cx| {
+                        buffer.undo_transaction(*transaction_id, cx)
+                    });
                 }
             }
         }
@@ -3454,6 +3456,16 @@ impl History {
         }
     }
 
+    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> Option<&Transaction> {
+        let ix = self
+            .undo_stack
+            .iter()
+            .rposition(|transaction| transaction.id == transaction_id)?;
+        let transaction = self.undo_stack.remove(ix);
+        self.redo_stack.push(transaction);
+        self.redo_stack.last()
+    }
+
     fn group(&mut self) -> Option<TransactionId> {
         let mut count = 0;
         let mut transactions = self.undo_stack.iter();

crates/language/src/buffer.rs

@@ -1668,14 +1668,14 @@ impl Buffer {
         }
     }
 
-    pub fn undo_and_forget(
+    pub fn undo_transaction(
         &mut self,
         transaction_id: TransactionId,
         cx: &mut ModelContext<Self>,
     ) -> bool {
         let was_dirty = self.is_dirty();
         let old_version = self.version.clone();
-        if let Some(operation) = self.text.undo_and_forget(transaction_id) {
+        if let Some(operation) = self.text.undo_transaction(transaction_id) {
             self.send_operation(Operation::Buffer(operation), cx);
             self.did_edit(&old_version, was_dirty, cx);
             true

crates/text/src/text.rs

@@ -264,7 +264,19 @@ impl History {
         }
     }
 
-    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> &[HistoryEntry] {
+    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> Option<&HistoryEntry> {
+        assert_eq!(self.transaction_depth, 0);
+
+        let entry_ix = self
+            .undo_stack
+            .iter()
+            .rposition(|entry| entry.transaction.id == transaction_id)?;
+        let entry = self.undo_stack.remove(entry_ix);
+        self.redo_stack.push(entry);
+        self.redo_stack.last()
+    }
+
+    fn remove_from_undo_until(&mut self, transaction_id: TransactionId) -> &[HistoryEntry] {
         assert_eq!(self.transaction_depth, 0);
 
         let redo_stack_start_len = self.redo_stack.len();
@@ -1207,19 +1219,20 @@ impl Buffer {
         }
     }
 
-    pub fn undo_and_forget(&mut self, transaction_id: TransactionId) -> Option<Operation> {
+    pub fn undo_transaction(&mut self, transaction_id: TransactionId) -> Option<Operation> {
-        if let Some(transaction) = self.history.forget(transaction_id) {
+        let transaction = self
-            self.undo_or_redo(transaction).log_err()
+            .history
-        } else {
+            .remove_from_undo(transaction_id)?
-            None
+            .transaction
-        }
+            .clone();
+        self.undo_or_redo(transaction).log_err()
     }
 
     #[allow(clippy::needless_collect)]
     pub fn undo_to_transaction(&mut self, transaction_id: TransactionId) -> Vec<Operation> {
         let transactions = self
             .history
-            .remove_from_undo(transaction_id)
+            .remove_from_undo_until(transaction_id)
             .iter()
             .map(|entry| entry.transaction.clone())
             .collect::<Vec<_>>();