ep: Fix e2e_latency tracking (#51852)

Oleksiy Syvokon and Ben Kunkle created 3 weeks ago

We were always sending zero before

Release Notes:

- N/A

---------

Co-authored-by: Ben Kunkle <[e-mail]>

Change summary

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(-)

Detailed changes

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;
             })

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()),
             }
         ]
     );

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,