edit prediction: Assign providers when client status changes (#28919)

Agus Zubiaga created 8 months ago

There was recently a change that caused the Zed Edit Prediction provider
to only be assigned when the client was connected. However, this check
happened too early, resulting in restored buffers never getting
registered. We'll now subscribe to client status changes and reassign
providers accordingly.

Release Notes:

- edit prediction: Fixed bug disabling prediction in restored buffers

Change summary

crates/zed/src/zed/inline_completion_registry.rs | 41 ++++++++++-------
1 file changed, 25 insertions(+), 16 deletions(-)

Detailed changes

crates/zed/src/zed/inline_completion_registry.rs

@@ -5,9 +5,11 @@ 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, ZetaInlineCompletionProvider};
 
@@ -54,24 +56,31 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     })
     .detach();
 
+    cx.on_action(clear_zeta_edit_history);
+
     let mut provider = all_language_settings(None, cx).edit_predictions.provider;
-    for (editor, window) in editors.borrow().iter() {
-        _ = window.update(cx, |_window, window, cx| {
-            _ = editor.update(cx, |editor, cx| {
-                assign_edit_prediction_provider(
-                    editor,
-                    provider,
-                    &client,
-                    user_store.clone(),
-                    window,
-                    cx,
-                );
-            })
-        });
-    }
+    cx.spawn({
+        let user_store = user_store.clone();
+        let editors = editors.clone();
+        let client = client.clone();
 
-    cx.on_action(clear_zeta_edit_history);
-    assign_edit_prediction_providers(&editors, provider, &client, user_store.clone(), cx);
+        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();
+            }
+        }
+    })
+    .detach();
 
     cx.observe_global::<SettingsStore>({
         let editors = editors.clone();