From a7a7335da4ebeb6ebd5c644f108edf2441f7d77b Mon Sep 17 00:00:00 2001
From: Agus Zubiaga <agus@zed.dev>
Date: Wed, 16 Apr 2025 18:09:25 -0600
Subject: [PATCH] edit prediction: Assign providers when client status changes
 (#28919)

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
---
 .../zed/src/zed/inline_completion_registry.rs | 41 +++++++++++--------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/crates/zed/src/zed/inline_completion_registry.rs b/crates/zed/src/zed/inline_completion_registry.rs
index 639644b087f3764925a243e7829e059297b6a252..f2e9d21b96ad54462a27f80e0c9c352a215eacdf 100644
--- a/crates/zed/src/zed/inline_completion_registry.rs
+++ b/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();