Reorganize edit prediction code and remove old experiments (#44187)

Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <agus@zed.dev>
Co-authored-by: Ben Kunkle <ben@zed.dev>

crates/zed/src/main.rs

@@ -581,7 +581,7 @@ pub fn main() {
         language_model::init(app_state.client.clone(), cx);
         language_models::init(app_state.user_store.clone(), app_state.client.clone(), cx);
         acp_tools::init(cx);
-        zeta2_tools::init(cx);
+        edit_prediction_ui::init(cx);
         web_search::init(cx);
         web_search_providers::init(app_state.client.clone(), cx);
         snippet_provider::init(cx);
@@ -640,7 +640,7 @@ pub fn main() {
         settings_ui::init(cx);
         keymap_editor::init(cx);
         extensions_ui::init(cx);
-        zeta::init(cx);
+        edit_prediction::init(cx);
         inspector_ui::init(app_state.clone(), cx);
         json_schema_store::init(cx);
         miniprofiler_ui::init(*STARTUP_TIME.get().unwrap(), cx);

crates/zed/src/zed.rs

@@ -401,8 +401,8 @@ pub fn initialize_workspace(
         unstable_version_notification(cx);
 
         let edit_prediction_menu_handle = PopoverMenuHandle::default();
-        let edit_prediction_button = cx.new(|cx| {
-            edit_prediction_button::EditPredictionButton::new(
+        let edit_prediction_ui = cx.new(|cx| {
+            edit_prediction_ui::EditPredictionButton::new(
                 app_state.fs.clone(),
                 app_state.user_store.clone(),
                 edit_prediction_menu_handle.clone(),
@@ -411,7 +411,7 @@ pub fn initialize_workspace(
             )
         });
         workspace.register_action({
-            move |_, _: &edit_prediction_button::ToggleMenu, window, cx| {
+            move |_, _: &edit_prediction_ui::ToggleMenu, window, cx| {
                 edit_prediction_menu_handle.toggle(window, cx);
             }
         });
@@ -450,7 +450,7 @@ pub fn initialize_workspace(
             status_bar.add_left_item(lsp_button, window, cx);
             status_bar.add_left_item(diagnostic_summary, window, cx);
             status_bar.add_left_item(activity_indicator, window, cx);
-            status_bar.add_right_item(edit_prediction_button, window, cx);
+            status_bar.add_right_item(edit_prediction_ui, window, cx);
             status_bar.add_right_item(active_buffer_language, window, cx);
             status_bar.add_right_item(active_toolchain_language, window, cx);
             status_bar.add_right_item(line_ending_indicator, window, cx);

crates/zed/src/zed/edit_prediction_registry.rs

@@ -1,7 +1,8 @@
 use client::{Client, UserStore};
-use codestral::CodestralCompletionProvider;
+use codestral::CodestralEditPredictionDelegate;
 use collections::HashMap;
-use copilot::{Copilot, CopilotCompletionProvider};
+use copilot::{Copilot, CopilotEditPredictionDelegate};
+use edit_prediction::{SweepFeatureFlag, ZedEditPredictionDelegate, Zeta2FeatureFlag};
 use editor::Editor;
 use feature_flags::FeatureFlagAppExt;
 use gpui::{AnyWindowHandle, App, AppContext as _, Context, Entity, WeakEntity};
@@ -12,9 +13,8 @@ use settings::{
     EXPERIMENTAL_ZETA2_EDIT_PREDICTION_PROVIDER_NAME, SettingsStore,
 };
 use std::{cell::RefCell, rc::Rc, sync::Arc};
-use supermaven::{Supermaven, SupermavenCompletionProvider};
+use supermaven::{Supermaven, SupermavenEditPredictionDelegate};
 use ui::Window;
-use zeta::{SweepFeatureFlag, Zeta2FeatureFlag, ZetaEditPredictionProvider};
 
 pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
@@ -59,7 +59,7 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     })
     .detach();
 
-    cx.on_action(clear_zeta_edit_history);
+    cx.on_action(clear_edit_prediction_store_edit_history);
 
     let mut provider = all_language_settings(None, cx).edit_predictions.provider;
     cx.subscribe(&user_store, {
@@ -100,9 +100,9 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     .detach();
 }
 
-fn clear_zeta_edit_history(_: &zeta::ClearHistory, cx: &mut App) {
-    if let Some(zeta) = zeta::Zeta::try_global(cx) {
-        zeta.update(cx, |zeta, _| zeta.clear_history());
+fn clear_edit_prediction_store_edit_history(_: &edit_prediction::ClearHistory, cx: &mut App) {
+    if let Some(ep_store) = edit_prediction::EditPredictionStore::try_global(cx) {
+        ep_store.update(cx, |ep_store, _| ep_store.clear_history());
     }
 }
 
@@ -176,7 +176,7 @@ fn assign_edit_prediction_provider(
 
     match provider {
         EditPredictionProvider::None => {
-            editor.set_edit_prediction_provider::<ZetaEditPredictionProvider>(None, window, cx);
+            editor.set_edit_prediction_provider::<ZedEditPredictionDelegate>(None, window, cx);
         }
         EditPredictionProvider::Copilot => {
             if let Some(copilot) = Copilot::global(cx) {
@@ -187,55 +187,61 @@ fn assign_edit_prediction_provider(
                         copilot.register_buffer(&buffer, cx);
                     });
                 }
-                let provider = cx.new(|_| CopilotCompletionProvider::new(copilot));
+                let provider = cx.new(|_| CopilotEditPredictionDelegate::new(copilot));
                 editor.set_edit_prediction_provider(Some(provider), window, cx);
             }
         }
         EditPredictionProvider::Supermaven => {
             if let Some(supermaven) = Supermaven::global(cx) {
-                let provider = cx.new(|_| SupermavenCompletionProvider::new(supermaven));
+                let provider = cx.new(|_| SupermavenEditPredictionDelegate::new(supermaven));
                 editor.set_edit_prediction_provider(Some(provider), window, cx);
             }
         }
         EditPredictionProvider::Codestral => {
             let http_client = client.http_client();
-            let provider = cx.new(|_| CodestralCompletionProvider::new(http_client));
+            let provider = cx.new(|_| CodestralEditPredictionDelegate::new(http_client));
             editor.set_edit_prediction_provider(Some(provider), window, cx);
         }
         value @ (EditPredictionProvider::Experimental(_) | EditPredictionProvider::Zed) => {
-            let zeta = zeta::Zeta::global(client, &user_store, cx);
+            let ep_store = edit_prediction::EditPredictionStore::global(client, &user_store, cx);
 
             if let Some(project) = editor.project()
                 && let Some(buffer) = &singleton_buffer
                 && buffer.read(cx).file().is_some()
             {
-                let has_model = zeta.update(cx, |zeta, cx| {
+                let has_model = ep_store.update(cx, |ep_store, cx| {
                     let model = if let EditPredictionProvider::Experimental(name) = value {
                         if name == EXPERIMENTAL_SWEEP_EDIT_PREDICTION_PROVIDER_NAME
                             && cx.has_flag::<SweepFeatureFlag>()
                         {
-                            zeta::ZetaEditPredictionModel::Sweep
+                            edit_prediction::EditPredictionModel::Sweep
                         } else if name == EXPERIMENTAL_ZETA2_EDIT_PREDICTION_PROVIDER_NAME
                             && cx.has_flag::<Zeta2FeatureFlag>()
                         {
-                            zeta::ZetaEditPredictionModel::Zeta2
+                            edit_prediction::EditPredictionModel::Zeta2
                         } else {
                             return false;
                         }
                     } else if user_store.read(cx).current_user().is_some() {
-                        zeta::ZetaEditPredictionModel::Zeta1
+                        edit_prediction::EditPredictionModel::Zeta1
                     } else {
                         return false;
                     };
 
-                    zeta.set_edit_prediction_model(model);
-                    zeta.register_buffer(buffer, project, cx);
+                    ep_store.set_edit_prediction_model(model);
+                    ep_store.register_buffer(buffer, project, cx);
                     true
                 });
 
                 if has_model {
                     let provider = cx.new(|cx| {
-                        ZetaEditPredictionProvider::new(project.clone(), &client, &user_store, cx)
+                        ZedEditPredictionDelegate::new(
+                            project.clone(),
+                            singleton_buffer,
+                            &client,
+                            &user_store,
+                            cx,
+                        )
                     });
                     editor.set_edit_prediction_provider(Some(provider), window, cx);
                 }