Ensure Edit Prediction provider is properly assigned on sign-in (#35885)

Marshall Bowers and Richard Feldman created 3 months ago

This PR fixes an issue where Edit Predictions would not be available in
buffers that were opened when the workspace loaded.

The issue was that there was a race condition between fetching/setting
the authenticated user state and when we assigned the Edit Prediction
provider to buffers that were already opened.

We now wait for the event that we emit when we have successfully loaded
the user in order to assign the Edit Prediction provider, as we'll know
the user has been loaded into the `UserStore` by that point.

Closes https://github.com/zed-industries/zed/issues/35883

Release Notes:

- Fixed an issue where Edit Predictions were not working in buffers that
were open when the workspace initially loaded.

Co-authored-by: Richard Feldman <[e-mail]>

Change summary

crates/client/src/user.rs                      | 39 ++++++++++++++-----
crates/zed/src/zed/edit_prediction_registry.rs | 29 +++++---------
2 files changed, 39 insertions(+), 29 deletions(-)

Detailed changes

crates/client/src/user.rs

@@ -226,17 +226,35 @@ impl UserStore {
                     match status {
                         Status::Authenticated | Status::Connected { .. } => {
                             if let Some(user_id) = client.user_id() {
-                                let response = client.cloud_client().get_authenticated_user().await;
-                                let mut current_user = None;
+                                let response = client
+                                    .cloud_client()
+                                    .get_authenticated_user()
+                                    .await
+                                    .log_err();
+
+                                let current_user_and_response = if let Some(response) = response {
+                                    let user = Arc::new(User {
+                                        id: user_id,
+                                        github_login: response.user.github_login.clone().into(),
+                                        avatar_uri: response.user.avatar_url.clone().into(),
+                                        name: response.user.name.clone(),
+                                    });
+
+                                    Some((user, response))
+                                } else {
+                                    None
+                                };
+                                current_user_tx
+                                    .send(
+                                        current_user_and_response
+                                            .as_ref()
+                                            .map(|(user, _)| user.clone()),
+                                    )
+                                    .await
+                                    .ok();
+
                                 cx.update(|cx| {
-                                    if let Some(response) = response.log_err() {
-                                        let user = Arc::new(User {
-                                            id: user_id,
-                                            github_login: response.user.github_login.clone().into(),
-                                            avatar_uri: response.user.avatar_url.clone().into(),
-                                            name: response.user.name.clone(),
-                                        });
-                                        current_user = Some(user.clone());
+                                    if let Some((user, response)) = current_user_and_response {
                                         this.update(cx, |this, cx| {
                                             this.by_github_login
                                                 .insert(user.github_login.clone(), user_id);
@@ -247,7 +265,6 @@ impl UserStore {
                                         anyhow::Ok(())
                                     }
                                 })??;
-                                current_user_tx.send(current_user).await.ok();
 
                                 this.update(cx, |_, cx| cx.notify())?;
                             }

crates/zed/src/zed/edit_prediction_registry.rs

@@ -5,11 +5,9 @@ use editor::Editor;
 use gpui::{AnyWindowHandle, App, AppContext as _, Context, Entity, WeakEntity};
 use language::language_settings::{EditPredictionProvider, all_language_settings};
 use settings::SettingsStore;
-use smol::stream::StreamExt;
 use std::{cell::RefCell, rc::Rc, sync::Arc};
 use supermaven::{Supermaven, SupermavenCompletionProvider};
 use ui::Window;
-use util::ResultExt;
 use workspace::Workspace;
 use zeta::{ProviderDataCollection, ZetaEditPredictionProvider};
 
@@ -59,25 +57,20 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     cx.on_action(clear_zeta_edit_history);
 
     let mut provider = all_language_settings(None, cx).edit_predictions.provider;
-    cx.spawn({
-        let user_store = user_store.clone();
+    cx.subscribe(&user_store, {
         let editors = editors.clone();
         let client = client.clone();
-
-        async move |cx| {
-            let mut status = client.status();
-            while let Some(_status) = status.next().await {
-                cx.update(|cx| {
-                    assign_edit_prediction_providers(
-                        &editors,
-                        provider,
-                        &client,
-                        user_store.clone(),
-                        cx,
-                    );
-                })
-                .log_err();
+        move |user_store, event, cx| match event {
+            client::user::Event::PrivateUserInfoUpdated => {
+                assign_edit_prediction_providers(
+                    &editors,
+                    provider,
+                    &client,
+                    user_store.clone(),
+                    cx,
+                );
             }
+            _ => {}
         }
     })
     .detach();