ssh: Fix abs paths in file history & repeated go-to-def (#19027)

Thorsten Ball and Bennet created

This fixes two things:

- Go-to-def to absolute paths (i.e. opening stdlib files) multiple times
(opening, dropping, and re-opening worktrees)
- Re-opening abs paths from the file picker history that were added
there by go-to-def


Release Notes:

- N/A

---------

Co-authored-by: Bennet <bennet@zed.dev>

Change summary

crates/file_finder/src/file_finder.rs        | 63 ++++++++-------
crates/project/src/project.rs                | 84 ++++++++++++++-------
crates/project/src/worktree_store.rs         | 10 +-
crates/proto/proto/zed.proto                 |  9 ++
crates/proto/src/proto.rs                    |  4 
crates/remote_server/src/headless_project.rs | 47 +++++++++++-
6 files changed, 144 insertions(+), 73 deletions(-)

Detailed changes

crates/file_finder/src/file_finder.rs 🔗

@@ -758,11 +758,8 @@ impl FileFinderDelegate {
         cx: &mut ViewContext<'_, Picker<Self>>,
     ) -> Task<()> {
         cx.spawn(|picker, mut cx| async move {
-            let Some((project, fs)) = picker
-                .update(&mut cx, |picker, cx| {
-                    let fs = Arc::clone(&picker.delegate.project.read(cx).fs());
-                    (picker.delegate.project.clone(), fs)
-                })
+            let Some(project) = picker
+                .update(&mut cx, |picker, _| picker.delegate.project.clone())
                 .log_err()
             else {
                 return;
@@ -770,31 +767,36 @@ impl FileFinderDelegate {
 
             let query_path = Path::new(query.path_query());
             let mut path_matches = Vec::new();
-            match fs.metadata(query_path).await.log_err() {
-                Some(Some(_metadata)) => {
-                    let update_result = project
-                        .update(&mut cx, |project, cx| {
-                            if let Some((worktree, relative_path)) =
-                                project.find_worktree(query_path, cx)
-                            {
-                                path_matches.push(ProjectPanelOrdMatch(PathMatch {
-                                    score: 1.0,
-                                    positions: Vec::new(),
-                                    worktree_id: worktree.read(cx).id().to_usize(),
-                                    path: Arc::from(relative_path),
-                                    path_prefix: "".into(),
-                                    is_dir: false, // File finder doesn't support directories
-                                    distance_to_relative_ancestor: usize::MAX,
-                                }));
-                            }
-                        })
-                        .log_err();
-                    if update_result.is_none() {
-                        return;
-                    }
+
+            let abs_file_exists = if let Ok(task) = project.update(&mut cx, |this, cx| {
+                this.abs_file_path_exists(query.path_query(), cx)
+            }) {
+                task.await
+            } else {
+                false
+            };
+
+            if abs_file_exists {
+                let update_result = project
+                    .update(&mut cx, |project, cx| {
+                        if let Some((worktree, relative_path)) =
+                            project.find_worktree(query_path, cx)
+                        {
+                            path_matches.push(ProjectPanelOrdMatch(PathMatch {
+                                score: 1.0,
+                                positions: Vec::new(),
+                                worktree_id: worktree.read(cx).id().to_usize(),
+                                path: Arc::from(relative_path),
+                                path_prefix: "".into(),
+                                is_dir: false, // File finder doesn't support directories
+                                distance_to_relative_ancestor: usize::MAX,
+                            }));
+                        }
+                    })
+                    .log_err();
+                if update_result.is_none() {
+                    return;
                 }
-                Some(None) => {}
-                None => return,
             }
 
             picker
@@ -888,7 +890,8 @@ impl PickerDelegate for FileFinderDelegate {
                         project
                             .worktree_for_id(history_item.project.worktree_id, cx)
                             .is_some()
-                            || (project.is_local() && history_item.absolute.is_some())
+                            || ((project.is_local() || project.is_via_ssh())
+                                && history_item.absolute.is_some())
                     }),
                     self.currently_opened_path.as_ref(),
                     None,

crates/project/src/project.rs 🔗

@@ -2242,6 +2242,15 @@ impl Project {
             return;
         }
 
+        if let Some(ssh) = &self.ssh_client {
+            ssh.read(cx)
+                .to_proto_client()
+                .send(proto::RemoveWorktree {
+                    worktree_id: id_to_remove.to_proto(),
+                })
+                .log_err();
+        }
+
         cx.notify();
     }
 
@@ -3070,7 +3079,7 @@ impl Project {
         }
     }
 
-    // Returns the resolved version of `path`, that was found in `buffer`, if it exists.
+    /// Returns the resolved version of `path`, that was found in `buffer`, if it exists.
     pub fn resolve_existing_file_path(
         &self,
         path: &str,
@@ -3079,38 +3088,53 @@ impl Project {
     ) -> Task<Option<ResolvedPath>> {
         let path_buf = PathBuf::from(path);
         if path_buf.is_absolute() || path.starts_with("~") {
-            if self.is_local() {
-                let expanded = PathBuf::from(shellexpand::tilde(&path).into_owned());
+            self.resolve_abs_file_path(path, cx)
+        } else {
+            self.resolve_path_in_worktrees(path_buf, buffer, cx)
+        }
+    }
 
-                let fs = self.fs.clone();
-                cx.background_executor().spawn(async move {
-                    let path = expanded.as_path();
-                    let exists = fs.is_file(path).await;
+    pub fn abs_file_path_exists(&self, path: &str, cx: &mut ModelContext<Self>) -> Task<bool> {
+        let resolve_task = self.resolve_abs_file_path(path, cx);
+        cx.background_executor().spawn(async move {
+            let resolved_path = resolve_task.await;
+            resolved_path.is_some()
+        })
+    }
 
-                    exists.then(|| ResolvedPath::AbsPath(expanded))
-                })
-            } else if let Some(ssh_client) = self.ssh_client.as_ref() {
-                let request =
-                    ssh_client
-                        .read(cx)
-                        .to_proto_client()
-                        .request(proto::CheckFileExists {
-                            project_id: SSH_PROJECT_ID,
-                            path: path.to_string(),
-                        });
-                cx.background_executor().spawn(async move {
-                    let response = request.await.log_err()?;
-                    if response.exists {
-                        Some(ResolvedPath::AbsPath(PathBuf::from(response.path)))
-                    } else {
-                        None
-                    }
-                })
-            } else {
-                return Task::ready(None);
-            }
+    fn resolve_abs_file_path(
+        &self,
+        path: &str,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Option<ResolvedPath>> {
+        if self.is_local() {
+            let expanded = PathBuf::from(shellexpand::tilde(&path).into_owned());
+
+            let fs = self.fs.clone();
+            cx.background_executor().spawn(async move {
+                let path = expanded.as_path();
+                let exists = fs.is_file(path).await;
+
+                exists.then(|| ResolvedPath::AbsPath(expanded))
+            })
+        } else if let Some(ssh_client) = self.ssh_client.as_ref() {
+            let request = ssh_client
+                .read(cx)
+                .to_proto_client()
+                .request(proto::CheckFileExists {
+                    project_id: SSH_PROJECT_ID,
+                    path: path.to_string(),
+                });
+            cx.background_executor().spawn(async move {
+                let response = request.await.log_err()?;
+                if response.exists {
+                    Some(ResolvedPath::AbsPath(PathBuf::from(response.path)))
+                } else {
+                    None
+                }
+            })
         } else {
-            self.resolve_path_in_worktrees(path_buf, buffer, cx)
+            return Task::ready(None);
         }
     }
 

crates/project/src/worktree_store.rs 🔗

@@ -237,10 +237,13 @@ impl WorktreeStore {
             .to_string_lossy()
             .to_string();
         cx.spawn(|this, mut cx| async move {
+            let this = this.upgrade().context("Dropped worktree store")?;
+
             let response = client
                 .request(proto::AddWorktree {
                     project_id: SSH_PROJECT_ID,
                     path: abs_path.clone(),
+                    visible,
                 })
                 .await?;
 
@@ -252,7 +255,7 @@ impl WorktreeStore {
 
             let worktree = cx.update(|cx| {
                 Worktree::remote(
-                    0,
+                    SSH_PROJECT_ID,
                     0,
                     proto::WorktreeMetadata {
                         id: response.worktree_id,
@@ -509,11 +512,6 @@ impl WorktreeStore {
         for worktree in &self.worktrees {
             if let Some(worktree) = worktree.upgrade() {
                 worktree.update(cx, |worktree, _| {
-                    println!(
-                        "worktree. is_local: {:?}, is_remote: {:?}",
-                        worktree.is_local(),
-                        worktree.is_remote()
-                    );
                     if let Some(worktree) = worktree.as_remote_mut() {
                         worktree.disconnected_from_host();
                     }

crates/proto/proto/zed.proto 🔗

@@ -282,7 +282,9 @@ message Envelope {
         CheckFileExists check_file_exists = 255;
         CheckFileExistsResponse check_file_exists_response = 256;
 
-        ShutdownRemoteServer shutdown_remote_server = 257; // current max
+        ShutdownRemoteServer shutdown_remote_server = 257;
+
+        RemoveWorktree remove_worktree = 258; // current max
     }
 
     reserved 87 to 88;
@@ -2432,6 +2434,7 @@ message GetLlmTokenResponse {
 message AddWorktree {
     uint64 project_id = 2;
     string path = 1;
+    bool visible = 3;
 }
 
 message AddWorktreeResponse {
@@ -2460,3 +2463,7 @@ message CheckFileExistsResponse {
 }
 
 message ShutdownRemoteServer {}
+
+message RemoveWorktree {
+    uint64 worktree_id = 1;
+}

crates/proto/src/proto.rs 🔗

@@ -364,6 +364,7 @@ messages!(
     (CheckFileExists, Background),
     (CheckFileExistsResponse, Background),
     (ShutdownRemoteServer, Foreground),
+    (RemoveWorktree, Foreground),
 );
 
 request_messages!(
@@ -486,7 +487,8 @@ request_messages!(
     (LspExtSwitchSourceHeader, LspExtSwitchSourceHeaderResponse),
     (AddWorktree, AddWorktreeResponse),
     (CheckFileExists, CheckFileExistsResponse),
-    (ShutdownRemoteServer, Ack)
+    (ShutdownRemoteServer, Ack),
+    (RemoveWorktree, Ack)
 );
 
 entity_messages!(

crates/remote_server/src/headless_project.rs 🔗

@@ -135,6 +135,8 @@ impl HeadlessProject {
         client.add_request_handler(cx.weak_model(), Self::handle_ping);
 
         client.add_model_request_handler(Self::handle_add_worktree);
+        client.add_request_handler(cx.weak_model(), Self::handle_remove_worktree);
+
         client.add_model_request_handler(Self::handle_open_buffer_by_path);
         client.add_model_request_handler(Self::handle_find_search_candidates);
 
@@ -238,7 +240,7 @@ impl HeadlessProject {
             .update(&mut cx.clone(), |this, _| {
                 Worktree::local(
                     Arc::from(canonicalized),
-                    true,
+                    message.payload.visible,
                     this.fs.clone(),
                     this.next_entry_id.clone(),
                     &mut cx,
@@ -246,14 +248,49 @@ impl HeadlessProject {
             })?
             .await?;
 
-        this.update(&mut cx, |this, cx| {
-            this.worktree_store.update(cx, |worktree_store, cx| {
-                worktree_store.add(&worktree, cx);
-            });
+        let response = this.update(&mut cx, |_, cx| {
             worktree.update(cx, |worktree, _| proto::AddWorktreeResponse {
                 worktree_id: worktree.id().to_proto(),
             })
+        })?;
+
+        // We spawn this asynchronously, so that we can send the response back
+        // *before* `worktree_store.add()` can send out UpdateProject requests
+        // to the client about the new worktree.
+        //
+        // That lets the client manage the reference/handles of the newly-added
+        // worktree, before getting interrupted by an UpdateProject request.
+        //
+        // This fixes the problem of the client sending the AddWorktree request,
+        // headless project sending out a project update, client receiving it
+        // and immediately dropping the reference of the new client, causing it
+        // to be dropped on the headless project, and the client only then
+        // receiving a response to AddWorktree.
+        cx.spawn(|mut cx| async move {
+            this.update(&mut cx, |this, cx| {
+                this.worktree_store.update(cx, |worktree_store, cx| {
+                    worktree_store.add(&worktree, cx);
+                });
+            })
+            .log_err();
         })
+        .detach();
+
+        Ok(response)
+    }
+
+    pub async fn handle_remove_worktree(
+        this: Model<Self>,
+        envelope: TypedEnvelope<proto::RemoveWorktree>,
+        mut cx: AsyncAppContext,
+    ) -> Result<proto::Ack> {
+        let worktree_id = WorktreeId::from_proto(envelope.payload.worktree_id);
+        this.update(&mut cx, |this, cx| {
+            this.worktree_store.update(cx, |worktree_store, cx| {
+                worktree_store.remove_worktree(worktree_id, cx);
+            });
+        })?;
+        Ok(proto::Ack {})
     }
 
     pub async fn handle_open_buffer_by_path(