ep: Send preferred experiment in a header (#54154)

Ben Kunkle created 

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Closes #ISSUE

Release Notes:

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

Change summary

crates/cloud_llm_client/src/cloud_llm_client.rs     |  3 +++
crates/edit_prediction/src/edit_prediction.rs       | 15 +++++++++++----
crates/edit_prediction/src/edit_prediction_tests.rs |  1 -
crates/edit_prediction/src/fim.rs                   |  1 -
crates/edit_prediction/src/mercury.rs               |  1 -
crates/edit_prediction/src/prediction.rs            |  1 -
crates/edit_prediction/src/zeta.rs                  |  4 +---
crates/edit_prediction_cli/src/load_project.rs      |  1 -
crates/zeta_prompt/src/zeta_prompt.rs               |  8 --------
9 files changed, 15 insertions(+), 20 deletions(-)

Detailed changes

crates/cloud_llm_client/src/cloud_llm_client.rs 🔗 

@@ -12,6 +12,9 @@ use uuid::Uuid;
 /// The name of the header used to indicate which version of Zed the client is running.
 pub const ZED_VERSION_HEADER_NAME: &str = "x-zed-version";
 
+/// The name of the header used to indicate which edit prediction experiment should be used.
+pub const PREFERRED_EXPERIMENT_HEADER_NAME: &str = "x-zed-preferred-experiment";
+
 /// The name of the header used to indicate when a request failed due to an
 /// expired LLM token.
 ///

crates/edit_prediction/src/edit_prediction.rs 🔗 

@@ -9,7 +9,8 @@ use cloud_llm_client::predict_edits_v3::{
 use cloud_llm_client::{
     EditPredictionRejectReason, EditPredictionRejection,
     MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
-    PredictEditsRequestTrigger, RejectEditPredictionsBodyRef, ZED_VERSION_HEADER_NAME,
+    PREFERRED_EXPERIMENT_HEADER_NAME, PredictEditsRequestTrigger, RejectEditPredictionsBodyRef,
+    ZED_VERSION_HEADER_NAME,
 };
 use collections::{HashMap, HashSet};
 use copilot::{Copilot, Reinstall, SignIn, SignOut};
@@ -2586,6 +2587,7 @@ impl EditPredictionStore {
 
     pub(crate) async fn send_v3_request(
         input: ZetaPromptInput,
+        preferred_experiment: Option<String>,
         client: Arc<Client>,
         llm_token: LlmApiToken,
         organization_id: Option<OrganizationId>,
@@ -2604,11 +2606,16 @@ impl EditPredictionStore {
 
         Self::send_api_request(
             |builder| {
-                let req = builder
+                let builder = builder
                     .uri(url.as_ref())
                     .header("Content-Encoding", "zstd")
-                    .header(PREDICT_EDITS_MODE_HEADER_NAME, mode.as_ref())
-                    .body(compressed.clone().into());
+                    .header(PREDICT_EDITS_MODE_HEADER_NAME, mode.as_ref());
+                let builder = if let Some(preferred_experiment) = preferred_experiment.as_deref() {
+                    builder.header(PREFERRED_EXPERIMENT_HEADER_NAME, preferred_experiment)
+                } else {
+                    builder
+                };
+                let req = builder.body(compressed.clone().into());
                 Ok(req?)
             },
             client,

crates/edit_prediction/src/edit_prediction_tests.rs 🔗 

@@ -2481,7 +2481,6 @@ async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
             excerpt_start_row: None,
             excerpt_ranges: Default::default(),
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,

crates/edit_prediction/src/fim.rs 🔗 

@@ -87,7 +87,6 @@ pub fn request_prediction(
             cursor_excerpt,
             excerpt_ranges: Default::default(),
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,

crates/edit_prediction/src/mercury.rs 🔗 

@@ -109,7 +109,6 @@ impl Mercury {
                     - excerpt_offset_range.start,
                 cursor_path: full_path.clone(),
                 cursor_excerpt,
-                experiment: None,
                 excerpt_start_row: Some(excerpt_point_range.start.row),
                 excerpt_ranges,
                 syntax_ranges: Some(syntax_ranges),

crates/edit_prediction/src/prediction.rs 🔗 

@@ -155,7 +155,6 @@ mod tests {
                 excerpt_start_row: None,
                 excerpt_ranges: Default::default(),
                 syntax_ranges: None,
-                experiment: None,
                 in_open_source_repo: false,
                 can_collect_data: false,
                 repo_url: None,

crates/edit_prediction/src/zeta.rs 🔗 

@@ -120,7 +120,6 @@ pub fn request_prediction_with_zeta(
                 diagnostic_search_range,
                 excerpt_path,
                 cursor_offset,
-                preferred_experiment,
                 is_open_source,
                 can_collect_data,
                 repo_url,
@@ -274,6 +273,7 @@ pub fn request_prediction_with_zeta(
                     // Use V3 endpoint - server handles model/version selection and suffix stripping
                     let (response, usage) = EditPredictionStore::send_v3_request(
                         prompt_input.clone(),
+                        preferred_experiment.clone(),
                         client,
                         llm_token,
                         organization_id,
@@ -535,7 +535,6 @@ pub fn zeta2_prompt_input(
     diagnostic_search_range: Range<Point>,
     excerpt_path: Arc<Path>,
     cursor_offset: usize,
-    preferred_experiment: Option<String>,
     is_open_source: bool,
     can_collect_data: bool,
     repo_url: Option<String>,
@@ -567,7 +566,6 @@ pub fn zeta2_prompt_input(
         active_buffer_diagnostics,
         excerpt_ranges,
         syntax_ranges: Some(syntax_ranges),
-        experiment: preferred_experiment,
         in_open_source_repo: is_open_source,
         can_collect_data,
         repo_url,

crates/zeta_prompt/src/zeta_prompt.rs 🔗 

@@ -51,9 +51,6 @@ pub struct ZetaPromptInput {
     /// instead of `excerpt_ranges`.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub syntax_ranges: Option<Vec<Range<usize>>>,
-    /// The name of the edit prediction model experiment to use.
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub experiment: Option<String>,
     #[serde(default)]
     pub in_open_source_repo: bool,
     #[serde(default)]
@@ -4503,7 +4500,6 @@ mod tests {
                 ..Default::default()
             },
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,
@@ -4534,7 +4530,6 @@ mod tests {
                 ..Default::default()
             },
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,
@@ -5163,7 +5158,6 @@ mod tests {
                 ..Default::default()
             },
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,
@@ -5228,7 +5222,6 @@ mod tests {
                 ..Default::default()
             },
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,
@@ -5288,7 +5281,6 @@ mod tests {
                 ..Default::default()
             },
             syntax_ranges: None,
-            experiment: None,
             in_open_source_repo: false,
             can_collect_data: false,
             repo_url: None,