Allow leaving worktree while opening a buffer

server/src/rpc.rs

@@ -1378,6 +1378,69 @@ mod tests {
         buffer_b.condition(&cx_b, |buf, _| buf.text() == text).await;
     }
 
+    #[gpui::test]
+    async fn test_leaving_worktree_while_opening_buffer(
+        mut cx_a: TestAppContext,
+        mut cx_b: TestAppContext,
+    ) {
+        cx_a.foreground().forbid_parking();
+        let lang_registry = Arc::new(LanguageRegistry::new());
+
+        // Connect to a server as 2 clients.
+        let mut server = TestServer::start().await;
+        let (client_a, _) = server.create_client(&mut cx_a, "user_a").await;
+        let (client_b, _) = server.create_client(&mut cx_b, "user_b").await;
+
+        // Share a local worktree as client A
+        let fs = Arc::new(FakeFs::new());
+        fs.insert_tree(
+            "/dir",
+            json!({
+                ".zed.toml": r#"collaborators = ["user_b"]"#,
+                "a.txt": "a-contents",
+            }),
+        )
+        .await;
+        let worktree_a = Worktree::open_local(
+            client_a.clone(),
+            "/dir".as_ref(),
+            fs,
+            lang_registry.clone(),
+            &mut cx_a.to_async(),
+        )
+        .await
+        .unwrap();
+        worktree_a
+            .read_with(&cx_a, |tree, _| tree.as_local().unwrap().scan_complete())
+            .await;
+        let worktree_id = worktree_a
+            .update(&mut cx_a, |tree, cx| tree.as_local_mut().unwrap().share(cx))
+            .await
+            .unwrap();
+
+        // Join that worktree as client B, and see that a guest has joined as client A.
+        let worktree_b = Worktree::open_remote(
+            client_b.clone(),
+            worktree_id,
+            lang_registry.clone(),
+            &mut cx_b.to_async(),
+        )
+        .await
+        .unwrap();
+        worktree_a
+            .condition(&cx_a, |tree, _| tree.peers().len() == 1)
+            .await;
+
+        let buffer_b = cx_b
+            .background()
+            .spawn(worktree_b.update(&mut cx_b, |worktree, cx| worktree.open_buffer("a.txt", cx)));
+        cx_b.update(|_| drop(worktree_b));
+        drop(buffer_b);
+        worktree_a
+            .condition(&cx_a, |tree, _| tree.peers().len() == 0)
+            .await;
+    }
+
     #[gpui::test]
     async fn test_peer_disconnection(mut cx_a: TestAppContext, cx_b: TestAppContext) {
         cx_a.foreground().forbid_parking();

zed/src/worktree.rs

@@ -1196,12 +1196,11 @@ impl RemoteWorktree {
         path: &Path,
         cx: &mut ModelContext<Worktree>,
     ) -> Task<Result<ModelHandle<Buffer>>> {
-        let handle = cx.handle();
         let mut existing_buffer = None;
         self.open_buffers.retain(|_buffer_id, buffer| {
             if let Some(buffer) = buffer.upgrade(cx.as_ref()) {
                 if let Some(file) = buffer.read(cx.as_ref()).file() {
-                    if file.worktree_id() == handle.id() && file.path.as_ref() == path {
+                    if file.worktree_id() == cx.model_id() && file.path.as_ref() == path {
                         existing_buffer = Some(buffer);
                     }
                 }
@@ -1215,21 +1214,27 @@ impl RemoteWorktree {
         let replica_id = self.replica_id;
         let remote_worktree_id = self.remote_id;
         let path = path.to_string_lossy().to_string();
-        cx.spawn(|this, mut cx| async move {
+        cx.spawn_weak(|this, mut cx| async move {
             if let Some(existing_buffer) = existing_buffer {
                 Ok(existing_buffer)
             } else {
                 let entry = this
+                    .upgrade(&cx)
+                    .ok_or_else(|| anyhow!("worktree was closed"))?
                     .read_with(&cx, |tree, _| tree.entry_for_path(&path).cloned())
                     .ok_or_else(|| anyhow!("file does not exist"))?;
-                let file = File::new(entry.id, handle, entry.path, entry.mtime);
-                let language = cx.read(|cx| file.select_language(cx));
                 let response = rpc
                     .request(proto::OpenBuffer {
                         worktree_id: remote_worktree_id as u64,
                         path,
                     })
                     .await?;
+
+                let this = this
+                    .upgrade(&cx)
+                    .ok_or_else(|| anyhow!("worktree was closed"))?;
+                let file = File::new(entry.id, this.clone(), entry.path, entry.mtime);
+                let language = cx.read(|cx| file.select_language(cx));
                 let remote_buffer = response.buffer.ok_or_else(|| anyhow!("empty buffer"))?;
                 let buffer_id = remote_buffer.id as usize;
                 let buffer = cx.add_model(|cx| {