From 7cf7dcac513322db1de5f644eb9ecca99e2f84db Mon Sep 17 00:00:00 2001
From: Oleksiy Syvokon <oleksiy@zed.dev>
Date: Wed, 18 Mar 2026 21:20:28 +0200
Subject: [PATCH] ep: Fix e2e_latency tracking (#51852)

We were always sending zero before

Release Notes:

- N/A

---------

Co-authored-by: Ben Kunkle <ben@zed.dev>
---
 crates/edit_prediction/src/edit_prediction.rs       | 5 +++--
 crates/edit_prediction/src/edit_prediction_tests.rs | 4 +++-
 crates/edit_prediction/src/zeta.rs                  | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs
index b8a631955ec4dc8f5494379eca3d9cffb4c2d05b..da2392d757675bb1993d2dacdf963fe24aa91430 100644
--- a/crates/edit_prediction/src/edit_prediction.rs
+++ b/crates/edit_prediction/src/edit_prediction.rs
@@ -2151,11 +2151,12 @@ impl EditPredictionStore {
                     let project_state = this.get_or_init_project(&project, cx);
                     let throttle = *select_throttle(project_state, request_trigger);
 
+                    let now = cx.background_executor().now();
                     throttle.and_then(|(last_entity, last_timestamp)| {
                         if throttle_entity != last_entity {
                             return None;
                         }
-                        (last_timestamp + throttle_timeout).checked_duration_since(Instant::now())
+                        (last_timestamp + throttle_timeout).checked_duration_since(now)
                     })
                 })
                 .ok()
@@ -2183,7 +2184,7 @@ impl EditPredictionStore {
                     return;
                 }
 
-                let new_refresh = (throttle_entity, Instant::now());
+                let new_refresh = (throttle_entity, cx.background_executor().now());
                 *select_throttle(project_state, request_trigger) = Some(new_refresh);
                 is_cancelled = false;
             })
diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs
index 5daa7ee4a0dea1384e002acefe1fb4b47d0d5f91..257dd255cbb6ae3a2ddb02019e56886d41676de2 100644
--- a/crates/edit_prediction/src/edit_prediction_tests.rs
+++ b/crates/edit_prediction/src/edit_prediction_tests.rs
@@ -1807,7 +1807,9 @@ async fn test_cancel_second_on_third_request(cx: &mut TestAppContext) {
                 reason: EditPredictionRejectReason::Replaced,
                 was_shown: false,
                 model_version: None,
-                e2e_latency_ms: Some(0),
+                // 2 throttle waits (for 2nd and 3rd requests) elapsed
+                // between this request's start and response.
+                e2e_latency_ms: Some(2 * EditPredictionStore::THROTTLE_TIMEOUT.as_millis()),
             }
         ]
     );
diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs
index e7d38df5c8e99b86303ca72a715e10acf22eb9b1..fdfe3ebcf06c8319f5ce00066fa279d79eda7eea 100644
--- a/crates/edit_prediction/src/zeta.rs
+++ b/crates/edit_prediction/src/zeta.rs
@@ -365,10 +365,10 @@ pub fn request_prediction_with_zeta(
     });
 
     cx.spawn(async move |this, cx| {
-        let request_duration = cx.background_executor().now() - request_start;
         let Some((id, prediction)) = handle_api_response(&this, request_task.await, cx)? else {
             return Ok(None);
         };
+        let request_duration = cx.background_executor().now() - request_start;
 
         let Some(Prediction {
             prompt_input: inputs,