From a6abc1feec2fdc26337db583fa7339b09116e578 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 26 Jan 2026 16:54:59 -0800
Subject: [PATCH] Send some traffic to zeta2 for testing (#47710)

Release Notes:

- N/A
---
 .../cloud_llm_client/src/cloud_llm_client.rs  |  1 +
 crates/edit_prediction/src/capture_example.rs |  5 +++++
 crates/edit_prediction/src/edit_prediction.rs | 20 +++++++++++++++++--
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/crates/cloud_llm_client/src/cloud_llm_client.rs b/crates/cloud_llm_client/src/cloud_llm_client.rs
index 2c5b2649000bb071b9d206d9d2c204f1eea9bda1..3920d8aad9dfc48e484f8b788d1e1a853591fac1 100644
--- a/crates/cloud_llm_client/src/cloud_llm_client.rs
+++ b/crates/cloud_llm_client/src/cloud_llm_client.rs
@@ -176,6 +176,7 @@ pub struct PredictEditsBody {
 
 #[derive(Default, Debug, Clone, Copy, Serialize, Deserialize)]
 pub enum PredictEditsRequestTrigger {
+    Testing,
     Diagnostics,
     Cli,
     #[default]
diff --git a/crates/edit_prediction/src/capture_example.rs b/crates/edit_prediction/src/capture_example.rs
index f5d6088c7cb70494ea89418fe2e0dbcd79ec0b57..a53862081359536bba498bc615e91d52935054ea 100644
--- a/crates/edit_prediction/src/capture_example.rs
+++ b/crates/edit_prediction/src/capture_example.rs
@@ -13,6 +13,7 @@ use std::{collections::hash_map, fmt::Write as _, ops::Range, path::Path, sync::
 use text::{BufferSnapshot as TextBufferSnapshot, Point};
 
 pub(crate) const DEFAULT_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS: u16 = 10;
+pub(crate) const ZETA2_TESTING_RATE_PER_10K_PREDICTION: u16 = 500;
 
 pub fn capture_example(
     project: Entity<Project>,
@@ -240,6 +241,10 @@ pub(crate) fn should_sample_edit_prediction_example_capture(cx: &App) -> bool {
         && rand::random::<u16>() % 10_000 < capture_rate
 }
 
+pub(crate) fn should_send_testing_zeta2_request() -> bool {
+    rand::random::<u16>() % 10_000 < ZETA2_TESTING_RATE_PER_10K_PREDICTION
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs
index 092bacff4067f91841ee644b8866ca120248630a..0ff03de6154878ef685988b9c583e1fc17d6d26a 100644
--- a/crates/edit_prediction/src/edit_prediction.rs
+++ b/crates/edit_prediction/src/edit_prediction.rs
@@ -71,7 +71,9 @@ pub mod zeta2;
 #[cfg(test)]
 mod edit_prediction_tests;
 
-use crate::capture_example::should_sample_edit_prediction_example_capture;
+use crate::capture_example::{
+    should_sample_edit_prediction_example_capture, should_send_testing_zeta2_request,
+};
 use crate::license_detection::LicenseDetectionWatcher;
 use crate::mercury::Mercury;
 use crate::onboarding_modal::ZedPredictModal;
@@ -192,6 +194,7 @@ pub enum EditPredictionModel {
     Mercury,
 }
 
+#[derive(Clone)]
 pub struct EditPredictionModelInput {
     project: Entity<Project>,
     buffer: Entity<Buffer>,
@@ -1816,7 +1819,20 @@ impl EditPredictionStore {
             }
         }
         let task = match self.edit_prediction_model {
-            EditPredictionModel::Zeta1 => zeta1::request_prediction_with_zeta1(self, inputs, cx),
+            EditPredictionModel::Zeta1 => {
+                if should_send_testing_zeta2_request() {
+                    let mut zeta2_inputs = inputs.clone();
+                    zeta2_inputs.trigger = PredictEditsRequestTrigger::Testing;
+                    zeta2::request_prediction_with_zeta2(
+                        self,
+                        zeta2_inputs,
+                        Default::default(),
+                        cx,
+                    )
+                    .detach();
+                }
+                zeta1::request_prediction_with_zeta1(self, inputs, cx)
+            }
             EditPredictionModel::Zeta2 { version } => {
                 zeta2::request_prediction_with_zeta2(self, inputs, version, cx)
             }