feat: support restoring on remote worktree

dino created

Change summary

crates/collab/src/rpc.rs             |   1 
crates/fs/src/fs.rs                  |   6 +
crates/project/src/project.rs        |  14 ++-
crates/project/src/worktree_store.rs |  16 ++++
crates/proto/proto/worktree.proto    |  11 ++
crates/proto/proto/zed.proto         |   4 
crates/proto/src/proto.rs            |   4 +
crates/worktree/src/worktree.rs      | 115 +++++++++++++++++++++--------
8 files changed, 131 insertions(+), 40 deletions(-)

Detailed changes

crates/collab/src/rpc.rs 🔗

@@ -348,6 +348,7 @@ impl Server {
             .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>)
             .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>)
             .add_request_handler(forward_mutating_project_request::<proto::TrashProjectEntry>)
+            .add_request_handler(forward_mutating_project_request::<proto::RestoreProjectEntry>)
             .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>)
             .add_request_handler(
                 forward_mutating_project_request::<proto::ExpandAllForProjectEntry>,

crates/fs/src/fs.rs 🔗

@@ -398,6 +398,12 @@ impl From<MTime> for proto::Timestamp {
 
 slotmap::new_key_type! { pub struct TrashId; }
 
+// TODO!: Should we convert these to `from_proto` and `to_proto` for the sake of
+// consistency with other `Id` types like:
+//
+// * `WorktreeId`
+// * `ChannelId`
+// * `ProjectId`
 impl TrashId {
     pub fn from_u64(value: u64) -> Self {
         KeyData::from_ffi(value).into()

crates/project/src/project.rs 🔗

@@ -2614,12 +2614,14 @@ impl Project {
         };
 
         cx.spawn(async move |_, cx| {
-            Worktree::restore_entry(trash_id, worktree, cx)
-                .await
-                .map(|rel_path_buf| ProjectPath {
-                    worktree_id: worktree_id,
-                    path: Arc::from(rel_path_buf.as_rel_path()),
-                })
+            let entry = worktree
+                .update(cx, |worktree, cx| worktree.restore_entry(trash_id, cx))
+                .await?;
+
+            Ok(ProjectPath {
+                worktree_id: worktree_id,
+                path: entry.path.clone(),
+            })
         })
     }
 

crates/project/src/worktree_store.rs 🔗

@@ -102,6 +102,7 @@ impl WorktreeStore {
         client.add_entity_request_handler(Self::handle_copy_project_entry);
         client.add_entity_request_handler(Self::handle_delete_project_entry);
         client.add_entity_request_handler(Self::handle_trash_project_entry);
+        client.add_entity_request_handler(Self::handle_restore_project_entry);
         client.add_entity_request_handler(Self::handle_expand_project_entry);
         client.add_entity_request_handler(Self::handle_expand_all_for_project_entry);
     }
@@ -1206,6 +1207,21 @@ impl WorktreeStore {
         Worktree::handle_delete_entry(worktree, envelope.payload, cx).await
     }
 
+    pub async fn handle_restore_project_entry(
+        this: Entity<Self>,
+        envelope: TypedEnvelope<proto::RestoreProjectEntry>,
+        mut cx: AsyncApp,
+    ) -> Result<proto::RestoreProjectEntryResponse> {
+        let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
+
+        let worktree = this.update(&mut cx, |this, cx| {
+            this.worktree_for_id(worktree_id, cx)
+                .context("worktree not found")
+        })?;
+
+        Worktree::handle_restore_entry(worktree, envelope.payload, cx).await
+    }
+
     pub async fn handle_rename_project_entry(
         this: Entity<Self>,
         request: proto::RenameProjectEntry,

crates/proto/proto/worktree.proto 🔗

@@ -125,6 +125,12 @@ message TrashProjectEntry {
   uint64 entry_id = 2;
 }
 
+message RestoreProjectEntry {
+  uint64 project_id = 1;
+  uint64 worktree_id = 2;
+  uint64 trash_id = 3;
+}
+
 message ExpandProjectEntry {
   uint64 project_id = 1;
   uint64 entry_id = 2;
@@ -154,6 +160,11 @@ message TrashProjectEntryResponse {
   uint64 worktree_scan_id = 2;
 }
 
+message RestoreProjectEntryResponse {
+  uint64 entry_id = 1;
+  uint64 worktree_scan_id = 2;
+}
+
 message UpdateWorktreeSettings {
   uint64 project_id = 1;
   uint64 worktree_id = 2;

crates/proto/proto/zed.proto 🔗

@@ -478,7 +478,9 @@ message Envelope {
     GitGetHeadSha git_get_head_sha = 440;
     GitGetHeadShaResponse git_get_head_sha_response = 441;
     TrashProjectEntry trash_project_entry = 442;
-    TrashProjectEntryResponse trash_project_entry_response = 443; // current max
+    TrashProjectEntryResponse trash_project_entry_response = 443;
+    RestoreProjectEntry restore_project_entry = 444;
+    RestoreProjectEntryResponse restore_project_entry_response = 445; // current max
   }
 
   reserved 87 to 88;

crates/proto/src/proto.rs 🔗

@@ -66,6 +66,8 @@ messages!(
     (DeleteProjectEntry, Foreground),
     (TrashProjectEntry, Foreground),
     (TrashProjectEntryResponse, Foreground),
+    (RestoreProjectEntry, Foreground),
+    (RestoreProjectEntryResponse, Foreground),
     (DownloadFileByPath, Background),
     (DownloadFileResponse, Background),
     (EndStream, Foreground),
@@ -390,6 +392,7 @@ request_messages!(
     (DeleteChannel, Ack),
     (DeleteProjectEntry, ProjectEntryResponse),
     (TrashProjectEntry, TrashProjectEntryResponse),
+    (RestoreProjectEntry, RestoreProjectEntryResponse),
     (DownloadFileByPath, DownloadFileResponse),
     (ExpandProjectEntry, ExpandProjectEntryResponse),
     (ExpandAllForProjectEntry, ExpandAllForProjectEntryResponse),
@@ -616,6 +619,7 @@ entity_messages!(
     GetFoldingRanges,
     DeleteProjectEntry,
     TrashProjectEntry,
+    RestoreProjectEntry,
     ExpandProjectEntry,
     ExpandAllForProjectEntry,
     FindSearchCandidates,

crates/worktree/src/worktree.rs 🔗

@@ -71,7 +71,7 @@ use text::{LineEnding, Rope};
 use util::{
     ResultExt, maybe,
     paths::{PathMatcher, PathStyle, SanitizedPath, home_dir},
-    rel_path::{RelPath, RelPathBuf},
+    rel_path::RelPath,
 };
 pub use worktree_settings::WorktreeSettings;
 
@@ -896,16 +896,14 @@ impl Worktree {
         Some(task)
     }
 
-    pub async fn restore_entry(
+    pub fn restore_entry(
+        &mut self,
         trash_id: TrashId,
-        worktree: Entity<Self>,
-        cx: &mut AsyncApp,
-    ) -> Result<RelPathBuf> {
-        let is_local = worktree.read_with(cx, |this, _| this.is_local());
-        if is_local {
-            LocalWorktree::restore_entry(trash_id, worktree, cx).await
-        } else {
-            RemoteWorktree::restore_entry(trash_id, worktree, cx).await
+        cx: &mut Context<'_, Worktree>,
+    ) -> Task<Result<Entry>> {
+        match self {
+            Worktree::Local(this) => this.restore_entry(trash_id, cx),
+            Worktree::Remote(this) => this.restore_entry(trash_id, cx),
         }
     }
 
@@ -1049,6 +1047,26 @@ impl Worktree {
         })
     }
 
+    pub async fn handle_restore_entry(
+        this: Entity<Self>,
+        request: proto::RestoreProjectEntry,
+        mut cx: AsyncApp,
+    ) -> Result<proto::RestoreProjectEntryResponse> {
+        let (scan_id, task) = this.update(&mut cx, |this, cx| {
+            (
+                this.scan_id(),
+                this.restore_entry(TrashId::from_u64(request.trash_id), cx),
+            )
+        });
+
+        let entry = task.await?;
+
+        Ok(proto::RestoreProjectEntryResponse {
+            entry_id: entry.id.to_proto(),
+            worktree_scan_id: scan_id as u64,
+        })
+    }
+
     pub async fn handle_expand_entry(
         this: Entity<Self>,
         request: proto::ExpandProjectEntry,
@@ -1784,32 +1802,31 @@ impl LocalWorktree {
         })
     }
 
-    pub async fn restore_entry(
-        trash_entry: TrashId,
-        this: Entity<Worktree>,
-        cx: &mut AsyncApp,
-    ) -> Result<RelPathBuf> {
-        let Some((fs, worktree_abs_path, path_style)) = this.read_with(cx, |this, _cx| {
-            let local_worktree = match this {
-                Worktree::Local(local_worktree) => local_worktree,
-                Worktree::Remote(_) => return None,
-            };
+    pub fn restore_entry(
+        &mut self,
+        trash_id: TrashId,
+        cx: &mut Context<'_, Worktree>,
+    ) -> Task<Result<Entry>> {
+        let fs = self.fs.clone();
+        let worktree_abs_path = self.abs_path().clone();
+        let path_style = self.path_style();
 
-            let fs = local_worktree.fs.clone();
-            let path_style = local_worktree.path_style();
-            Some((fs, Arc::clone(local_worktree.abs_path()), path_style))
-        }) else {
-            return Err(anyhow!("Localworktree should not change into a remote one"));
-        };
+        cx.spawn(async move |this, cx| {
+            let path_buf = fs.restore(trash_id).await?;
+            let path = path_buf
+                .strip_prefix(worktree_abs_path)
+                .context("Could not strip prefix")?;
+            let path = Arc::from(RelPath::new(&path, path_style)?.as_ref());
 
-        let path_buf = fs.restore(trash_entry).await?;
-        let path = path_buf
-            .strip_prefix(worktree_abs_path)
-            .context("Could not strip prefix")?;
-        let path = RelPath::new(&path, path_style)?;
-        let path = path.into_owned();
+            let entry = this
+                .update(cx, |this, cx| {
+                    this.as_local_mut().unwrap().refresh_entry(path, None, cx)
+                })?
+                .await?
+                .context("Entry not found after restore")?;
 
-        Ok(path)
+            Ok(entry)
+        })
     }
 
     pub fn copy_external_entries(
@@ -2230,6 +2247,38 @@ impl RemoteWorktree {
         })
     }
 
+    fn restore_entry(&mut self, trash_id: TrashId, cx: &Context<Worktree>) -> Task<Result<Entry>> {
+        let project_id = self.project_id();
+        let worktree_id = self.id().to_proto();
+        let trash_id = trash_id.to_u64();
+
+        let request = self.client.request(proto::RestoreProjectEntry {
+            project_id,
+            worktree_id,
+            trash_id,
+        });
+
+        cx.spawn(async move |this, cx| {
+            let response = request.await?;
+            let scan_id = response.worktree_scan_id as usize;
+            let entry_id = ProjectEntryId(response.entry_id as usize);
+
+            let (task, entry) = this.update(cx, |worktree, cx| {
+                // TODO!(dino): Remove `entry_for_id(entry_id).unwrap()` call,
+                // avoid unwrapping.
+                let remote_worktree = worktree.as_remote_mut().unwrap();
+                let entry = remote_worktree.entry_for_id(entry_id).unwrap().clone();
+                let proto_entry = proto::Entry::from(&entry);
+                let task = remote_worktree.insert_entry(proto_entry, scan_id, cx);
+
+                (task, entry)
+            })?;
+
+            task.await?;
+            Ok(entry)
+        })
+    }
+
     fn copy_external_entries(
         &self,
         target_directory: Arc<RelPath>,