From 315a92091b91b4448b4c95ecf9e3dc3fa1bd7a62 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Fri, 8 Aug 2025 14:10:09 -0400
Subject: [PATCH] Ensure Edit Prediction provider is properly assigned on
 sign-in (#35885)

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 <oss@rtfeldman.com>
---
 crates/client/src/user.rs                     | 39 +++++++++++++------
 .../zed/src/zed/edit_prediction_registry.rs   | 29 ++++++--------
 2 files changed, 39 insertions(+), 29 deletions(-)

diff --git a/crates/client/src/user.rs b/crates/client/src/user.rs
index 9f76dd7ad08034e814d4be0ff825f37415a3bccd..faf46945d888d3a3da18f69a16cdcc11009e1937 100644
--- a/crates/client/src/user.rs
+++ b/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())?;
                             }
diff --git a/crates/zed/src/zed/edit_prediction_registry.rs b/crates/zed/src/zed/edit_prediction_registry.rs
index b9f561c0e7884f3d7c8579cd7b65fa1c787b5291..da4b6e78c62db447a6d9669dd378ffe9d6fb84d2 100644
--- a/crates/zed/src/zed/edit_prediction_registry.rs
+++ b/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();