edit prediction: Report early-rejected predictions and fix cancel bug (#43585)

Many prediction requests end up being rejected early without ever being
set as the current prediction. Before this change, those cases weren’t
reported as rejections because the `request_prediction_with_*` functions
simply returned `Ok(None)`.

With this update, whenever we get a successful response from the
provider, we will return at least the `id`, allowing it to be properly
reported. The request now also includes a “reject reason,” since the
different variants carry distinct implications for prediction quality.

All of these scenarios are now covered by tests. While adding them, I
also found and fixed a bug where some cancelled predictions were
incorrectly being set as the current one.

Release Notes:

- N/A

---------

Co-authored-by: MrSubidubi <dev@bahn.sh>

crates/cloud_llm_client/src/cloud_llm_client.rs

@@ -200,12 +200,31 @@ pub struct RejectEditPredictionsBody {
     pub rejections: Vec<EditPredictionRejection>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct EditPredictionRejection {
     pub request_id: String,
+    #[serde(default)]
+    pub reason: EditPredictionRejectReason,
     pub was_shown: bool,
 }
 
+#[derive(Default, Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
+pub enum EditPredictionRejectReason {
+    /// New requests were triggered before this one completed
+    Canceled,
+    /// No edits returned
+    Empty,
+    /// Edits returned, but none remained after interpolation
+    InterpolatedEmpty,
+    /// The new prediction was preferred over the current one
+    Replaced,
+    /// The current prediction was preferred over the new one
+    CurrentPreferred,
+    /// The current prediction was discarded
+    #[default]
+    Discarded,
+}
+
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum CompletionMode {

crates/zeta/src/prediction.rs

@@ -5,6 +5,7 @@ use std::{
     time::{Duration, Instant},
 };
 
+use cloud_llm_client::EditPredictionRejectReason;
 use gpui::{AsyncApp, Entity, SharedString};
 use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot};
 use serde::Serialize;
@@ -24,13 +25,71 @@ impl std::fmt::Display for EditPredictionId {
     }
 }
 
+/// A prediction response that was returned from the provider, whether it was ultimately valid or not.
+pub struct EditPredictionResult {
+    pub id: EditPredictionId,
+    pub prediction: Result<EditPrediction, EditPredictionRejectReason>,
+}
+
+impl EditPredictionResult {
+    pub async fn new(
+        id: EditPredictionId,
+        edited_buffer: &Entity<Buffer>,
+        edited_buffer_snapshot: &BufferSnapshot,
+        edits: Arc<[(Range<Anchor>, Arc<str>)]>,
+        buffer_snapshotted_at: Instant,
+        response_received_at: Instant,
+        inputs: EditPredictionInputs,
+        cx: &mut AsyncApp,
+    ) -> Self {
+        if edits.is_empty() {
+            return Self {
+                id,
+                prediction: Err(EditPredictionRejectReason::Empty),
+            };
+        }
+
+        let Some((edits, snapshot, edit_preview_task)) = edited_buffer
+            .read_with(cx, |buffer, cx| {
+                let new_snapshot = buffer.snapshot();
+                let edits: Arc<[_]> =
+                    interpolate_edits(&edited_buffer_snapshot, &new_snapshot, edits)?.into();
+
+                Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
+            })
+            .ok()
+            .flatten()
+        else {
+            return Self {
+                id,
+                prediction: Err(EditPredictionRejectReason::InterpolatedEmpty),
+            };
+        };
+
+        let edit_preview = edit_preview_task.await;
+
+        Self {
+            id: id.clone(),
+            prediction: Ok(EditPrediction {
+                id,
+                edits,
+                snapshot,
+                edit_preview,
+                inputs,
+                buffer: edited_buffer.clone(),
+                buffer_snapshotted_at,
+                response_received_at,
+            }),
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct EditPrediction {
     pub id: EditPredictionId,
     pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
     pub snapshot: BufferSnapshot,
     pub edit_preview: EditPreview,
-    // We keep a reference to the buffer so that we do not need to reload it from disk when applying the prediction.
     pub buffer: Entity<Buffer>,
     pub buffer_snapshotted_at: Instant,
     pub response_received_at: Instant,
@@ -46,40 +105,6 @@ pub struct EditPredictionInputs {
 }
 
 impl EditPrediction {
-    pub async fn new(
-        id: EditPredictionId,
-        edited_buffer: &Entity<Buffer>,
-        edited_buffer_snapshot: &BufferSnapshot,
-        edits: Arc<[(Range<Anchor>, Arc<str>)]>,
-        buffer_snapshotted_at: Instant,
-        response_received_at: Instant,
-        inputs: EditPredictionInputs,
-        cx: &mut AsyncApp,
-    ) -> Option<Self> {
-        let (edits, snapshot, edit_preview_task) = edited_buffer
-            .read_with(cx, |buffer, cx| {
-                let new_snapshot = buffer.snapshot();
-                let edits: Arc<[_]> =
-                    interpolate_edits(&edited_buffer_snapshot, &new_snapshot, edits)?.into();
-
-                Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
-            })
-            .ok()??;
-
-        let edit_preview = edit_preview_task.await;
-
-        Some(EditPrediction {
-            id,
-            edits,
-            snapshot,
-            edit_preview,
-            inputs,
-            buffer: edited_buffer.clone(),
-            buffer_snapshotted_at,
-            response_received_at,
-        })
-    }
-
     pub fn interpolate(
         &self,
         new_snapshot: &TextBufferSnapshot,

crates/zeta/src/provider.rs

@@ -1,6 +1,7 @@
 use std::{cmp, sync::Arc, time::Duration};
 
 use client::{Client, UserStore};
+use cloud_llm_client::EditPredictionRejectReason;
 use edit_prediction::{DataCollectionState, Direction, EditPredictionProvider};
 use gpui::{App, Entity, prelude::*};
 use language::ToPoint as _;
@@ -132,7 +133,11 @@ impl EditPredictionProvider for ZetaEditPredictionProvider {
 
     fn discard(&mut self, cx: &mut Context<Self>) {
         self.zeta.update(cx, |zeta, cx| {
-            zeta.discard_current_prediction(&self.project, cx);
+            zeta.reject_current_prediction(
+                EditPredictionRejectReason::Discarded,
+                &self.project,
+                cx,
+            );
         });
     }
 
@@ -169,7 +174,11 @@ impl EditPredictionProvider for ZetaEditPredictionProvider {
 
         let Some(edits) = prediction.interpolate(&snapshot) else {
             self.zeta.update(cx, |zeta, cx| {
-                zeta.discard_current_prediction(&self.project, cx);
+                zeta.reject_current_prediction(
+                    EditPredictionRejectReason::InterpolatedEmpty,
+                    &self.project,
+                    cx,
+                );
             });
             return None;
         };

crates/zeta/src/sweep_ai.rs

@@ -18,7 +18,7 @@ use std::{
     time::Instant,
 };
 
-use crate::{EditPrediction, EditPredictionId, EditPredictionInputs};
+use crate::{EditPredictionId, EditPredictionInputs, prediction::EditPredictionResult};
 
 const SWEEP_API_URL: &str = "https://autocomplete.sweep.dev/backend/next_edit_autocomplete";
 
@@ -45,7 +45,7 @@ impl SweepAi {
         recent_paths: &VecDeque<ProjectPath>,
         diagnostic_search_range: Range<Point>,
         cx: &mut App,
-    ) -> Task<Result<Option<EditPrediction>>> {
+    ) -> Task<Result<Option<EditPredictionResult>>> {
         let debug_info = self.debug_info.clone();
         let Some(api_token) = self.api_token.clone() else {
             return Task::ready(Ok(None));
@@ -242,8 +242,8 @@ impl SweepAi {
 
         cx.spawn(async move |cx| {
             let (id, edits, old_snapshot, response_received_at, inputs) = result.await?;
-            anyhow::Ok(
-                EditPrediction::new(
+            anyhow::Ok(Some(
+                EditPredictionResult::new(
                     EditPredictionId(id.into()),
                     &buffer,
                     &old_snapshot,
@@ -254,7 +254,7 @@ impl SweepAi {
                     cx,
                 )
                 .await,
-            )
+            ))
         })
     }
 }

crates/zeta/src/zeta1.rs

@@ -4,7 +4,7 @@ use std::{fmt::Write, ops::Range, path::Path, sync::Arc, time::Instant};
 
 use crate::{
     EditPredictionId, ZedUpdateRequiredError, Zeta,
-    prediction::{EditPrediction, EditPredictionInputs},
+    prediction::{EditPredictionInputs, EditPredictionResult},
 };
 use anyhow::{Context as _, Result};
 use cloud_llm_client::{
@@ -36,7 +36,7 @@ pub(crate) fn request_prediction_with_zeta1(
     position: language::Anchor,
     events: Vec<Arc<Event>>,
     cx: &mut Context<Zeta>,
-) -> Task<Result<Option<EditPrediction>>> {
+) -> Task<Result<Option<EditPredictionResult>>> {
     let buffer = buffer.clone();
     let buffer_snapshotted_at = Instant::now();
     let client = zeta.client.clone();
@@ -216,7 +216,7 @@ pub(crate) fn request_prediction_with_zeta1(
             );
         }
 
-        edit_prediction
+        edit_prediction.map(Some)
     })
 }
 
@@ -229,7 +229,7 @@ fn process_completion_response(
     buffer_snapshotted_at: Instant,
     received_response_at: Instant,
     cx: &AsyncApp,
-) -> Task<Result<Option<EditPrediction>>> {
+) -> Task<Result<EditPredictionResult>> {
     let snapshot = snapshot.clone();
     let request_id = prediction_response.request_id;
     let output_excerpt = prediction_response.output_excerpt;
@@ -246,8 +246,9 @@ fn process_completion_response(
             .await?
             .into();
 
-        Ok(EditPrediction::new(
-            EditPredictionId(request_id.into()),
+        let id = EditPredictionId(request_id.into());
+        Ok(EditPredictionResult::new(
+            id,
             &buffer,
             &snapshot,
             edits,

crates/zeta/src/zeta_tests.rs

@@ -538,7 +538,7 @@ async fn run_edit_prediction(
     let prediction_task = zeta.update(cx, |zeta, cx| {
         zeta.request_prediction(&project, buffer, cursor, cx)
     });
-    prediction_task.await.unwrap().unwrap()
+    prediction_task.await.unwrap().unwrap().prediction.unwrap()
 }
 
 async fn make_test_zeta(

crates/zeta_cli/src/predict.rs

@@ -235,7 +235,10 @@ pub async fn perform_predict(
     let mut result = Arc::into_inner(result).unwrap().into_inner().unwrap();
 
     result.diff = prediction
-        .and_then(|prediction| prediction.edit_preview.as_unified_diff(&prediction.edits))
+        .and_then(|prediction| {
+            let prediction = prediction.prediction.ok()?;
+            prediction.edit_preview.as_unified_diff(&prediction.edits)
+        })
         .unwrap_or_default();
 
     anyhow::Ok(result)