diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs
index 2c0407b56b2b5c7baffc1094dfaf1cd4baf6393f..d73fdc9b39e350ff0697cdb5cdf1ec7d0c866a72 100644
--- a/crates/edit_prediction/src/edit_prediction.rs
+++ b/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();
diff --git a/crates/edit_prediction_ui/src/edit_prediction_button.rs b/crates/edit_prediction_ui/src/edit_prediction_button.rs
index b00a229164d480d38312ca97cac31a23010f8b69..5e9ec08b96442dc8c10e89aa43f891e7743f85ef 100644
--- a/crates/edit_prediction_ui/src/edit_prediction_button.rs
+++ b/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(),