From b0525a26a684ab2822baaa6c94ee5acd7f50a1d4 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Sun, 16 Nov 2025 16:51:13 -0300 Subject: [PATCH] Report automatically discarded zeta predictions (#42761) We weren't reporting predictions that were generated but never made it out of the provider, such as predictions that failed to interpolate, and those that are cancelled because another request completes before it. Release Notes: - N/A --------- Co-authored-by: Max Brunsfeld Co-authored-by: Ben Kunkle --- .../zed/src/zed/edit_prediction_registry.rs | 3 +- crates/zeta/src/zeta.rs | 105 ++++++++++++------ 2 files changed, 74 insertions(+), 34 deletions(-) diff --git a/crates/zed/src/zed/edit_prediction_registry.rs b/crates/zed/src/zed/edit_prediction_registry.rs index fd16478b5a..74b6687f62 100644 --- a/crates/zed/src/zed/edit_prediction_registry.rs +++ b/crates/zed/src/zed/edit_prediction_registry.rs @@ -251,11 +251,12 @@ fn assign_edit_prediction_provider( }); } - let provider = cx.new(|_| { + let provider = cx.new(|cx| { zeta::ZetaEditPredictionProvider::new( zeta, project.clone(), singleton_buffer, + cx, ) }); editor.set_edit_prediction_provider(Some(provider), window, cx); diff --git a/crates/zeta/src/zeta.rs b/crates/zeta/src/zeta.rs index 577ca77c13..c2ef5cb826 100644 --- a/crates/zeta/src/zeta.rs +++ b/crates/zeta/src/zeta.rs @@ -1283,6 +1283,7 @@ struct CurrentEditPrediction { buffer_id: EntityId, completion: EditPrediction, was_shown: bool, + was_accepted: bool, } impl CurrentEditPrediction { @@ -1310,7 +1311,7 @@ impl CurrentEditPrediction { struct PendingCompletion { id: usize, - _task: Task<()>, + task: Task<()>, } #[derive(Debug, Clone, Copy)] @@ -1386,6 +1387,7 @@ pub struct ZetaEditPredictionProvider { zeta: Entity, singleton_buffer: Option>, pending_completions: ArrayVec, + canceled_completions: HashMap>, next_pending_completion_id: usize, current_completion: Option, last_request_timestamp: Instant, @@ -1399,17 +1401,34 @@ impl ZetaEditPredictionProvider { zeta: Entity, project: Entity, singleton_buffer: Option>, + cx: &mut Context, ) -> Self { + cx.on_release(|this, cx| { + this.take_current_edit_prediction(cx); + }) + .detach(); + Self { zeta, singleton_buffer, pending_completions: ArrayVec::new(), + canceled_completions: HashMap::default(), next_pending_completion_id: 0, current_completion: None, last_request_timestamp: Instant::now(), project, } } + + fn take_current_edit_prediction(&mut self, cx: &mut App) { + if let Some(completion) = self.current_completion.take() { + if !completion.was_accepted { + self.zeta.update(cx, |zeta, cx| { + zeta.discard_completion(completion.completion.id, completion.was_shown, cx); + }); + } + } + } } impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider { @@ -1531,42 +1550,65 @@ impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider { buffer_id: buffer.entity_id(), completion, was_shown: false, + was_accepted: false, }) }) } Err(error) => Err(error), }; + + let discarded = this + .update(cx, |this, cx| { + if this + .pending_completions + .first() + .is_some_and(|completion| completion.id == pending_completion_id) + { + this.pending_completions.remove(0); + } else { + if let Some(discarded) = this.pending_completions.drain(..).next() { + this.canceled_completions + .insert(discarded.id, discarded.task); + } + } + + let canceled = this.canceled_completions.remove(&pending_completion_id); + + if canceled.is_some() + && let Ok(Some(new_completion)) = &completion + { + this.zeta.update(cx, |zeta, cx| { + zeta.discard_completion(new_completion.completion.id, false, cx); + }); + return true; + } + + cx.notify(); + false + }) + .ok() + .unwrap_or(true); + + if discarded { + return; + } + let Some(new_completion) = completion .context("edit prediction failed") .log_err() .flatten() else { - this.update(cx, |this, cx| { - if this.pending_completions[0].id == pending_completion_id { - this.pending_completions.remove(0); - } else { - this.pending_completions.clear(); - } - - cx.notify(); - }) - .ok(); return; }; this.update(cx, |this, cx| { - if this.pending_completions[0].id == pending_completion_id { - this.pending_completions.remove(0); - } else { - this.pending_completions.clear(); - } - if let Some(old_completion) = this.current_completion.as_ref() { let snapshot = buffer.read(cx).snapshot(); if new_completion.should_replace_completion(old_completion, &snapshot) { this.zeta.update(cx, |zeta, cx| { zeta.completion_shown(&new_completion.completion, cx); }); + this.take_current_edit_prediction(cx); this.current_completion = Some(new_completion); } } else { @@ -1586,13 +1628,16 @@ impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider { if self.pending_completions.len() <= 1 { self.pending_completions.push(PendingCompletion { id: pending_completion_id, - _task: task, + task, }); } else if self.pending_completions.len() == 2 { - self.pending_completions.pop(); + if let Some(discarded) = self.pending_completions.pop() { + self.canceled_completions + .insert(discarded.id, discarded.task); + } self.pending_completions.push(PendingCompletion { id: pending_completion_id, - _task: task, + task, }); } } @@ -1608,14 +1653,12 @@ impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider { } fn accept(&mut self, cx: &mut Context) { - let completion_id = self - .current_completion - .as_ref() - .map(|completion| completion.completion.id); - if let Some(completion_id) = completion_id { + let completion = self.current_completion.as_mut(); + if let Some(completion) = completion { + completion.was_accepted = true; self.zeta .update(cx, |zeta, cx| { - zeta.accept_edit_prediction(completion_id, cx) + zeta.accept_edit_prediction(completion.completion.id, cx) }) .detach(); } @@ -1624,11 +1667,7 @@ impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider { fn discard(&mut self, cx: &mut Context) { self.pending_completions.clear(); - if let Some(completion) = self.current_completion.take() { - self.zeta.update(cx, |zeta, cx| { - zeta.discard_completion(completion.completion.id, completion.was_shown, cx); - }); - } + self.take_current_edit_prediction(cx); } fn did_show(&mut self, _cx: &mut Context) { @@ -1651,13 +1690,13 @@ impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider { // Invalidate previous completion if it was generated for a different buffer. if *buffer_id != buffer.entity_id() { - self.current_completion.take(); + self.take_current_edit_prediction(cx); return None; } let buffer = buffer.read(cx); let Some(edits) = completion.interpolate(&buffer.snapshot()) else { - self.current_completion.take(); + self.take_current_edit_prediction(cx); return None; };