ep: Indicate active experiment (#50870)

Ben Kunkle created 1 month ago

Closes #ISSUE

Before you mark this PR as ready for review, make sure that you have:
- [ ] Added a solid test coverage and/or screenshots from doing manual
testing
- [ ] Done a self-review taking into account security and performance
aspects
- [ ] Aligned any UI changes with the [UI
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)

Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/edit_prediction/src/edit_prediction.rs           | 9 +++++++++
crates/edit_prediction_ui/src/edit_prediction_button.rs | 4 +++-
2 files changed, 12 insertions(+), 1 deletion(-)

Detailed changes

crates/edit_prediction/src/edit_prediction.rs

@@ -795,6 +795,15 @@ impl EditPredictionStore {
         &self.available_experiments
     }
 
+    pub fn active_experiment(&self) -> Option<&str> {
+        self.preferred_experiment.as_deref().or_else(|| {
+            self.shown_predictions
+                .iter()
+                .find_map(|p| p.model_version.as_ref())
+                .and_then(|model_version| model_version.strip_prefix("zeta2:"))
+        })
+    }
+
     pub fn refresh_available_experiments(&mut self, cx: &mut Context<Self>) {
         let client = self.client.clone();
         let llm_token = self.llm_token.clone();

crates/edit_prediction_ui/src/edit_prediction_button.rs

@@ -1198,6 +1198,7 @@ impl EditPredictionButton {
                     let store = store.read(cx);
                     let experiments = store.available_experiments().to_vec();
                     let preferred = store.preferred_experiment().map(|s| s.to_owned());
+                    let active = store.active_experiment().map(|s| s.to_owned());
 
                     let preferred_for_submenu = preferred.clone();
                     menu = menu
@@ -1219,7 +1220,8 @@ impl EditPredictionButton {
                                 },
                             );
                             for experiment in &experiments {
-                                let is_selected = preferred.as_deref() == Some(experiment.as_str());
+                                let is_selected = active.as_deref() == Some(experiment.as_str())
+                                    || preferred.as_deref() == Some(experiment.as_str());
                                 let experiment_name = experiment.clone();
                                 menu = menu.toggleable_entry(
                                     experiment.clone(),