From 992448b560f0e8ee00323c177564ac8ecce7b509 Mon Sep 17 00:00:00 2001
From: Danilo Leal <67129314+danilo-leal@users.noreply.github.com>
Date: Wed, 29 Oct 2025 22:16:13 -0300
Subject: [PATCH] edit prediction: Add ability to switch providers from the
 status bar menu (#41504)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Closes https://github.com/zed-industries/zed/issues/41500
Release Notes:
- Added the ability to switch between configured edit prediction
providers through the status bar menu.
---
 crates/codestral/src/codestral.rs             |   8 +
 .../src/edit_prediction_button.rs             | 164 ++++++++++++++----
 .../src/provider/copilot_chat.rs              |   5 +-
 .../language_models/src/provider/mistral.rs   |   4 +-
 crates/zed/src/zed.rs                         |   1 +
 5 files changed, 144 insertions(+), 38 deletions(-)
diff --git a/crates/codestral/src/codestral.rs b/crates/codestral/src/codestral.rs
index fe6b6678c99181facc4269df187c32c5a72ab565..e439cfb974fb55f4d30e5eb4be5c0dfa0d77c3d3 100644
--- a/crates/codestral/src/codestral.rs
+++ b/crates/codestral/src/codestral.rs
@@ -66,6 +66,14 @@ impl CodestralCompletionProvider {
         Self::api_key(cx).is_some()
     }
 
+    /// This is so we can immediately show Codestral as a provider users can
+    /// switch to in the edit prediction menu, if the API has been added
+    pub fn ensure_api_key_loaded(http_client: Arc, cx: &mut App) {
+        MistralLanguageModelProvider::global(http_client, cx)
+            .load_codestral_api_key(cx)
+            .detach();
+    }
+
     fn api_key(cx: &App) -> Option> {
         MistralLanguageModelProvider::try_global(cx)
             .and_then(|provider| provider.codestral_api_key(CODESTRAL_API_URL, cx))
diff --git a/crates/edit_prediction_button/src/edit_prediction_button.rs b/crates/edit_prediction_button/src/edit_prediction_button.rs
index 8b9bfc1c50092b65892cfcee9f4da1aeb2a0993e..70c861ab1112630c2e3293cb54a4e96c6754b3bd 100644
--- a/crates/edit_prediction_button/src/edit_prediction_button.rs
+++ b/crates/edit_prediction_button/src/edit_prediction_button.rs
@@ -1,5 +1,5 @@
 use anyhow::Result;
-use client::{UserStore, zed_urls};
+use client::{Client, UserStore, zed_urls};
 use cloud_llm_client::UsageLimit;
 use codestral::CodestralCompletionProvider;
 use copilot::{Copilot, Status};
@@ -192,6 +192,7 @@ impl Render for EditPredictionButton {
                                 Some(ContextMenu::build(window, cx, |menu, _, _| {
                                     let fs = fs.clone();
                                     let activate_url = activate_url.clone();
+
                                     menu.entry("Sign In", None, move |_, cx| {
                                         cx.open_url(activate_url.as_str())
                                     })
@@ -244,15 +245,8 @@ impl Render for EditPredictionButton {
                             } else {
                                 Some(ContextMenu::build(window, cx, |menu, _, _| {
                                     let fs = fs.clone();
-                                    menu.entry("Use Zed AI instead", None, move |_, cx| {
-                                        set_completion_provider(
-                                            fs.clone(),
-                                            cx,
-                                            EditPredictionProvider::Zed,
-                                        )
-                                    })
-                                    .separator()
-                                    .entry(
+
+                                    menu.entry(
                                         "Configure Codestral API Key",
                                         None,
                                         move |window, cx| {
@@ -262,6 +256,18 @@ impl Render for EditPredictionButton {
                                             );
                                         },
                                     )
+                                    .separator()
+                                    .entry(
+                                        "Use Zed AI instead",
+                                        None,
+                                        move |_, cx| {
+                                            set_completion_provider(
+                                                fs.clone(),
+                                                cx,
+                                                EditPredictionProvider::Zed,
+                                            )
+                                        },
+                                    )
                                 }))
                             }
                         })
@@ -412,6 +418,7 @@ impl EditPredictionButton {
         fs: Arc,
         user_store: Entity,
         popover_menu_handle: PopoverMenuHandle,
+        client: Arc,
         cx: &mut Context,
     ) -> Self {
         if let Some(copilot) = Copilot::global(cx) {
@@ -421,6 +428,8 @@ impl EditPredictionButton {
         cx.observe_global::(move |_, cx| cx.notify())
             .detach();
 
+        CodestralCompletionProvider::ensure_api_key_loaded(client.http_client(), cx);
+
         Self {
             editor_subscription: None,
             editor_enabled: None,
@@ -435,6 +444,89 @@ impl EditPredictionButton {
         }
     }
 
+    fn get_available_providers(&self, cx: &App) -> Vec {
+        let mut providers = Vec::new();
+
+        providers.push(EditPredictionProvider::Zed);
+
+        if let Some(copilot) = Copilot::global(cx) {
+            if matches!(copilot.read(cx).status(), Status::Authorized) {
+                providers.push(EditPredictionProvider::Copilot);
+            }
+        }
+
+        if let Some(supermaven) = Supermaven::global(cx) {
+            if let Supermaven::Spawned(agent) = supermaven.read(cx) {
+                if matches!(agent.account_status, AccountStatus::Ready) {
+                    providers.push(EditPredictionProvider::Supermaven);
+                }
+            }
+        }
+
+        if CodestralCompletionProvider::has_api_key(cx) {
+            providers.push(EditPredictionProvider::Codestral);
+        }
+
+        providers
+    }
+
+    fn add_provider_switching_section(
+        &self,
+        mut menu: ContextMenu,
+        current_provider: EditPredictionProvider,
+        cx: &App,
+    ) -> ContextMenu {
+        let available_providers = self.get_available_providers(cx);
+
+        let other_providers: Vec<_> = available_providers
+            .into_iter()
+            .filter(|p| *p != current_provider && *p != EditPredictionProvider::None)
+            .collect();
+
+        if !other_providers.is_empty() {
+            menu = menu.separator().header("Switch Providers");
+
+            for provider in other_providers {
+                let fs = self.fs.clone();
+
+                menu = match provider {
+                    EditPredictionProvider::Zed => menu.item(
+                        ContextMenuEntry::new("Zed AI")
+                            .documentation_aside(
+                                DocumentationSide::Left,
+                                DocumentationEdge::Top,
+                                |_| {
+                                    Label::new("Zed's edit prediction is powered by Zeta, an open-source, dataset mode.")
+                                        .into_any_element()
+                                },
+                            )
+                            .handler(move |_, cx| {
+                                set_completion_provider(fs.clone(), cx, provider);
+                            }),
+                    ),
+                    EditPredictionProvider::Copilot => {
+                        menu.entry("GitHub Copilot", None, move |_, cx| {
+                            set_completion_provider(fs.clone(), cx, provider);
+                        })
+                    }
+                    EditPredictionProvider::Supermaven => {
+                        menu.entry("Supermaven", None, move |_, cx| {
+                            set_completion_provider(fs.clone(), cx, provider);
+                        })
+                    }
+                    EditPredictionProvider::Codestral => {
+                        menu.entry("Codestral", None, move |_, cx| {
+                            set_completion_provider(fs.clone(), cx, provider);
+                        })
+                    }
+                    EditPredictionProvider::None => continue,
+                };
+            }
+        }
+
+        menu
+    }
+
     pub fn build_copilot_start_menu(
         &mut self,
         window: &mut Window,
@@ -572,8 +664,10 @@ impl EditPredictionButton {
         }
 
         menu = menu.separator().header("Privacy");
+
         if let Some(provider) = &self.edit_prediction_provider {
             let data_collection = provider.data_collection_state(cx);
+
             if data_collection.is_supported() {
                 let provider = provider.clone();
                 let enabled = data_collection.is_enabled();
@@ -691,7 +785,7 @@ impl EditPredictionButton {
                     }
                 }),
         ).item(
-            ContextMenuEntry::new("View Documentation")
+            ContextMenuEntry::new("View Docs")
                 .icon(IconName::FileGeneric)
                 .icon_color(Color::Muted)
                 .handler(move |_, cx| {
@@ -711,6 +805,7 @@ impl EditPredictionButton {
         if let Some(editor_focus_handle) = self.editor_focus_handle.clone() {
             menu = menu
                 .separator()
+                .header("Actions")
                 .entry(
                     "Predict Edit at Cursor",
                     Some(Box::new(ShowEditPrediction)),
@@ -721,7 +816,11 @@ impl EditPredictionButton {
                         }
                     },
                 )
-                .context(editor_focus_handle);
+                .context(editor_focus_handle)
+                .when(
+                    cx.has_flag::(),
+                    |this| this.action("Rate Completions", RateCompletions.boxed_clone()),
+                );
         }
 
         menu
@@ -733,15 +832,11 @@ impl EditPredictionButton {
         cx: &mut Context,
     ) -> Entity {
         ContextMenu::build(window, cx, |menu, window, cx| {
-            self.build_language_settings_menu(menu, window, cx)
-                .separator()
-                .entry("Use Zed AI instead", None, {
-                    let fs = self.fs.clone();
-                    move |_window, cx| {
-                        set_completion_provider(fs.clone(), cx, EditPredictionProvider::Zed)
-                    }
-                })
-                .separator()
+            let menu = self.build_language_settings_menu(menu, window, cx);
+            let menu =
+                self.add_provider_switching_section(menu, EditPredictionProvider::Copilot, cx);
+
+            menu.separator()
                 .link(
                     "Go to Copilot Settings",
                     OpenBrowser {
@@ -759,8 +854,11 @@ impl EditPredictionButton {
         cx: &mut Context,
     ) -> Entity {
         ContextMenu::build(window, cx, |menu, window, cx| {
-            self.build_language_settings_menu(menu, window, cx)
-                .separator()
+            let menu = self.build_language_settings_menu(menu, window, cx);
+            let menu =
+                self.add_provider_switching_section(menu, EditPredictionProvider::Supermaven, cx);
+
+            menu.separator()
                 .action("Sign Out", supermaven::SignOut.boxed_clone())
         })
     }
@@ -770,14 +868,12 @@ impl EditPredictionButton {
         window: &mut Window,
         cx: &mut Context,
     ) -> Entity {
-        let fs = self.fs.clone();
         ContextMenu::build(window, cx, |menu, window, cx| {
-            self.build_language_settings_menu(menu, window, cx)
-                .separator()
-                .entry("Use Zed AI instead", None, move |_, cx| {
-                    set_completion_provider(fs.clone(), cx, EditPredictionProvider::Zed)
-                })
-                .separator()
+            let menu = self.build_language_settings_menu(menu, window, cx);
+            let menu =
+                self.add_provider_switching_section(menu, EditPredictionProvider::Codestral, cx);
+
+            menu.separator()
                 .entry("Configure Codestral API Key", None, move |window, cx| {
                     window.dispatch_action(zed_actions::agent::OpenSettings.boxed_clone(), cx);
                 })
@@ -872,10 +968,10 @@ impl EditPredictionButton {
                     .separator();
             }
 
-            self.build_language_settings_menu(menu, window, cx).when(
-                cx.has_flag::(),
-                |this| this.action("Rate Completions", RateCompletions.boxed_clone()),
-            )
+            let menu = self.build_language_settings_menu(menu, window, cx);
+            let menu = self.add_provider_switching_section(menu, EditPredictionProvider::Zed, cx);
+
+            menu
         })
     }
 
diff --git a/crates/language_models/src/provider/copilot_chat.rs b/crates/language_models/src/provider/copilot_chat.rs
index 1941bd903951420266ba5c4609cb34c15130224e..6c665a0c1f06aa44e2b86f96517f7998fc02f4d3 100644
--- a/crates/language_models/src/provider/copilot_chat.rs
+++ b/crates/language_models/src/provider/copilot_chat.rs
@@ -1377,11 +1377,12 @@ impl Render for ConfigurationView {
 
                         v_flex().gap_2().child(Label::new(LABEL)).child(
                             Button::new("sign_in", "Sign in to use GitHub Copilot")
+                                .full_width()
+                                .style(ButtonStyle::Outlined)
                                 .icon_color(Color::Muted)
                                 .icon(IconName::Github)
                                 .icon_position(IconPosition::Start)
-                                .icon_size(IconSize::Medium)
-                                .full_width()
+                                .icon_size(IconSize::Small)
                                 .on_click(|_, window, cx| copilot::initiate_sign_in(window, cx)),
                         )
                     }
diff --git a/crates/language_models/src/provider/mistral.rs b/crates/language_models/src/provider/mistral.rs
index 66527792ff0b82348457fd28ae04dba60d10de5b..acd4a1c768e0d6ffdffbc3d69dcdc2bfd37fa928 100644
--- a/crates/language_models/src/provider/mistral.rs
+++ b/crates/language_models/src/provider/mistral.rs
@@ -753,9 +753,9 @@ struct ConfigurationView {
 impl ConfigurationView {
     fn new(state: Entity, window: &mut Window, cx: &mut Context) -> Self {
         let api_key_editor =
-            cx.new(|cx| InputField::new(window, cx, "0aBCDEFGhIjKLmNOpqrSTUVwxyzabCDE1f2"));
+            cx.new(|cx| InputField::new(window, cx, "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"));
         let codestral_api_key_editor =
-            cx.new(|cx| InputField::new(window, cx, "0aBCDEFGhIjKLmNOpqrSTUVwxyzabCDE1f2"));
+            cx.new(|cx| InputField::new(window, cx, "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"));
 
         cx.observe(&state, |_, _, cx| {
             cx.notify();
diff --git a/crates/zed/src/zed.rs b/crates/zed/src/zed.rs
index f71299997930040c848dd6f5c2819185cf8fee81..d712f782ca78745a94ce22c9a57900a8b8e42863 100644
--- a/crates/zed/src/zed.rs
+++ b/crates/zed/src/zed.rs
@@ -388,6 +388,7 @@ pub fn initialize_workspace(
                 app_state.fs.clone(),
                 app_state.user_store.clone(),
                 edit_prediction_menu_handle.clone(),
+                app_state.client.clone(),
                 cx,
             )
         });