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 | 43 +++++++++++--------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/crates/zed/src/zed/inline_completion_registry.rs b/crates/zed/src/zed/inline_completion_registry.rs
index 639644b087..f2e9d21b96 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();
 
-    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.on_action(clear_zeta_edit_history);
-    assign_edit_prediction_providers(&editors, provider, &client, user_store.clone(), cx);
+
+    let mut provider = all_language_settings(None, cx).edit_predictions.provider;
+    cx.spawn({
+        let user_store = user_store.clone();
+        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();
+            }
+        }
+    })
+    .detach();
 
     cx.observe_global::<SettingsStore>({
         let editors = editors.clone();