From 9743fe7dfd8e66fe125634c41fb65d0a67ee963c Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Tue, 17 Feb 2026 17:43:40 -0800 Subject: [PATCH] Put back logic for passing data collection choice to cloud (#49426) Release Notes: - N/A --- crates/edit_prediction/src/edit_prediction.rs | 4 ++++ crates/edit_prediction/src/edit_prediction_tests.rs | 1 + crates/edit_prediction/src/mercury.rs | 1 + crates/edit_prediction/src/ollama.rs | 2 ++ crates/edit_prediction/src/prediction.rs | 1 + crates/edit_prediction/src/sweep_ai.rs | 1 + crates/edit_prediction/src/zeta2.rs | 5 +++++ crates/edit_prediction_cli/src/format_prompt.rs | 1 + crates/zeta_prompt/src/zeta_prompt.rs | 6 ++++++ 9 files changed, 22 insertions(+) diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 0ff5af7e8d9445f1ff7978d55c1b91fcc36622dd..ea7708c5769fd20b891a9a7c518a4ab7563b6401 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -2209,6 +2209,10 @@ impl EditPredictionStore { .is_some_and(|watcher| watcher.is_project_open_source()) } + pub(crate) fn is_data_collection_enabled(&self, cx: &App) -> bool { + self.data_collection_choice.is_enabled(cx) + } + fn load_data_collection_choice() -> DataCollectionChoice { let choice = KEY_VALUE_STORE .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE) diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index bb5edbdcd4cb667cc622286720812e45a23ca2c0..242a2bf3fff5f0eb87b183ec6c65280cbe75256a 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -1707,6 +1707,7 @@ async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }, buffer_snapshotted_at: Instant::now(), response_received_at: Instant::now(), diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs index 91c33f0fb663fa54cb94b302fb23f3db16378222..9d546ff8272c0d8acd4a22285c3aa069ea4c525a 100644 --- a/crates/edit_prediction/src/mercury.rs +++ b/crates/edit_prediction/src/mercury.rs @@ -100,6 +100,7 @@ impl Mercury { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; let prompt = build_prompt(&inputs); diff --git a/crates/edit_prediction/src/ollama.rs b/crates/edit_prediction/src/ollama.rs index c372c73a01990596db7a7d4551808788739fd9d8..8de90ba67ee2a9ecac3e52d7fee101b8db84c54e 100644 --- a/crates/edit_prediction/src/ollama.rs +++ b/crates/edit_prediction/src/ollama.rs @@ -172,6 +172,7 @@ impl Ollama { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; (prompt, stop_tokens, Some(editable_offset_range), inputs) @@ -201,6 +202,7 @@ impl Ollama { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; let prefix = inputs.cursor_excerpt[..inputs.cursor_offset_in_excerpt].to_string(); diff --git a/crates/edit_prediction/src/prediction.rs b/crates/edit_prediction/src/prediction.rs index 3d87edb14ab775ef7ee8da2a8faa31efb79ec899..750b1a435ae4a7a281ef41973e1f6d0d2158445e 100644 --- a/crates/edit_prediction/src/prediction.rs +++ b/crates/edit_prediction/src/prediction.rs @@ -161,6 +161,7 @@ mod tests { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }, buffer_snapshotted_at: Instant::now(), response_received_at: Instant::now(), diff --git a/crates/edit_prediction/src/sweep_ai.rs b/crates/edit_prediction/src/sweep_ai.rs index eb8ee8fe68c9b4458663e196cfb45e1ffadaa0ce..1253916487894d757c74293c21f4ace1c681cd11 100644 --- a/crates/edit_prediction/src/sweep_ai.rs +++ b/crates/edit_prediction/src/sweep_ai.rs @@ -222,6 +222,7 @@ impl SweepAi { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; send_started_event( diff --git a/crates/edit_prediction/src/zeta2.rs b/crates/edit_prediction/src/zeta2.rs index c9a7847704ae5dc35116079868870bbdf4ee0fdd..2d6e4fdce6a794bb6a64c830196a4fce187b2866 100644 --- a/crates/edit_prediction/src/zeta2.rs +++ b/crates/edit_prediction/src/zeta2.rs @@ -66,6 +66,8 @@ pub fn request_prediction_with_zeta2( && events.iter().all(|event| event.in_open_source_repo()) && related_files.iter().all(|file| file.in_open_source_repo); + let can_collect_data = is_open_source && store.is_data_collection_enabled(cx); + let request_task = cx.background_spawn({ async move { let zeta_version = raw_config @@ -83,6 +85,7 @@ pub fn request_prediction_with_zeta2( zeta_version, preferred_model, is_open_source, + can_collect_data, ); if prompt_input_contains_special_tokens(&prompt_input, zeta_version) { @@ -265,6 +268,7 @@ pub fn zeta2_prompt_input( zeta_format: ZetaFormat, preferred_model: Option, is_open_source: bool, + can_collect_data: bool, ) -> (std::ops::Range, zeta_prompt::ZetaPromptInput) { let cursor_point = cursor_offset.to_point(snapshot); @@ -309,6 +313,7 @@ pub fn zeta2_prompt_input( excerpt_ranges: Some(excerpt_ranges), preferred_model, in_open_source_repo: is_open_source, + can_collect_data, }; (editable_offset_range, prompt_input) } diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs index 702af3091a2d5f292806034ad875af861c484e2f..3797f84ac8f4dddaa66b6a10bd1b68fb46cfd5e5 100644 --- a/crates/edit_prediction_cli/src/format_prompt.rs +++ b/crates/edit_prediction_cli/src/format_prompt.rs @@ -100,6 +100,7 @@ pub async fn run_format_prompt( .captured_prompt_input .as_ref() .map_or(false, |input| input.in_open_source_repo), + can_collect_data: false, }; let prompt = format_zeta_prompt(&input, version); let prefill = zeta_prompt::get_prefill(&input, version); diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 53de7b387ff6a92801e4482eef809f44a23ff7fa..d8376b5ce312db1166667b89d69bfe9476543a8d 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -65,6 +65,8 @@ pub struct ZetaPromptInput { pub preferred_model: Option, #[serde(default)] pub in_open_source_repo: bool, + #[serde(default)] + pub can_collect_data: bool, } #[derive( @@ -1105,6 +1107,7 @@ mod tests { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, } } @@ -1481,6 +1484,7 @@ mod tests { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; let prompt = zeta1::format_zeta1_from_input(&input, 15..41, 0..excerpt.len()); @@ -1535,6 +1539,7 @@ mod tests { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; let prompt = zeta1::format_zeta1_from_input(&input, 0..28, 0..28); @@ -1584,6 +1589,7 @@ mod tests { excerpt_ranges: None, preferred_model: None, in_open_source_repo: false, + can_collect_data: false, }; let prompt = zeta1::format_zeta1_from_input(&input, editable_range, context_range);