From 739d5ca3735cd2ad69517b8d661186ac1f349356 Mon Sep 17 00:00:00 2001
From: Kirill Bulatov <kirill@zed.dev>
Date: Thu, 25 May 2023 18:07:38 +0300
Subject: [PATCH] Have proper undo for both client and host

---
 crates/collab/src/tests/integration_tests.rs | 45 ++++++++++++++++++--
 crates/editor/src/editor.rs                  | 30 ++++++++++---
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/crates/collab/src/tests/integration_tests.rs b/crates/collab/src/tests/integration_tests.rs
index ecbce24cf1b8740272bd3da13dfc653a428ccab0..d771f969d8899e0183d95db17a42389dbe706dd0 100644
--- a/crates/collab/src/tests/integration_tests.rs
+++ b/crates/collab/src/tests/integration_tests.rs
@@ -7464,9 +7464,6 @@ async fn test_on_input_format_from_host_to_guest(
             }]))
         },
     );
-    // .next()
-    // .await
-    // .unwrap();
 
     // Open the buffer on the guest and see that the formattings worked
     let buffer_b = project_b
@@ -7486,6 +7483,27 @@ async fn test_on_input_format_from_host_to_guest(
     buffer_b.read_with(cx_b, |buffer, _| {
         assert_eq!(buffer.text(), "fn main() { a>~< }")
     });
+
+    // Undo should remove LSP edits first
+    editor_a.update(cx_a, |editor, cx| {
+        assert_eq!(editor.text(cx), "fn main() { a>~< }");
+        editor.undo(&Undo, cx);
+        assert_eq!(editor.text(cx), "fn main() { a> }");
+    });
+    cx_b.foreground().run_until_parked();
+    buffer_b.read_with(cx_b, |buffer, _| {
+        assert_eq!(buffer.text(), "fn main() { a> }")
+    });
+
+    editor_a.update(cx_a, |editor, cx| {
+        assert_eq!(editor.text(cx), "fn main() { a> }");
+        editor.undo(&Undo, cx);
+        assert_eq!(editor.text(cx), "fn main() { a }");
+    });
+    cx_b.foreground().run_until_parked();
+    buffer_b.read_with(cx_b, |buffer, _| {
+        assert_eq!(buffer.text(), "fn main() { a }")
+    });
 }
 
 #[gpui::test(iterations = 10)]
@@ -7595,6 +7613,27 @@ async fn test_on_input_format_from_guest_to_host(
     buffer_a.read_with(cx_a, |buffer, _| {
         assert_eq!(buffer.text(), "fn main() { a:~: }")
     });
+
+    // Undo should remove LSP edits first
+    editor_b.update(cx_b, |editor, cx| {
+        assert_eq!(editor.text(cx), "fn main() { a:~: }");
+        editor.undo(&Undo, cx);
+        assert_eq!(editor.text(cx), "fn main() { a: }");
+    });
+    cx_a.foreground().run_until_parked();
+    buffer_a.read_with(cx_a, |buffer, _| {
+        assert_eq!(buffer.text(), "fn main() { a: }")
+    });
+
+    editor_b.update(cx_b, |editor, cx| {
+        assert_eq!(editor.text(cx), "fn main() { a: }");
+        editor.undo(&Undo, cx);
+        assert_eq!(editor.text(cx), "fn main() { a }");
+    });
+    cx_a.foreground().run_until_parked();
+    buffer_a.read_with(cx_a, |buffer, _| {
+        assert_eq!(buffer.text(), "fn main() { a }")
+    });
 }
 
 #[derive(Debug, Eq, PartialEq)]
diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs
index d944b9712456dba5ed40c8cdec34db52a16d0f5d..5a504a610cd63529e80e0b1a6f872f0fe08fb8bd 100644
--- a/crates/editor/src/editor.rs
+++ b/crates/editor/src/editor.rs
@@ -2524,15 +2524,33 @@ impl Editor {
             .buffer
             .read(cx)
             .text_anchor_for_position(position.clone(), cx)?;
+
+        // OnTypeFormatting retuns a list of edits, no need to pass them between Zed instances,
+        // hence we do LSP request & edit on host side only — add formats to host's history.
+        let push_to_lsp_host_history = true;
+        // If this is not the host, append its history with new edits.
+        let push_to_client_history = project.read(cx).is_remote();
+
         let on_type_formatting = project.update(cx, |project, cx| {
-            project.on_type_format(buffer, buffer_position, input, true, cx)
+            project.on_type_format(
+                buffer.clone(),
+                buffer_position,
+                input,
+                push_to_lsp_host_history,
+                cx,
+            )
         });
-
         Some(cx.spawn(|editor, mut cx| async move {
-            on_type_formatting.await?;
-            editor.update(&mut cx, |editor, cx| {
-                editor.refresh_document_highlights(cx);
-            })?;
+            if let Some(transaction) = on_type_formatting.await? {
+                if push_to_client_history {
+                    buffer.update(&mut cx, |buffer, _| {
+                        buffer.push_transaction(transaction, Instant::now());
+                    });
+                }
+                editor.update(&mut cx, |editor, cx| {
+                    editor.refresh_document_highlights(cx);
+                })?;
+            }
             Ok(())
         }))
     }