Improve icon selection for edit prediction providers (#47911)

Release Notes:

- N/A

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>

Cargo.lock

@@ -3302,6 +3302,7 @@ dependencies = [
  "futures 0.3.31",
  "gpui",
  "http_client",
+ "icons",
  "language",
  "language_models",
  "log",
@@ -3734,6 +3735,7 @@ dependencies = [
  "futures 0.3.31",
  "gpui",
  "http_client",
+ "icons",
  "indoc",
  "language",
  "log",
@@ -5444,6 +5446,7 @@ version = "0.1.0"
 dependencies = [
  "client",
  "gpui",
+ "icons",
  "language",
  "text",
 ]

crates/codestral/Cargo.toml

@@ -15,6 +15,7 @@ edit_prediction.workspace = true
 futures.workspace = true
 gpui.workspace = true
 http_client.workspace = true
+icons.workspace = true
 language.workspace = true
 language_models.workspace = true
 log.workspace = true

crates/codestral/src/codestral.rs

@@ -1,9 +1,10 @@
 use anyhow::Result;
 use edit_prediction::cursor_excerpt;
-use edit_prediction_types::{EditPrediction, EditPredictionDelegate};
+use edit_prediction_types::{EditPrediction, EditPredictionDelegate, EditPredictionIconSet};
 use futures::AsyncReadExt;
 use gpui::{App, Context, Entity, Task};
 use http_client::HttpClient;
+use icons::IconName;
 use language::{
     language_settings::all_language_settings, Anchor, Buffer, BufferSnapshot, EditPreview, ToPoint,
 };
@@ -172,6 +173,10 @@ impl EditPredictionDelegate for CodestralEditPredictionDelegate {
         true
     }
 
+    fn icons(&self, _cx: &App) -> EditPredictionIconSet {
+        EditPredictionIconSet::new(IconName::AiMistral)
+    }
+
     fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
         Self::api_key(cx).is_some()
     }

crates/copilot/Cargo.toml

@@ -32,6 +32,7 @@ fs.workspace = true
 futures.workspace = true
 gpui.workspace = true
 edit_prediction_types.workspace = true
+icons.workspace = true
 language.workspace = true
 log.workspace = true
 lsp.workspace = true

crates/copilot/src/copilot_edit_prediction_delegate.rs

@@ -6,8 +6,11 @@ use crate::{
     },
 };
 use anyhow::Result;
-use edit_prediction_types::{EditPrediction, EditPredictionDelegate, interpolate_edits};
+use edit_prediction_types::{
+    EditPrediction, EditPredictionDelegate, EditPredictionIconSet, interpolate_edits,
+};
 use gpui::{App, Context, Entity, Task};
+use icons::IconName;
 use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, ToPointUtf16};
 use std::{ops::Range, sync::Arc, time::Duration};
 
@@ -50,6 +53,12 @@ impl EditPredictionDelegate for CopilotEditPredictionDelegate {
         true
     }
 
+    fn icons(&self, _cx: &App) -> EditPredictionIconSet {
+        EditPredictionIconSet::new(IconName::Copilot)
+            .with_disabled(IconName::CopilotDisabled)
+            .with_error(IconName::CopilotError)
+    }
+
     fn is_refreshing(&self, _cx: &App) -> bool {
         self.pending_refresh.is_some() && self.completion.is_none()
     }

crates/edit_prediction/src/edit_prediction.rs

@@ -655,6 +655,29 @@ impl EditPredictionStore {
         self.edit_prediction_model = model;
     }
 
+    pub fn icons(&self) -> edit_prediction_types::EditPredictionIconSet {
+        use ui::IconName;
+        match self.edit_prediction_model {
+            EditPredictionModel::Sweep => {
+                edit_prediction_types::EditPredictionIconSet::new(IconName::SweepAi)
+                    .with_disabled(IconName::SweepAiDisabled)
+                    .with_up(IconName::SweepAiUp)
+                    .with_down(IconName::SweepAiDown)
+                    .with_error(IconName::SweepAiError)
+            }
+            EditPredictionModel::Mercury => {
+                edit_prediction_types::EditPredictionIconSet::new(IconName::Inception)
+            }
+            EditPredictionModel::Zeta1 | EditPredictionModel::Zeta2 { .. } => {
+                edit_prediction_types::EditPredictionIconSet::new(IconName::ZedPredict)
+                    .with_disabled(IconName::ZedPredictDisabled)
+                    .with_up(IconName::ZedPredictUp)
+                    .with_down(IconName::ZedPredictDown)
+                    .with_error(IconName::ZedPredictError)
+            }
+        }
+    }
+
     pub fn has_sweep_api_token(&self, cx: &App) -> bool {
         self.sweep_ai.api_token.read(cx).has_key()
     }

crates/edit_prediction/src/zed_edit_prediction_delegate.rs

@@ -2,10 +2,13 @@ use std::{cmp, sync::Arc};
 
 use client::{Client, UserStore};
 use cloud_llm_client::EditPredictionRejectReason;
-use edit_prediction_types::{DataCollectionState, EditPredictionDelegate, SuggestionDisplayType};
+use edit_prediction_types::{
+    DataCollectionState, EditPredictionDelegate, EditPredictionIconSet, SuggestionDisplayType,
+};
 use gpui::{App, Entity, prelude::*};
 use language::{Buffer, ToPoint as _};
 use project::Project;
+use ui::prelude::*;
 
 use crate::{BufferEditPrediction, EditPredictionModel, EditPredictionStore};
 
@@ -58,6 +61,24 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate {
         true
     }
 
+    fn icons(&self, cx: &App) -> EditPredictionIconSet {
+        match self.store.read(cx).edit_prediction_model {
+            EditPredictionModel::Sweep => EditPredictionIconSet::new(IconName::SweepAi)
+                .with_disabled(IconName::SweepAiDisabled)
+                .with_up(IconName::SweepAiUp)
+                .with_down(IconName::SweepAiDown)
+                .with_error(IconName::SweepAiError),
+            EditPredictionModel::Mercury => EditPredictionIconSet::new(IconName::Inception),
+            EditPredictionModel::Zeta1 | EditPredictionModel::Zeta2 { .. } => {
+                EditPredictionIconSet::new(IconName::ZedPredict)
+                    .with_disabled(IconName::ZedPredictDisabled)
+                    .with_up(IconName::ZedPredictUp)
+                    .with_down(IconName::ZedPredictDown)
+                    .with_error(IconName::ZedPredictError)
+            }
+        }
+    }
+
     fn data_collection_state(&self, cx: &App) -> DataCollectionState {
         if let Some(buffer) = &self.singleton_buffer
             && let Some(file) = buffer.read(cx).file()

crates/edit_prediction_types/Cargo.toml

@@ -14,5 +14,6 @@ path = "src/edit_prediction_types.rs"
 [dependencies]
 client.workspace = true
 gpui.workspace = true
+icons.workspace = true
 language.workspace = true
 text.workspace = true

crates/edit_prediction_types/src/edit_prediction_types.rs

@@ -2,8 +2,50 @@ use std::{ops::Range, sync::Arc};
 
 use client::EditPredictionUsage;
 use gpui::{App, Context, Entity, SharedString};
+use icons::IconName;
 use language::{Anchor, Buffer, OffsetRangeExt};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct EditPredictionIconSet {
+    pub base: IconName,
+    pub disabled: IconName,
+    pub up: IconName,
+    pub down: IconName,
+    pub error: IconName,
+}
+
+impl EditPredictionIconSet {
+    pub fn new(base: IconName) -> Self {
+        Self {
+            base,
+            disabled: IconName::ZedPredictDisabled,
+            up: IconName::ZedPredictUp,
+            down: IconName::ZedPredictDown,
+            error: IconName::ZedPredictError,
+        }
+    }
+
+    pub fn with_disabled(mut self, disabled: IconName) -> Self {
+        self.disabled = disabled;
+        self
+    }
+
+    pub fn with_up(mut self, up: IconName) -> Self {
+        self.up = up;
+        self
+    }
+
+    pub fn with_down(mut self, down: IconName) -> Self {
+        self.down = down;
+        self
+    }
+
+    pub fn with_error(mut self, error: IconName) -> Self {
+        self.error = error;
+        self
+    }
+}
+
 /// Represents a predicted cursor position after an edit is applied.
 ///
 /// Since the cursor may be positioned inside newly inserted text that doesn't
@@ -110,6 +152,8 @@ pub trait EditPredictionDelegate: 'static + Sized {
         true
     }
 
+    fn icons(&self, cx: &App) -> EditPredictionIconSet;
+
     fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
         DataCollectionState::Unsupported
     }
@@ -156,6 +200,7 @@ pub trait EditPredictionDelegateHandle {
     fn show_predictions_in_menu(&self) -> bool;
     fn show_tab_accept_marker(&self) -> bool;
     fn supports_jump_to_edit(&self) -> bool;
+    fn icons(&self, cx: &App) -> EditPredictionIconSet;
     fn data_collection_state(&self, cx: &App) -> DataCollectionState;
     fn usage(&self, cx: &App) -> Option<EditPredictionUsage>;
     fn toggle_data_collection(&self, cx: &mut App);
@@ -202,6 +247,10 @@ where
         T::supports_jump_to_edit()
     }
 
+    fn icons(&self, cx: &App) -> EditPredictionIconSet {
+        self.read(cx).icons(cx)
+    }
+
     fn data_collection_state(&self, cx: &App) -> DataCollectionState {
         self.read(cx).data_collection_state(cx)
     }

crates/edit_prediction_ui/src/edit_prediction_button.rs

@@ -308,6 +308,13 @@ impl Render for EditPredictionButton {
             }
             provider @ (EditPredictionProvider::Experimental(_) | EditPredictionProvider::Zed) => {
                 let enabled = self.editor_enabled.unwrap_or(true);
+                let icons = self
+                    .edit_prediction_provider
+                    .as_ref()
+                    .map(|p| p.icons(cx))
+                    .unwrap_or_else(|| {
+                        edit_prediction_types::EditPredictionIconSet::new(IconName::ZedPredict)
+                    });
 
                 let ep_icon;
                 let tooltip_meta;
@@ -317,19 +324,19 @@ impl Render for EditPredictionButton {
                     EditPredictionProvider::Experimental(
                         EXPERIMENTAL_SWEEP_EDIT_PREDICTION_PROVIDER_NAME,
                     ) => {
-                        ep_icon = IconName::SweepAi;
+                        missing_token = edit_prediction::EditPredictionStore::try_global(cx)
+                            .is_some_and(|ep_store| !ep_store.read(cx).has_sweep_api_token(cx));
+                        ep_icon = if enabled { icons.base } else { icons.disabled };
                         tooltip_meta = if missing_token {
                             "Missing API key for Sweep"
                         } else {
                             "Powered by Sweep"
                         };
-                        missing_token = edit_prediction::EditPredictionStore::try_global(cx)
-                            .is_some_and(|ep_store| !ep_store.read(cx).has_sweep_api_token(cx));
                     }
                     EditPredictionProvider::Experimental(
                         EXPERIMENTAL_MERCURY_EDIT_PREDICTION_PROVIDER_NAME,
                     ) => {
-                        ep_icon = IconName::Inception;
+                        ep_icon = if enabled { icons.base } else { icons.disabled };
                         missing_token = edit_prediction::EditPredictionStore::try_global(cx)
                             .is_some_and(|ep_store| !ep_store.read(cx).has_mercury_api_token(cx));
                         tooltip_meta = if missing_token {
@@ -339,11 +346,7 @@ impl Render for EditPredictionButton {
                         };
                     }
                     _ => {
-                        ep_icon = if enabled {
-                            IconName::ZedPredict
-                        } else {
-                            IconName::ZedPredictDisabled
-                        };
+                        ep_icon = if enabled { icons.base } else { icons.disabled };
                         tooltip_meta = "Powered by Zeta"
                     }
                 };
@@ -849,10 +852,17 @@ impl EditPredictionButton {
         );
 
         if !self.editor_enabled.unwrap_or(true) {
+            let icons = self
+                .edit_prediction_provider
+                .as_ref()
+                .map(|p| p.icons(cx))
+                .unwrap_or_else(|| {
+                    edit_prediction_types::EditPredictionIconSet::new(IconName::ZedPredict)
+                });
             menu = menu.item(
                 ContextMenuEntry::new("This file is excluded.")
                     .disabled(true)
-                    .icon(IconName::ZedPredictDisabled)
+                    .icon(icons.disabled)
                     .icon_size(IconSize::Small),
             );
         }

crates/edit_prediction_ui/src/rate_prediction_modal.rs

@@ -839,20 +839,21 @@ impl Render for RatePredictionsModal {
                     .border_color(border_color)
                     .flex_shrink_0()
                     .overflow_hidden()
-                    .child(
+                    .child({
+                        let icons = self.ep_store.read(cx).icons();
                         h_flex()
                             .h_8()
                             .px_2()
                             .justify_between()
                             .border_b_1()
                             .border_color(border_color)
-                            .child(Icon::new(IconName::ZedPredict).size(IconSize::Small))
+                            .child(Icon::new(icons.base).size(IconSize::Small))
                             .child(
                                 Label::new("From most recent to oldest")
                                     .color(Color::Muted)
                                     .size(LabelSize::Small),
-                            ),
-                    )
+                            )
+                    })
                     .child(
                         div()
                             .id("completion_list")

crates/editor/src/edit_prediction_tests.rs

@@ -1,9 +1,12 @@
-use edit_prediction_types::{EditPredictionDelegate, PredictedCursorPosition};
+use edit_prediction_types::{
+    EditPredictionDelegate, EditPredictionIconSet, PredictedCursorPosition,
+};
 use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
 use indoc::indoc;
 use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
 use std::{ops::Range, sync::Arc};
 use text::{Point, ToOffset};
+use ui::prelude::*;
 
 use crate::{
     AcceptEditPrediction, EditPrediction, MenuEditPredictionsPolicy, editor_tests::init_test,
@@ -589,6 +592,10 @@ impl EditPredictionDelegate for FakeEditPredictionDelegate {
         true
     }
 
+    fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
+        EditPredictionIconSet::new(IconName::ZedPredict)
+    }
+
     fn is_enabled(
         &self,
         _buffer: &gpui::Entity<language::Buffer>,
@@ -656,6 +663,10 @@ impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
         false
     }
 
+    fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
+        EditPredictionIconSet::new(IconName::ZedPredict)
+    }
+
     fn is_enabled(
         &self,
         _buffer: &gpui::Entity<language::Buffer>,

crates/editor/src/editor.rs

@@ -9765,6 +9765,7 @@ impl Editor {
 
         let keybind = self.render_edit_prediction_accept_keybind(window, cx);
         let has_keybind = keybind.is_some();
+        let icons = Self::get_prediction_provider_icons(&self.edit_prediction_provider, cx);
 
         h_flex()
             .id("ep-line-popover")
@@ -9783,7 +9784,7 @@ impl Editor {
                 el.bg(status_colors.error_background)
                     .border_color(status_colors.error.opacity(0.6))
                     .pl_2()
-                    .child(Icon::new(IconName::ZedPredictError).color(Color::Error))
+                    .child(Icon::new(icons.error).color(Color::Error))
                     .cursor_default()
                     .hoverable_tooltip(move |_window, cx| {
                         cx.new(|_| MissingEditPredictionKeybindingTooltip).into()
@@ -9823,6 +9824,7 @@ impl Editor {
     ) -> Stateful<Div> {
         let keybind = self.render_edit_prediction_accept_keybind(window, cx);
         let has_keybind = keybind.is_some();
+        let icons = Self::get_prediction_provider_icons(&self.edit_prediction_provider, cx);
 
         let file_name = snapshot
             .file()
@@ -9845,7 +9847,7 @@ impl Editor {
                 el.bg(status_colors.error_background)
                     .border_color(status_colors.error.opacity(0.6))
                     .pl_2()
-                    .child(Icon::new(IconName::ZedPredictError).color(Color::Error))
+                    .child(Icon::new(icons.error).color(Color::Error))
                     .cursor_default()
                     .hoverable_tooltip(move |_window, cx| {
                         cx.new(|_| MissingEditPredictionKeybindingTooltip).into()
@@ -9887,16 +9889,13 @@ impl Editor {
         let editor_bg_color = cx.theme().colors().editor_background;
         editor_bg_color.blend(accent_color.opacity(0.6))
     }
-    fn get_prediction_provider_icon_name(
+    fn get_prediction_provider_icons(
         provider: &Option<RegisteredEditPredictionDelegate>,
-    ) -> IconName {
+        cx: &App,
+    ) -> edit_prediction_types::EditPredictionIconSet {
         match provider {
-            Some(provider) => match provider.provider.name() {
-                "copilot" => IconName::Copilot,
-                "supermaven" => IconName::Supermaven,
-                _ => IconName::ZedPredict,
-            },
-            None => IconName::ZedPredict,
+            Some(provider) => provider.provider.icons(cx),
+            None => edit_prediction_types::EditPredictionIconSet::new(IconName::ZedPredict),
         }
     }
 
@@ -9911,7 +9910,7 @@ impl Editor {
         cx: &mut Context<Editor>,
     ) -> Option<AnyElement> {
         let provider = self.edit_prediction_provider.as_ref()?;
-        let provider_icon = Self::get_prediction_provider_icon_name(&self.edit_prediction_provider);
+        let icons = Self::get_prediction_provider_icons(&self.edit_prediction_provider, cx);
 
         let is_refreshing = provider.provider.is_refreshing(cx);
 
@@ -9941,16 +9940,16 @@ impl Editor {
                                     use text::ToPoint as _;
                                     if target.text_anchor.to_point(snapshot).row > cursor_point.row
                                     {
-                                        Icon::new(IconName::ZedPredictDown)
+                                        Icon::new(icons.down)
                                     } else {
-                                        Icon::new(IconName::ZedPredictUp)
+                                        Icon::new(icons.up)
                                     }
                                 }
                                 EditPrediction::MoveOutside { .. } => {
                                     // TODO [zeta2] custom icon for external jump?
-                                    Icon::new(provider_icon)
+                                    Icon::new(icons.base)
                                 }
-                                EditPrediction::Edit { .. } => Icon::new(provider_icon),
+                                EditPrediction::Edit { .. } => Icon::new(icons.base),
                             }))
                             .child(
                                 h_flex()
@@ -10017,11 +10016,11 @@ impl Editor {
                     cx,
                 )?,
 
-                None => pending_completion_container(provider_icon)
+                None => pending_completion_container(icons.base)
                     .child(Label::new("...").size(LabelSize::Small)),
             },
 
-            None => pending_completion_container(provider_icon)
+            None => pending_completion_container(icons.base)
                 .child(Label::new("...").size(LabelSize::Small)),
         };
 
@@ -10131,6 +10130,8 @@ impl Editor {
             .map(|provider| provider.provider.supports_jump_to_edit())
             .unwrap_or(true);
 
+        let icons = Self::get_prediction_provider_icons(&self.edit_prediction_provider, cx);
+
         match &completion.completion {
             EditPrediction::MoveWithin {
                 target, snapshot, ..
@@ -10146,9 +10147,9 @@ impl Editor {
                         .flex_1()
                         .child(
                             if target.text_anchor.to_point(snapshot).row > cursor_point.row {
-                                Icon::new(IconName::ZedPredictDown)
+                                Icon::new(icons.down)
                             } else {
-                                Icon::new(IconName::ZedPredictUp)
+                                Icon::new(icons.up)
                             },
                         )
                         .child(Label::new("Jump to Edit")),
@@ -10164,7 +10165,7 @@ impl Editor {
                         .px_2()
                         .gap_2()
                         .flex_1()
-                        .child(Icon::new(IconName::ZedPredict))
+                        .child(Icon::new(icons.base))
                         .child(Label::new(format!("Jump to {file_name}"))),
                 )
             }
@@ -10197,9 +10198,7 @@ impl Editor {
                     render_relative_row_jump("", cursor_point.row, first_edit_row)
                         .into_any_element()
                 } else {
-                    let icon_name =
-                        Editor::get_prediction_provider_icon_name(&self.edit_prediction_provider);
-                    Icon::new(icon_name).into_any_element()
+                    Icon::new(icons.base).into_any_element()
                 };
 
                 Some(

crates/icons/src/icons.rs

@@ -222,6 +222,10 @@ pub enum IconName {
     SupermavenInit,
     SwatchBook,
     SweepAi,
+    SweepAiDisabled,
+    SweepAiDown,
+    SweepAiError,
+    SweepAiUp,
     Tab,
     Terminal,
     TerminalAlt,

crates/supermaven/src/supermaven_edit_prediction_delegate.rs

@@ -1,6 +1,6 @@
 use crate::{Supermaven, SupermavenCompletionStateId};
 use anyhow::Result;
-use edit_prediction_types::{EditPrediction, EditPredictionDelegate};
+use edit_prediction_types::{EditPrediction, EditPredictionDelegate, EditPredictionIconSet};
 use futures::StreamExt as _;
 use gpui::{App, Context, Entity, EntityId, Task};
 use language::{Anchor, Buffer, BufferSnapshot};
@@ -11,6 +11,7 @@ use std::{
     time::Duration,
 };
 use text::{ToOffset, ToPoint};
+use ui::prelude::*;
 use unicode_segmentation::UnicodeSegmentation;
 
 pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(75);
@@ -126,6 +127,12 @@ impl EditPredictionDelegate for SupermavenEditPredictionDelegate {
         false
     }
 
+    fn icons(&self, _cx: &App) -> EditPredictionIconSet {
+        EditPredictionIconSet::new(IconName::Supermaven)
+            .with_disabled(IconName::SupermavenDisabled)
+            .with_error(IconName::SupermavenError)
+    }
+
     fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
         self.supermaven.read(cx).is_enabled()
     }