zeta: Check whether data collection is allowed for recent edit history (#37680)

Also:

* Adds tests for can_collect_data.
* Temporarily removes collection of diagnostics.

Release Notes:

- Edit Prediction: Fixed a bug where requests were marked eligible for
data collection despite the recent edit history in the request involving
files that may not be open source. The requests affected by this bug
will not be used in training data.

Cargo.lock

@@ -20841,6 +20841,7 @@ dependencies = [
  "language_model",
  "log",
  "menu",
+ "parking_lot",
  "postage",
  "project",
  "rand 0.9.1",

crates/zed/src/zed/edit_prediction_registry.rs

@@ -8,7 +8,7 @@ use settings::SettingsStore;
 use std::{cell::RefCell, rc::Rc, sync::Arc};
 use supermaven::{Supermaven, SupermavenCompletionProvider};
 use ui::Window;
-use zeta::{ProviderDataCollection, ZetaEditPredictionProvider};
+use zeta::ZetaEditPredictionProvider;
 
 pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
@@ -214,11 +214,8 @@ fn assign_edit_prediction_provider(
                     });
                 }
 
-                let data_collection =
-                    ProviderDataCollection::new(zeta.clone(), singleton_buffer, cx);
-
                 let provider =
-                    cx.new(|_| zeta::ZetaEditPredictionProvider::new(zeta, data_collection));
+                    cx.new(|_| zeta::ZetaEditPredictionProvider::new(zeta, singleton_buffer));
 
                 editor.set_edit_prediction_provider(Some(provider), window, cx);
             }

crates/zeta/Cargo.toml

@@ -72,6 +72,7 @@ gpui = { workspace = true, features = ["test-support"] }
 http_client = { workspace = true, features = ["test-support"] }
 indoc.workspace = true
 language = { workspace = true, features = ["test-support"] }
+parking_lot.workspace = true
 reqwest_client = { workspace = true, features = ["test-support"] }
 rpc = { workspace = true, features = ["test-support"] }
 settings = { workspace = true, features = ["test-support"] }

crates/zeta/src/input_excerpt.rs

@@ -1,6 +1,6 @@
 use crate::{
     CURSOR_MARKER, EDITABLE_REGION_END_MARKER, EDITABLE_REGION_START_MARKER, START_OF_FILE_MARKER,
-    tokens_for_bytes,
+    guess_token_count,
 };
 use language::{BufferSnapshot, Point};
 use std::{fmt::Write, ops::Range};
@@ -22,7 +22,7 @@ pub fn excerpt_for_cursor_position(
     let mut remaining_edit_tokens = editable_region_token_limit;
 
     while let Some(parent) = snapshot.syntax_ancestor(scope_range.clone()) {
-        let parent_tokens = tokens_for_bytes(parent.byte_range().len());
+        let parent_tokens = guess_token_count(parent.byte_range().len());
         let parent_point_range = Point::new(
             parent.start_position().row as u32,
             parent.start_position().column as u32,
@@ -99,7 +99,7 @@ fn expand_range(
         if remaining_tokens > 0 && expanded_range.start.row > 0 {
             expanded_range.start.row -= 1;
             let line_tokens =
-                tokens_for_bytes(snapshot.line_len(expanded_range.start.row) as usize);
+                guess_token_count(snapshot.line_len(expanded_range.start.row) as usize);
             remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
             expanded = true;
         }
@@ -107,7 +107,7 @@ fn expand_range(
         if remaining_tokens > 0 && expanded_range.end.row < snapshot.max_point().row {
             expanded_range.end.row += 1;
             expanded_range.end.column = snapshot.line_len(expanded_range.end.row);
-            let line_tokens = tokens_for_bytes(expanded_range.end.column as usize);
+            let line_tokens = guess_token_count(expanded_range.end.column as usize);
             remaining_tokens = remaining_tokens.saturating_sub(line_tokens);
             expanded = true;
         }

crates/zeta/src/license_detection.rs

@@ -358,7 +358,6 @@ impl LicenseDetectionWatcher {
 
 #[cfg(test)]
 mod tests {
-
     use fs::FakeFs;
     use gpui::TestAppContext;
     use serde_json::json;

crates/zeta_cli/src/main.rs

@@ -189,30 +189,17 @@ async fn get_context(
         Some(events) => events.read_to_string().await?,
         None => String::new(),
     };
-    // Enable gathering extra data not currently needed for edit predictions
-    let can_collect_data = true;
-    let git_info = None;
-    let mut gather_context_output = cx
-        .update(|cx| {
-            gather_context(
-                &project,
-                full_path_str,
-                &snapshot,
-                clipped_cursor,
-                move || events,
-                can_collect_data,
-                git_info,
-                cx,
-            )
-        })?
-        .await;
-
-    // Disable data collection for these requests, as this is currently just used for evals
-    if let Ok(gather_context_output) = gather_context_output.as_mut() {
-        gather_context_output.body.can_collect_data = false
-    }
-
-    gather_context_output
+    let prompt_for_events = move || (events, 0);
+    cx.update(|cx| {
+        gather_context(
+            full_path_str,
+            &snapshot,
+            clipped_cursor,
+            prompt_for_events,
+            cx,
+        )
+    })?
+    .await
 }
 
 pub async fn open_buffer_with_language_server(