zeta: Only send outline and diagnostics when data collection is enabled (#35896)

Michael Sloan created 3 months ago

This data is not currently used by edit predictions - it is only useful
when `can_collect_data == true`.

Release Notes:

- N/A

Change summary

crates/zeta/src/zeta.rs     |  8 ++++++--
crates/zeta_cli/src/main.rs | 38 ++++++++++++++++++++++++--------------
2 files changed, 30 insertions(+), 16 deletions(-)

Detailed changes

crates/zeta/src/zeta.rs

@@ -1211,7 +1211,7 @@ pub fn gather_context(
     let local_lsp_store =
         project.and_then(|project| project.read(cx).lsp_store().read(cx).as_local());
     let diagnostic_groups: Vec<(String, serde_json::Value)> =
-        if let Some(local_lsp_store) = local_lsp_store {
+        if can_collect_data && let Some(local_lsp_store) = local_lsp_store {
             snapshot
                 .diagnostic_groups(None)
                 .into_iter()
@@ -1245,7 +1245,11 @@ pub fn gather_context(
                 MAX_CONTEXT_TOKENS,
             );
             let input_events = make_events_prompt();
-            let input_outline = prompt_for_outline(&snapshot);
+            let input_outline = if can_collect_data {
+                prompt_for_outline(&snapshot)
+            } else {
+                String::new()
+            };
             let editable_range = input_excerpt.editable_range.to_offset(&snapshot);
 
             let body = PredictEditsBody {

crates/zeta_cli/src/main.rs

@@ -171,21 +171,31 @@ async fn get_context(
         Some(events) => events.read_to_string().await?,
         None => String::new(),
     };
-    let can_collect_data = false;
+    // Enable gathering extra data not currently needed for edit predictions
+    let can_collect_data = true;
     let git_info = None;
-    cx.update(|cx| {
-        gather_context(
-            project.as_ref(),
-            full_path_str,
-            &snapshot,
-            clipped_cursor,
-            move || events,
-            can_collect_data,
-            git_info,
-            cx,
-        )
-    })?
-    .await
+    let mut gather_context_output = cx
+        .update(|cx| {
+            gather_context(
+                project.as_ref(),
+                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
+    match gather_context_output.as_mut() {
+        Ok(gather_context_output) => gather_context_output.body.can_collect_data = false,
+        Err(_) => {}
+    }
+
+    gather_context_output
 }
 
 pub async fn open_buffer_with_language_server(