Remove remote worktrees and close their buffers when host unshares

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) => {