zeta: Revised data-collection onboarding experience (#24031)

Release Notes:

- N/A

---------

Co-authored-by: Danilo <danilo@zed.dev>
Co-authored-by: Danilo Leal <daniloleal09@gmail.com>
Co-authored-by: João Marcos <marcospb19@hotmail.com>

crates/zed/src/main.rs

@@ -439,7 +439,6 @@ fn main() {
         inline_completion_registry::init(
             app_state.client.clone(),
             app_state.user_store.clone(),
-            app_state.fs.clone(),
             cx,
         );
         let prompt_builder = PromptBuilder::load(app_state.fs.clone(), stdout_is_a_pty(), cx);

crates/zed/src/zed.rs

@@ -176,7 +176,6 @@ pub fn initialize_workspace(
                 workspace.weak_handle(),
                 app_state.fs.clone(),
                 app_state.user_store.clone(),
-                app_state.client.clone(),
                 popover_menu_handle.clone(),
                 cx,
             )

crates/zed/src/zed/inline_completion_registry.rs

@@ -1,21 +1,17 @@
-use std::{cell::RefCell, rc::Rc, sync::Arc};
-
 use client::{Client, UserStore};
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::{Editor, EditorMode};
 use feature_flags::{FeatureFlagAppExt, PredictEditsFeatureFlag};
-use fs::Fs;
 use gpui::{AnyWindowHandle, App, AppContext, Context, Entity, WeakEntity};
 use language::language_settings::{all_language_settings, InlineCompletionProvider};
 use settings::SettingsStore;
+use std::{cell::RefCell, rc::Rc, sync::Arc};
 use supermaven::{Supermaven, SupermavenCompletionProvider};
 use ui::Window;
-use workspace::Workspace;
-use zed_predict_onboarding::ZedPredictModal;
 use zeta::ProviderDataCollection;
 
-pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, fs: Arc<dyn Fs>, cx: &mut App) {
+pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
     cx.observe_new({
         let editors = editors.clone();
@@ -96,7 +92,6 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, fs: Arc<dyn Fs>,
         let editors = editors.clone();
         let client = client.clone();
         let user_store = user_store.clone();
-        let fs = fs.clone();
         move |cx| {
             let new_provider = all_language_settings(None, cx).inline_completions.provider;
             if new_provider != provider {
@@ -120,21 +115,10 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, fs: Arc<dyn Fs>,
                                 return;
                             };
 
-                            let Some(Some(workspace)) = window
-                                .update(cx, |_, window, _| window.root().flatten())
-                                .ok()
-                            else {
-                                return;
-                            };
-
                             window
                                 .update(cx, |_, window, cx| {
-                                    ZedPredictModal::toggle(
-                                        workspace,
-                                        user_store.clone(),
-                                        client.clone(),
-                                        fs.clone(),
-                                        window,
+                                    window.dispatch_action(
+                                        Box::new(zed_actions::OpenZedPredictOnboarding),
                                         cx,
                                     );
                                 })
@@ -228,6 +212,7 @@ fn assign_inline_completion_provider(
     window: &mut Window,
     cx: &mut Context<Editor>,
 ) {
+    // TODO: Do we really want to collect data only for singleton buffers?
     let singleton_buffer = editor.buffer().read(cx).as_singleton();
 
     match provider {
@@ -255,7 +240,23 @@ fn assign_inline_completion_provider(
             if cx.has_flag::<PredictEditsFeatureFlag>()
                 || (cfg!(debug_assertions) && client.status().borrow().is_connected())
             {
-                let zeta = zeta::Zeta::register(client.clone(), user_store, cx);
+                let mut worktree = None;
+
+                if let Some(buffer) = &singleton_buffer {
+                    if let Some(file) = buffer.read(cx).file() {
+                        let id = file.worktree_id(cx);
+                        if let Some(inner_worktree) = editor
+                            .project
+                            .as_ref()
+                            .and_then(|project| project.read(cx).worktree_for_id(id, cx))
+                        {
+                            worktree = Some(inner_worktree);
+                        }
+                    }
+                }
+
+                let zeta = zeta::Zeta::register(worktree, client.clone(), user_store, cx);
+
                 if let Some(buffer) = &singleton_buffer {
                     if buffer.read(cx).file().is_some() {
                         zeta.update(cx, |zeta, cx| {
@@ -264,12 +265,8 @@ fn assign_inline_completion_provider(
                     }
                 }
 
-                let data_collection = ProviderDataCollection::new(
-                    zeta.clone(),
-                    window.root::<Workspace>().flatten(),
-                    singleton_buffer,
-                    cx,
-                );
+                let data_collection =
+                    ProviderDataCollection::new(zeta.clone(), singleton_buffer, cx);
 
                 let provider =
                     cx.new(|_| zeta::ZetaInlineCompletionProvider::new(zeta, data_collection));