Start separating authentication from connection to collab (#35471)

This pull request should be idempotent, but lays the groundwork for
avoiding to connect to collab in order to interact with AI features
provided by Zed.

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <git@maxdeviant.com>
Co-authored-by: Richard Feldman <oss@rtfeldman.com>

crates/zed/src/main.rs

@@ -5,7 +5,7 @@ use agent_ui::AgentPanel;
 use anyhow::{Context as _, Result};
 use clap::{Parser, command};
 use cli::FORCE_CLI_MODE_ENV_VAR_NAME;
-use client::{Client, CloudUserStore, ProxySettings, UserStore, parse_zed_link};
+use client::{Client, ProxySettings, UserStore, parse_zed_link};
 use collab_ui::channel_view::ChannelView;
 use collections::HashMap;
 use db::kvp::{GLOBAL_KEY_VALUE_STORE, KEY_VALUE_STORE};
@@ -42,7 +42,7 @@ use theme::{
     ActiveTheme, IconThemeNotFoundError, SystemAppearance, ThemeNotFoundError, ThemeRegistry,
     ThemeSettings,
 };
-use util::{ConnectionResult, ResultExt, TryFutureExt, maybe};
+use util::{ResultExt, TryFutureExt, maybe};
 use uuid::Uuid;
 use welcome::{FIRST_OPEN, show_welcome_view};
 use workspace::{
@@ -457,8 +457,6 @@ pub fn main() {
         language::init(cx);
         languages::init(languages.clone(), node_runtime.clone(), cx);
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
 
         language_extension::init(
@@ -518,7 +516,6 @@ pub fn main() {
             languages: languages.clone(),
             client: client.clone(),
             user_store: user_store.clone(),
-            cloud_user_store,
             fs: fs.clone(),
             build_window_options,
             workspace_store,
@@ -556,12 +553,7 @@ pub fn main() {
         );
         supermaven::init(app_state.client.clone(), cx);
         language_model::init(app_state.client.clone(), cx);
-        language_models::init(
-            app_state.user_store.clone(),
-            app_state.cloud_user_store.clone(),
-            app_state.client.clone(),
-            cx,
-        );
+        language_models::init(app_state.user_store.clone(), app_state.client.clone(), cx);
         agent_settings::init(cx);
         agent_servers::init(cx);
         web_search::init(cx);
@@ -569,7 +561,7 @@ pub fn main() {
         snippet_provider::init(cx);
         inline_completion_registry::init(
             app_state.client.clone(),
-            app_state.cloud_user_store.clone(),
+            app_state.user_store.clone(),
             cx,
         );
         let prompt_builder = PromptBuilder::load(app_state.fs.clone(), stdout_is_a_pty(), cx);
@@ -690,17 +682,9 @@ pub fn main() {
 
         cx.spawn({
             let client = app_state.client.clone();
-            async move |cx| match authenticate(client, &cx).await {
-                ConnectionResult::Timeout => log::error!("Timeout during initial auth"),
-                ConnectionResult::ConnectionReset => {
-                    log::error!("Connection reset during initial auth")
-                }
-                ConnectionResult::Result(r) => {
-                    r.log_err();
-                }
-            }
+            async move |cx| authenticate(client, &cx).await
         })
-        .detach();
+        .detach_and_log_err(cx);
 
         let urls: Vec<_> = args
             .paths_or_urls
@@ -850,15 +834,7 @@ fn handle_open_request(request: OpenRequest, app_state: Arc<AppState>, cx: &mut
                 let client = app_state.client.clone();
                 // we continue even if authentication fails as join_channel/ open channel notes will
                 // show a visible error message.
-                match authenticate(client, &cx).await {
-                    ConnectionResult::Timeout => {
-                        log::error!("Timeout during open request handling")
-                    }
-                    ConnectionResult::ConnectionReset => {
-                        log::error!("Connection reset during open request handling")
-                    }
-                    ConnectionResult::Result(r) => r?,
-                };
+                authenticate(client, &cx).await.log_err();
 
                 if let Some(channel_id) = request.join_channel {
                     cx.update(|cx| {
@@ -908,18 +884,18 @@ fn handle_open_request(request: OpenRequest, app_state: Arc<AppState>, cx: &mut
     }
 }
 
-async fn authenticate(client: Arc<Client>, cx: &AsyncApp) -> ConnectionResult<()> {
+async fn authenticate(client: Arc<Client>, cx: &AsyncApp) -> Result<()> {
     if stdout_is_a_pty() {
         if client::IMPERSONATE_LOGIN.is_some() {
-            return client.authenticate_and_connect(false, cx).await;
+            client.sign_in_with_optional_connect(false, cx).await?;
         } else if client.has_credentials(cx).await {
-            return client.authenticate_and_connect(true, cx).await;
+            client.sign_in_with_optional_connect(true, cx).await?;
         }
     } else if client.has_credentials(cx).await {
-        return client.authenticate_and_connect(true, cx).await;
+        client.sign_in_with_optional_connect(true, cx).await?;
     }
 
-    ConnectionResult::Result(Ok(()))
+    Ok(())
 }
 
 async fn system_id() -> Result<IdType> {

crates/zed/src/zed.rs

@@ -336,7 +336,7 @@ pub fn initialize_workspace(
         let edit_prediction_button = cx.new(|cx| {
             inline_completion_button::InlineCompletionButton::new(
                 app_state.fs.clone(),
-                app_state.cloud_user_store.clone(),
+                app_state.user_store.clone(),
                 inline_completion_menu_handle.clone(),
                 cx,
             )
@@ -4488,12 +4488,7 @@ mod tests {
             );
             image_viewer::init(cx);
             language_model::init(app_state.client.clone(), cx);
-            language_models::init(
-                app_state.user_store.clone(),
-                app_state.cloud_user_store.clone(),
-                app_state.client.clone(),
-                cx,
-            );
+            language_models::init(app_state.user_store.clone(), app_state.client.clone(), cx);
             web_search::init(cx);
             web_search_providers::init(app_state.client.clone(), cx);
             let prompt_builder = PromptBuilder::load(app_state.fs.clone(), false, cx);

crates/zed/src/zed/component_preview.rs

@@ -139,8 +139,7 @@ impl ComponentPreview {
         let project_clone = project.clone();
 
         cx.spawn_in(window, async move |entity, cx| {
-            let thread_store_future =
-                load_preview_thread_store(workspace_clone.clone(), project_clone.clone(), cx);
+            let thread_store_future = load_preview_thread_store(project_clone.clone(), cx);
             let text_thread_store_future =
                 load_preview_text_thread_store(workspace_clone.clone(), project_clone.clone(), cx);
 

crates/zed/src/zed/component_preview/preview_support/active_thread.rs

@@ -12,22 +12,19 @@ use ui::{App, Window};
 use workspace::Workspace;
 
 pub fn load_preview_thread_store(
-    workspace: WeakEntity<Workspace>,
     project: Entity<Project>,
     cx: &mut AsyncApp,
 ) -> Task<Result<Entity<ThreadStore>>> {
-    workspace
-        .update(cx, |workspace, cx| {
-            ThreadStore::load(
-                project.clone(),
-                workspace.app_state().cloud_user_store.clone(),
-                cx.new(|_| ToolWorkingSet::default()),
-                None,
-                Arc::new(PromptBuilder::new(None).unwrap()),
-                cx,
-            )
-        })
-        .unwrap_or(Task::ready(Err(anyhow!("workspace dropped"))))
+    cx.update(|cx| {
+        ThreadStore::load(
+            project.clone(),
+            cx.new(|_| ToolWorkingSet::default()),
+            None,
+            Arc::new(PromptBuilder::new(None).unwrap()),
+            cx,
+        )
+    })
+    .unwrap_or(Task::ready(Err(anyhow!("workspace dropped"))))
 }
 
 pub fn load_preview_text_thread_store(

crates/zed/src/zed/inline_completion_registry.rs

@@ -1,4 +1,4 @@
-use client::{Client, CloudUserStore, DisableAiSettings};
+use client::{Client, DisableAiSettings, UserStore};
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::Editor;
@@ -13,12 +13,12 @@ use util::ResultExt;
 use workspace::Workspace;
 use zeta::{ProviderDataCollection, ZetaInlineCompletionProvider};
 
-pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &mut App) {
+pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
     cx.observe_new({
         let editors = editors.clone();
         let client = client.clone();
-        let cloud_user_store = cloud_user_store.clone();
+        let user_store = user_store.clone();
         move |editor: &mut Editor, window, cx: &mut Context<Editor>| {
             if !editor.mode().is_full() {
                 return;
@@ -48,7 +48,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
                 editor,
                 provider,
                 &client,
-                cloud_user_store.clone(),
+                user_store.clone(),
                 window,
                 cx,
             );
@@ -60,7 +60,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
 
     let mut provider = all_language_settings(None, cx).edit_predictions.provider;
     cx.spawn({
-        let cloud_user_store = cloud_user_store.clone();
+        let user_store = user_store.clone();
         let editors = editors.clone();
         let client = client.clone();
 
@@ -72,7 +72,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
                         &editors,
                         provider,
                         &client,
-                        cloud_user_store.clone(),
+                        user_store.clone(),
                         cx,
                     );
                 })
@@ -85,12 +85,12 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
     cx.observe_global::<SettingsStore>({
         let editors = editors.clone();
         let client = client.clone();
-        let cloud_user_store = cloud_user_store.clone();
+        let user_store = user_store.clone();
         move |cx| {
             let new_provider = all_language_settings(None, cx).edit_predictions.provider;
 
             if new_provider != provider {
-                let tos_accepted = cloud_user_store.read(cx).has_accepted_tos();
+                let tos_accepted = user_store.read(cx).has_accepted_terms_of_service();
 
                 telemetry::event!(
                     "Edit Prediction Provider Changed",
@@ -104,7 +104,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
                     &editors,
                     provider,
                     &client,
-                    cloud_user_store.clone(),
+                    user_store.clone(),
                     cx,
                 );
 
@@ -145,7 +145,7 @@ fn assign_edit_prediction_providers(
     editors: &Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>>,
     provider: EditPredictionProvider,
     client: &Arc<Client>,
-    cloud_user_store: Entity<CloudUserStore>,
+    user_store: Entity<UserStore>,
     cx: &mut App,
 ) {
     for (editor, window) in editors.borrow().iter() {
@@ -155,7 +155,7 @@ fn assign_edit_prediction_providers(
                     editor,
                     provider,
                     &client,
-                    cloud_user_store.clone(),
+                    user_store.clone(),
                     window,
                     cx,
                 );
@@ -210,7 +210,7 @@ fn assign_edit_prediction_provider(
     editor: &mut Editor,
     provider: EditPredictionProvider,
     client: &Arc<Client>,
-    cloud_user_store: Entity<CloudUserStore>,
+    user_store: Entity<UserStore>,
     window: &mut Window,
     cx: &mut Context<Editor>,
 ) {
@@ -241,7 +241,7 @@ fn assign_edit_prediction_provider(
             }
         }
         EditPredictionProvider::Zed => {
-            if cloud_user_store.read(cx).is_authenticated() {
+            if user_store.read(cx).current_user().is_some() {
                 let mut worktree = None;
 
                 if let Some(buffer) = &singleton_buffer {
@@ -263,7 +263,7 @@ fn assign_edit_prediction_provider(
                     .map(|workspace| workspace.downgrade());
 
                 let zeta =
-                    zeta::Zeta::register(workspace, worktree, client.clone(), cloud_user_store, cx);
+                    zeta::Zeta::register(workspace, worktree, client.clone(), user_store, cx);
 
                 if let Some(buffer) = &singleton_buffer {
                     if buffer.read(cx).file().is_some() {