From cc5eb2406691765ff624d217bc32b07519941280 Mon Sep 17 00:00:00 2001
From: Michael Sloan <michael@zed.dev>
Date: Tue, 12 Aug 2025 00:47:54 -0600
Subject: [PATCH] zeta: Add latency telemetry for 1% of edit predictions
 (#36020)

Release Notes:

- N/A

Co-authored-by: Oleksiy <oleksiy@zed.dev>
---
 Cargo.lock              |  1 +
 crates/zeta/Cargo.toml  |  3 ++-
 crates/zeta/src/zeta.rs | 24 ++++++++++++++++++++++--
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8d22eeafab..79bce189e2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -20923,6 +20923,7 @@ dependencies = [
  "menu",
  "postage",
  "project",
+ "rand 0.8.5",
  "regex",
  "release_channel",
  "reqwest_client",
diff --git a/crates/zeta/Cargo.toml b/crates/zeta/Cargo.toml
index 9f1d02b790..ee76308ff3 100644
--- a/crates/zeta/Cargo.toml
+++ b/crates/zeta/Cargo.toml
@@ -26,6 +26,7 @@ collections.workspace = true
 command_palette_hooks.workspace = true
 copilot.workspace = true
 db.workspace = true
+edit_prediction.workspace = true
 editor.workspace = true
 feature_flags.workspace = true
 fs.workspace = true
@@ -33,13 +34,13 @@ futures.workspace = true
 gpui.workspace = true
 http_client.workspace = true
 indoc.workspace = true
-edit_prediction.workspace = true
 language.workspace = true
 language_model.workspace = true
 log.workspace = true
 menu.workspace = true
 postage.workspace = true
 project.workspace = true
+rand.workspace = true
 regex.workspace = true
 release_channel.workspace = true
 serde.workspace = true
diff --git a/crates/zeta/src/zeta.rs b/crates/zeta/src/zeta.rs
index 6900082003..1a6a8c2934 100644
--- a/crates/zeta/src/zeta.rs
+++ b/crates/zeta/src/zeta.rs
@@ -429,6 +429,7 @@ impl Zeta {
                 body,
                 editable_range,
             } = gather_task.await?;
+            let done_gathering_context_at = Instant::now();
 
             log::debug!(
                 "Events:\n{}\nExcerpt:\n{:?}",
@@ -481,6 +482,7 @@ impl Zeta {
                 }
             };
 
+            let received_response_at = Instant::now();
             log::debug!("completion response: {}", &response.output_excerpt);
 
             if let Some(usage) = usage {
@@ -492,7 +494,7 @@ impl Zeta {
                 .ok();
             }
 
-            Self::process_completion_response(
+            let edit_prediction = Self::process_completion_response(
                 response,
                 buffer,
                 &snapshot,
@@ -505,7 +507,25 @@ impl Zeta {
                 buffer_snapshotted_at,
                 &cx,
             )
-            .await
+            .await;
+
+            let finished_at = Instant::now();
+
+            // record latency for ~1% of requests
+            if rand::random::<u8>() <= 2 {
+                telemetry::event!(
+                    "Edit Prediction Request",
+                    context_latency = done_gathering_context_at
+                        .duration_since(buffer_snapshotted_at)
+                        .as_millis(),
+                    request_latency = received_response_at
+                        .duration_since(done_gathering_context_at)
+                        .as_millis(),
+                    process_latency = finished_at.duration_since(received_response_at).as_millis()
+                );
+            }
+
+            edit_prediction
         })
     }