From cbf5dd1f23c671d5e561d71d70c69412e7bbd997 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Thu, 7 Aug 2025 18:14:25 -0400
Subject: [PATCH] client: Only connect to Collab automatically for Zed staff
 (#35827)

This PR makes it so that only Zed staff connect to Collab automatically.

Anyone else can connect to Collab manually when they want to collaborate
(but this is not required for using Zed's LLM features).

Release Notes:

- N/A

---------

Co-authored-by: Richard <richard@zed.dev>
---
 crates/client/src/client.rs               | 45 ++++++++++++++++++-----
 crates/feature_flags/src/feature_flags.rs | 24 ++++++++++++
 2 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/crates/client/src/client.rs b/crates/client/src/client.rs
index 9d58692c0d..12ea4bcd3e 100644
--- a/crates/client/src/client.rs
+++ b/crates/client/src/client.rs
@@ -16,6 +16,7 @@ use clock::SystemClock;
 use cloud_api_client::CloudApiClient;
 use cloud_api_client::websocket_protocol::MessageToClient;
 use credentials_provider::CredentialsProvider;
+use feature_flags::FeatureFlagAppExt as _;
 use futures::{
     AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,
     channel::oneshot, future::BoxFuture,
@@ -964,25 +965,51 @@ impl Client {
         Ok(())
     }
 
-    /// Performs a sign-in and also connects to Collab.
+    /// Performs a sign-in and also (optionally) connects to Collab.
     ///
-    /// This is called in places where we *don't* need to connect in the future. We will replace these calls with calls
-    /// to `sign_in` when we're ready to remove auto-connection to Collab.
+    /// Only Zed staff automatically connect to Collab.
     pub async fn sign_in_with_optional_connect(
         self: &Arc<Self>,
         try_provider: bool,
         cx: &AsyncApp,
     ) -> Result<()> {
+        let (is_staff_tx, is_staff_rx) = oneshot::channel::<bool>();
+        let mut is_staff_tx = Some(is_staff_tx);
+        cx.update(|cx| {
+            cx.on_flags_ready(move |state, _cx| {
+                if let Some(is_staff_tx) = is_staff_tx.take() {
+                    is_staff_tx.send(state.is_staff).log_err();
+                }
+            })
+            .detach();
+        })
+        .log_err();
+
         let credentials = self.sign_in(try_provider, cx).await?;
 
         self.connect_to_cloud(cx).await.log_err();
 
-        let connect_result = match self.connect_with_credentials(credentials, cx).await {
-            ConnectionResult::Timeout => Err(anyhow!("connection timed out")),
-            ConnectionResult::ConnectionReset => Err(anyhow!("connection reset")),
-            ConnectionResult::Result(result) => result.context("client auth and connect"),
-        };
-        connect_result.log_err();
+        cx.update(move |cx| {
+            cx.spawn({
+                let client = self.clone();
+                async move |cx| {
+                    let is_staff = is_staff_rx.await?;
+                    if is_staff {
+                        match client.connect_with_credentials(credentials, cx).await {
+                            ConnectionResult::Timeout => Err(anyhow!("connection timed out")),
+                            ConnectionResult::ConnectionReset => Err(anyhow!("connection reset")),
+                            ConnectionResult::Result(result) => {
+                                result.context("client auth and connect")
+                            }
+                        }
+                    } else {
+                        Ok(())
+                    }
+                }
+            })
+            .detach_and_log_err(cx);
+        })
+        .log_err();
 
         Ok(())
     }
diff --git a/crates/feature_flags/src/feature_flags.rs b/crates/feature_flags/src/feature_flags.rs
index 631bafc841..ef357adf35 100644
--- a/crates/feature_flags/src/feature_flags.rs
+++ b/crates/feature_flags/src/feature_flags.rs
@@ -158,6 +158,11 @@ where
     }
 }
 
+#[derive(Debug)]
+pub struct OnFlagsReady {
+    pub is_staff: bool,
+}
+
 pub trait FeatureFlagAppExt {
     fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag;
 
@@ -169,6 +174,10 @@ pub trait FeatureFlagAppExt {
     fn has_flag<T: FeatureFlag>(&self) -> bool;
     fn is_staff(&self) -> bool;
 
+    fn on_flags_ready<F>(&mut self, callback: F) -> Subscription
+    where
+        F: FnMut(OnFlagsReady, &mut App) + 'static;
+
     fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
     where
         F: FnMut(bool, &mut App) + 'static;
@@ -198,6 +207,21 @@ impl FeatureFlagAppExt for App {
             .unwrap_or(false)
     }
 
+    fn on_flags_ready<F>(&mut self, mut callback: F) -> Subscription
+    where
+        F: FnMut(OnFlagsReady, &mut App) + 'static,
+    {
+        self.observe_global::<FeatureFlags>(move |cx| {
+            let feature_flags = cx.global::<FeatureFlags>();
+            callback(
+                OnFlagsReady {
+                    is_staff: feature_flags.staff,
+                },
+                cx,
+            );
+        })
+    }
+
     fn observe_flag<T: FeatureFlag, F>(&mut self, mut callback: F) -> Subscription
     where
         F: FnMut(bool, &mut App) + 'static,