From bb400433629f115636bf727c1438445855a58f4f 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/inline_completion_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 c360470277..38f52359b7 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(),
-                                            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/inline_completion_registry.rs b/crates/zed/src/zed/inline_completion_registry.rs
index 378c172078..acd0be0291 100644
--- a/crates/zed/src/zed/inline_completion_registry.rs
+++ b/crates/zed/src/zed/inline_completion_registry.rs
@@ -6,11 +6,9 @@ use gpui::{AnyWindowHandle, App, AppContext as _, Context, Entity, WeakEntity};
 use language::language_settings::{EditPredictionProvider, all_language_settings};
 use project::DisableAiSettings;
 use settings::{Settings as _, 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, ZetaInlineCompletionProvider};
 
@@ -60,25 +58,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();