Remove remote worktrees and close their buffers when host unshares

server/src/rpc.rs

@@ -968,10 +968,12 @@ mod tests {
         editor::{Editor, EditorStyle, Insert},
         fs::{FakeFs, Fs as _},
         language::LanguageRegistry,
+        people_panel::JoinWorktree,
         rpc::{self, Client, Credentials, EstablishConnectionError},
         settings,
         test::FakeHttpClient,
         user::UserStore,
+        workspace::Workspace,
         worktree::Worktree,
     };
     use zrpc::Peer;
@@ -1087,6 +1089,94 @@ mod tests {
             .await;
     }
 
+    #[gpui::test]
+    async fn test_unshare_worktree(mut cx_a: TestAppContext, mut cx_b: TestAppContext) {
+        cx_b.update(zed::workspace::init);
+        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, user_store_b) = server.create_client(&mut cx_b, "user_b").await;
+        let app_state_b = zed::AppState {
+            rpc: client_b,
+            user_store: user_store_b,
+            ..Arc::try_unwrap(cx_b.update(zed::test::test_app_state))
+                .ok()
+                .unwrap()
+        };
+
+        cx_a.foreground().forbid_parking();
+
+        // Share a local worktree as client A
+        let fs = Arc::new(FakeFs::new());
+        fs.insert_tree(
+            "/a",
+            json!({
+                ".zed.toml": r#"collaborators = ["user_b"]"#,
+                "a.txt": "a-contents",
+                "b.txt": "b-contents",
+            }),
+        )
+        .await;
+        let worktree_a = Worktree::open_local(
+            client_a.clone(),
+            "/a".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 remote_worktree_id = worktree_a
+            .update(&mut cx_a, |tree, cx| tree.as_local_mut().unwrap().share(cx))
+            .await
+            .unwrap();
+
+        let (window_b, workspace_b) = cx_b.add_window(|cx| Workspace::new(&app_state_b, cx));
+        cx_b.update(|cx| {
+            cx.dispatch_action(
+                window_b,
+                vec![workspace_b.id()],
+                &JoinWorktree(remote_worktree_id),
+            );
+        });
+        workspace_b
+            .condition(&cx_b, |workspace, _| workspace.worktrees().len() == 1)
+            .await;
+
+        let local_worktree_id_b = workspace_b.read_with(&cx_b, |workspace, cx| {
+            let active_pane = workspace.active_pane().read(cx);
+            assert!(active_pane.active_item().is_none());
+            workspace.worktrees().iter().next().unwrap().id()
+        });
+        workspace_b
+            .update(&mut cx_b, |worktree, cx| {
+                worktree.open_entry((local_worktree_id_b, Path::new("a.txt").into()), cx)
+            })
+            .unwrap()
+            .await;
+        workspace_b.read_with(&cx_b, |workspace, cx| {
+            let active_pane = workspace.active_pane().read(cx);
+            assert!(active_pane.active_item().is_some());
+        });
+
+        worktree_a.update(&mut cx_a, |tree, cx| {
+            tree.as_local_mut().unwrap().unshare(cx);
+        });
+        workspace_b
+            .condition(&cx_b, |workspace, _| workspace.worktrees().len() == 0)
+            .await;
+        workspace_b.read_with(&cx_b, |workspace, cx| {
+            let active_pane = workspace.active_pane().read(cx);
+            assert!(active_pane.active_item().is_none());
+        });
+    }
+
     #[gpui::test]
     async fn test_propagate_saves_and_fs_changes_in_shared_worktree(
         mut cx_a: TestAppContext,

zed/src/workspace.rs

@@ -13,7 +13,7 @@ use crate::{
     settings::Settings,
     user,
     util::TryFutureExt as _,
-    worktree::{File, Worktree},
+    worktree::{self, File, Worktree},
     AppState, Authenticate,
 };
 use anyhow::Result;
@@ -886,9 +886,27 @@ impl Workspace {
                 rpc.authenticate_and_connect(&cx).await?;
                 let worktree =
                     Worktree::open_remote(rpc.clone(), worktree_id, languages, &mut cx).await?;
-                this.update(&mut cx, |workspace, cx| {
+                this.update(&mut cx, |this, cx| {
                     cx.observe(&worktree, |_, _, cx| cx.notify()).detach();
-                    workspace.worktrees.insert(worktree);
+                    cx.subscribe(&worktree, move |this, _, event, cx| match event {
+                        worktree::Event::Closed => {
+                            this.worktrees.retain(|worktree| {
+                                worktree.update(cx, |worktree, cx| {
+                                    if let Some(worktree) = worktree.as_remote_mut() {
+                                        if worktree.remote_id() == worktree_id {
+                                            worktree.close_all_buffers(cx);
+                                            return false;
+                                        }
+                                    }
+                                    true
+                                })
+                            });
+
+                            cx.notify();
+                        }
+                    })
+                    .detach();
+                    this.worktrees.insert(worktree);
                     cx.notify();
                 });
 

zed/src/worktree.rs

@@ -62,8 +62,12 @@ pub enum Worktree {
     Remote(RemoteWorktree),
 }
 
+pub enum Event {
+    Closed,
+}
+
 impl Entity for Worktree {
-    type Event = ();
+    type Event = Event;
 
     fn release(&mut self, cx: &mut MutableAppContext) {
         match self {
@@ -223,6 +227,7 @@ impl Worktree {
                     rpc.subscribe_to_entity(remote_id, cx, Self::handle_update),
                     rpc.subscribe_to_entity(remote_id, cx, Self::handle_update_buffer),
                     rpc.subscribe_to_entity(remote_id, cx, Self::handle_buffer_saved),
+                    rpc.subscribe_to_entity(remote_id, cx, Self::handle_unshare),
                 ];
 
                 Worktree::Remote(RemoteWorktree {
@@ -522,6 +527,16 @@ impl Worktree {
         Ok(())
     }
 
+    pub fn handle_unshare(
+        &mut self,
+        _: TypedEnvelope<proto::UnshareWorktree>,
+        _: Arc<rpc::Client>,
+        cx: &mut ModelContext<Self>,
+    ) -> Result<()> {
+        cx.emit(Event::Closed);
+        Ok(())
+    }
+
     fn poll_snapshot(&mut self, cx: &mut ModelContext<Self>) {
         match self {
             Self::Local(worktree) => {