From c5f3a1d367207c7ce24ec47f9d00c54409be105c Mon Sep 17 00:00:00 2001
From: Ben Kunkle <ben@zed.dev>
Date: Thu, 5 Mar 2026 15:19:49 -0600
Subject: [PATCH] ep: Indicate active experiment (#50870)

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 ...
---
 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(-)

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(),