ep: Always collect for staff (#46176)

Agus Zubiaga , Max Brunsfeld , and Oleksiy Syvokon created 4 days ago

Release Notes:

- N/A

Co-authored-by: Max Brunsfeld <[e-mail]>
Co-authored-by: Oleksiy Syvokon <[e-mail]>

Change summary

crates/edit_prediction/src/edit_prediction.rs              | 23 +++----
crates/edit_prediction/src/zed_edit_prediction_delegate.rs |  2 
crates/edit_prediction/src/zeta1.rs                        |  2 
crates/edit_prediction_ui/src/edit_prediction_button.rs    |  1 
4 files changed, 13 insertions(+), 15 deletions(-)

Detailed changes

crates/edit_prediction/src/edit_prediction.rs

@@ -1642,7 +1642,7 @@ impl EditPredictionStore {
         let can_collect_example = snapshot
             .file()
             .is_some_and(|file| self.can_collect_file(&project, file, cx))
-            && self.can_collect_events(&inputs.events);
+            && self.can_collect_events(&inputs.events, cx);
 
         if can_collect_example && should_sample_edit_prediction_example_capture(cx) {
             let events_for_capture =
@@ -2010,11 +2010,11 @@ impl EditPredictionStore {
     }
 
     fn can_collect_file(&self, project: &Entity<Project>, file: &Arc<dyn File>, cx: &App) -> bool {
-        self.data_collection_choice.is_enabled() && self.is_file_open_source(project, file, cx)
+        self.data_collection_choice.is_enabled(cx) && self.is_file_open_source(project, file, cx)
     }
 
-    fn can_collect_events(&self, events: &[Arc<zeta_prompt::Event>]) -> bool {
-        if !self.data_collection_choice.is_enabled() {
+    fn can_collect_events(&self, events: &[Arc<zeta_prompt::Event>], cx: &App) -> bool {
+        if !self.data_collection_choice.is_enabled(cx) {
             return false;
         }
         events.iter().all(|event| {
@@ -2048,10 +2048,11 @@ impl EditPredictionStore {
     fn toggle_data_collection_choice(&mut self, cx: &mut Context<Self>) {
         self.data_collection_choice = self.data_collection_choice.toggle();
         let new_choice = self.data_collection_choice;
+        let is_enabled = new_choice.is_enabled(cx);
         db::write_and_log(cx, move || {
             KEY_VALUE_STORE.write_kvp(
                 ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
-                new_choice.is_enabled().to_string(),
+                is_enabled.to_string(),
             )
         });
     }
@@ -2139,20 +2140,16 @@ pub enum DataCollectionChoice {
 }
 
 impl DataCollectionChoice {
-    pub fn is_enabled(self) -> bool {
+    pub fn is_enabled(self, cx: &App) -> bool {
+        if cx.is_staff() {
+            return true;
+        }
         match self {
             Self::Enabled => true,
             Self::NotAnswered | Self::Disabled => false,
         }
     }
 
-    pub fn is_answered(self) -> bool {
-        match self {
-            Self::Enabled | Self::Disabled => true,
-            Self::NotAnswered => false,
-        }
-    }
-
     #[must_use]
     pub fn toggle(&self) -> DataCollectionChoice {
         match self {

crates/edit_prediction/src/zed_edit_prediction_delegate.rs

@@ -66,7 +66,7 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate {
                 self.store
                     .read(cx)
                     .is_file_open_source(&self.project, file, cx);
-            if self.store.read(cx).data_collection_choice.is_enabled() {
+            if self.store.read(cx).data_collection_choice.is_enabled(cx) {
                 DataCollectionState::Enabled {
                     is_project_open_source,
                 }

crates/edit_prediction/src/zeta1.rs

@@ -103,7 +103,7 @@ pub(crate) fn request_prediction_with_zeta1(
         let included_events = &events[events.len() - included_events_count..events.len()];
         body.can_collect_data = can_collect_file
             && this
-                .read_with(cx, |this, _| this.can_collect_events(included_events))
+                .read_with(cx, |this, cx| this.can_collect_events(included_events, cx))
                 .unwrap_or(false);
         if body.can_collect_data {
             body.git_info = git_info;

crates/edit_prediction_ui/src/edit_prediction_button.rs

@@ -778,6 +778,7 @@ impl EditPredictionButton {
                             .toggleable(IconPosition::Start, data_collection.is_enabled())
                             .icon(icon_name)
                             .icon_color(icon_color)
+                            .disabled(cx.is_staff())
                             .documentation_aside(DocumentationSide::Left, move |cx| {
                                 let (msg, label_color, icon_name, icon_color) = match (is_open_source, is_collecting) {
                                     (true, true) => (