Send trust messages only to the clients related to the current WorktreeStore (#47005)

Release Notes:

- N/A

crates/collab/src/tests/editor_tests.rs

@@ -4699,10 +4699,10 @@ async fn test_remote_project_worktree_trust(cx_a: &mut TestAppContext, cx_b: &mu
     };
 
     cx_a.update(|cx| {
-        project::trusted_worktrees::init(HashMap::default(), None, None, cx);
+        project::trusted_worktrees::init(HashMap::default(), cx);
     });
     cx_b.update(|cx| {
-        project::trusted_worktrees::init(HashMap::default(), None, None, cx);
+        project::trusted_worktrees::init(HashMap::default(), cx);
     });
 
     let mut server = TestServer::start(cx_a.executor()).await;

crates/collab/src/tests/remote_editing_collaboration_tests.rs

@@ -857,11 +857,11 @@ async fn test_slow_adapter_startup_retries(
 async fn test_ssh_remote_worktree_trust(cx_a: &mut TestAppContext, server_cx: &mut TestAppContext) {
     cx_a.update(|cx| {
         release_channel::init(semver::Version::new(0, 0, 0), cx);
-        project::trusted_worktrees::init(HashMap::default(), None, None, cx);
+        project::trusted_worktrees::init(HashMap::default(), cx);
     });
     server_cx.update(|cx| {
         release_channel::init(semver::Version::new(0, 0, 0), cx);
-        project::trusted_worktrees::init(HashMap::default(), None, None, cx);
+        project::trusted_worktrees::init(HashMap::default(), cx);
     });
 
     let mut server = TestServer::start(cx_a.executor().clone()).await;

crates/editor/src/editor_tests.rs

@@ -30174,7 +30174,7 @@ async fn test_tab_list_indent(cx: &mut TestAppContext) {
 #[gpui::test]
 async fn test_local_worktree_trust(cx: &mut TestAppContext) {
     init_test(cx, |_| {});
-    cx.update(|cx| project::trusted_worktrees::init(HashMap::default(), None, None, cx));
+    cx.update(|cx| project::trusted_worktrees::init(HashMap::default(), cx));
 
     cx.update(|cx| {
         SettingsStore::update_global(cx, |store, cx| {

crates/project/src/trusted_worktrees.rs

@@ -55,16 +55,9 @@ use util::debug_panic;
 
 use crate::{project_settings::ProjectSettings, worktree_store::WorktreeStore};
 
-pub fn init(
-    db_trusted_paths: DbTrustedPaths,
-    downstream_client: Option<(AnyProtoClient, ProjectId)>,
-    upstream_client: Option<(AnyProtoClient, ProjectId)>,
-    cx: &mut App,
-) {
+pub fn init(db_trusted_paths: DbTrustedPaths, cx: &mut App) {
     if TrustedWorktrees::try_get_global(cx).is_none() {
-        let trusted_worktrees = cx.new(|_| {
-            TrustedWorktreesStore::new(db_trusted_paths, downstream_client, upstream_client)
-        });
+        let trusted_worktrees = cx.new(|_| TrustedWorktreesStore::new(db_trusted_paths));
         cx.set_global(TrustedWorktrees(trusted_worktrees))
     }
 }
@@ -80,21 +73,21 @@ pub fn track_worktree_trust(
     match TrustedWorktrees::try_get_global(cx) {
         Some(trusted_worktrees) => {
             trusted_worktrees.update(cx, |trusted_worktrees, cx| {
-                if let Some(downstream_client) = downstream_client {
-                    trusted_worktrees.downstream_clients.push(downstream_client);
-                }
-                if let Some(upstream_client) = upstream_client.clone() {
-                    trusted_worktrees.upstream_clients.push(upstream_client);
-                }
-                trusted_worktrees.add_worktree_store(worktree_store, remote_host, cx);
+                trusted_worktrees.add_worktree_store(
+                    worktree_store.clone(),
+                    remote_host,
+                    downstream_client,
+                    upstream_client.clone(),
+                    cx,
+                );
 
                 if let Some((upstream_client, upstream_project_id)) = upstream_client {
                     let trusted_paths = trusted_worktrees
                         .trusted_paths
-                        .iter()
-                        .flat_map(|(_, paths)| {
-                            paths.iter().map(|trusted_path| trusted_path.to_proto())
-                        })
+                        .get(&worktree_store.downgrade())
+                        .into_iter()
+                        .flatten()
+                        .map(|trusted_path| trusted_path.to_proto())
                         .collect::<Vec<_>>();
                     if !trusted_paths.is_empty() {
                         upstream_client
@@ -129,15 +122,20 @@ impl TrustedWorktrees {
 /// or a certain worktree had been trusted.
 #[derive(Debug)]
 pub struct TrustedWorktreesStore {
-    downstream_clients: Vec<(AnyProtoClient, ProjectId)>,
-    upstream_clients: Vec<(AnyProtoClient, ProjectId)>,
-    worktree_stores: HashMap<WeakEntity<WorktreeStore>, Option<RemoteHostLocation>>,
+    worktree_stores: HashMap<WeakEntity<WorktreeStore>, StoreData>,
     db_trusted_paths: DbTrustedPaths,
     trusted_paths: TrustedPaths,
     restricted: HashMap<WeakEntity<WorktreeStore>, HashSet<WorktreeId>>,
     worktree_trust_serialization: Task<()>,
 }
 
+#[derive(Debug, Default)]
+struct StoreData {
+    upstream_client: Option<(AnyProtoClient, ProjectId)>,
+    downstream_client: Option<(AnyProtoClient, ProjectId)>,
+    host: Option<RemoteHostLocation>,
+}
+
 /// An identifier of a host to split the trust questions by.
 /// Each trusted data change and event is done for a particular host.
 /// A host may contain more than one worktree or even project open concurrently.
@@ -226,36 +224,9 @@ type TrustedPaths = HashMap<WeakEntity<WorktreeStore>, HashSet<PathTrust>>;
 pub type DbTrustedPaths = HashMap<Option<RemoteHostLocation>, HashSet<PathBuf>>;
 
 impl TrustedWorktreesStore {
-    fn new(
-        db_trusted_paths: DbTrustedPaths,
-        downstream_client: Option<(AnyProtoClient, ProjectId)>,
-        upstream_client: Option<(AnyProtoClient, ProjectId)>,
-    ) -> Self {
-        if let Some((upstream_client, upstream_project_id)) = &upstream_client {
-            let trusted_paths = db_trusted_paths
-                .iter()
-                .flat_map(|(_, paths)| {
-                    paths
-                        .iter()
-                        .cloned()
-                        .map(PathTrust::AbsPath)
-                        .map(|trusted_path| trusted_path.to_proto())
-                })
-                .collect::<Vec<_>>();
-            if !trusted_paths.is_empty() {
-                upstream_client
-                    .send(proto::TrustWorktrees {
-                        project_id: upstream_project_id.0,
-                        trusted_paths,
-                    })
-                    .ok();
-            }
-        }
-
+    fn new(db_trusted_paths: DbTrustedPaths) -> Self {
         Self {
             db_trusted_paths,
-            downstream_clients: downstream_client.into_iter().collect(),
-            upstream_clients: upstream_client.into_iter().collect(),
             trusted_paths: HashMap::default(),
             worktree_stores: HashMap::default(),
             restricted: HashMap::default(),
@@ -406,25 +377,26 @@ impl TrustedWorktreesStore {
             );
         }
 
-        cx.emit(TrustedWorktreesEvent::Trusted(
-            weak_worktree_store,
-            trusted_paths.clone(),
-        ));
-
-        for (upstream_client, upstream_project_id) in &self.upstream_clients {
-            let trusted_paths = trusted_paths
-                .iter()
-                .map(|trusted_path| trusted_path.to_proto())
-                .collect::<Vec<_>>();
-            if !trusted_paths.is_empty() {
-                upstream_client
-                    .send(proto::TrustWorktrees {
-                        project_id: upstream_project_id.0,
-                        trusted_paths,
-                    })
-                    .ok();
+        if let Some(store_data) = self.worktree_stores.get(&weak_worktree_store) {
+            if let Some((upstream_client, upstream_project_id)) = &store_data.upstream_client {
+                let trusted_paths = trusted_paths
+                    .iter()
+                    .map(|trusted_path| trusted_path.to_proto())
+                    .collect::<Vec<_>>();
+                if !trusted_paths.is_empty() {
+                    upstream_client
+                        .send(proto::TrustWorktrees {
+                            project_id: upstream_project_id.0,
+                            trusted_paths,
+                        })
+                        .ok();
+                }
             }
         }
+        cx.emit(TrustedWorktreesEvent::Trusted(
+            weak_worktree_store,
+            trusted_paths,
+        ));
     }
 
     /// Restricts certain entities on this host.
@@ -533,26 +505,29 @@ impl TrustedWorktreesStore {
             .or_default()
             .insert(worktree_id);
         log::info!("Worktree {worktree_path:?} is not trusted");
+        if let Some(store_data) = self.worktree_stores.get(&weak_worktree_store) {
+            if let Some((downstream_client, downstream_project_id)) = &store_data.downstream_client
+            {
+                downstream_client
+                    .send(proto::RestrictWorktrees {
+                        project_id: downstream_project_id.0,
+                        worktree_ids: vec![worktree_id.to_proto()],
+                    })
+                    .ok();
+            }
+            if let Some((upstream_client, upstream_project_id)) = &store_data.upstream_client {
+                upstream_client
+                    .send(proto::RestrictWorktrees {
+                        project_id: upstream_project_id.0,
+                        worktree_ids: vec![worktree_id.to_proto()],
+                    })
+                    .ok();
+            }
+        }
         cx.emit(TrustedWorktreesEvent::Restricted(
             weak_worktree_store,
             HashSet::from_iter([PathTrust::Worktree(worktree_id)]),
         ));
-        for (downstream_client, downstream_project_id) in &self.downstream_clients {
-            downstream_client
-                .send(proto::RestrictWorktrees {
-                    project_id: downstream_project_id.0,
-                    worktree_ids: vec![worktree_id.to_proto()],
-                })
-                .ok();
-        }
-        for (upstream_client, upstream_project_id) in &self.upstream_clients {
-            upstream_client
-                .send(proto::RestrictWorktrees {
-                    project_id: upstream_project_id.0,
-                    worktree_ids: vec![worktree_id.to_proto()],
-                })
-                .ok();
-        }
         false
     }
 
@@ -625,7 +600,7 @@ impl TrustedWorktreesStore {
             .trusted_paths
             .iter()
             .filter_map(|(worktree_store, paths)| {
-                let host = self.worktree_stores.get(&worktree_store)?.clone();
+                let host = self.worktree_stores.get(&worktree_store)?.host.clone();
                 let abs_paths = paths
                     .iter()
                     .flat_map(|path| match path {
@@ -656,13 +631,21 @@ impl TrustedWorktreesStore {
         &mut self,
         worktree_store: Entity<WorktreeStore>,
         remote_host: Option<RemoteHostLocation>,
+        downstream_client: Option<(AnyProtoClient, ProjectId)>,
+        upstream_client: Option<(AnyProtoClient, ProjectId)>,
         cx: &mut Context<Self>,
     ) {
         self.worktree_stores
             .retain(|worktree_store, _| worktree_store.is_upgradable());
         let weak_worktree_store = worktree_store.downgrade();
-        self.worktree_stores
-            .insert(weak_worktree_store.clone(), remote_host.clone());
+        self.worktree_stores.insert(
+            weak_worktree_store.clone(),
+            StoreData {
+                host: remote_host.clone(),
+                downstream_client,
+                upstream_client,
+            },
+        );
 
         let mut new_trusted_paths = HashSet::default();
         if let Some(db_trusted_paths) = self.db_trusted_paths.get(&remote_host) {
@@ -734,7 +717,7 @@ mod tests {
         cx: &mut TestAppContext,
     ) -> Entity<TrustedWorktreesStore> {
         cx.update(|cx| {
-            init(HashMap::default(), None, None, cx);
+            init(HashMap::default(), cx);
             track_worktree_trust(worktree_store, None, None, None, cx);
             TrustedWorktrees::try_get_global(cx).expect("global should be set")
         })

crates/remote_server/src/unix.rs

@@ -1,7 +1,7 @@
 use crate::HeadlessProject;
 use crate::headless_project::HeadlessAppState;
 use anyhow::{Context as _, Result, anyhow};
-use client::{ProjectId, ProxySettings};
+use client::ProxySettings;
 use collections::HashMap;
 use project::trusted_worktrees;
 use util::ResultExt;
@@ -441,12 +441,7 @@ pub fn execute_run(
 
         log::info!("gpui app started, initializing server");
         let session = start_server(listeners, log_rx, cx, is_wsl_interop);
-        trusted_worktrees::init(
-            HashMap::default(),
-            Some((session.clone(), ProjectId(REMOTE_SERVER_PROJECT_ID))),
-            None,
-            cx,
-        );
+        trusted_worktrees::init(HashMap::default(), cx);
 
         GitHostingProviderRegistry::set_global(git_hosting_provider_registry, cx);
         git_hosting_providers::init(cx);

crates/zed/src/main.rs

@@ -423,7 +423,7 @@ fn main() {
                 HashMap::default()
             }
         };
-        trusted_worktrees::init(db_trusted_paths, None, None, cx);
+        trusted_worktrees::init(db_trusted_paths, cx);
         menu::init();
         zed_actions::init();