From 8d11484a35252646543f0a4b758f539f531702cc Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Mon, 6 Jul 2026 17:12:33 -0500 Subject: [PATCH] Split interpolate failure rejection reason (#60499) ## Summary - Added `InterpolateFailed` as a distinct edit prediction rejection reason. - Kept `InterpolatedEmpty` for successful interpolation that leaves no edits. - Updated interpolation callers to classify failed interpolation separately from empty edit vectors. Release Notes: - N/A --- .../cloud_llm_client/src/cloud_llm_client.rs | 2 + crates/codestral/src/codestral.rs | 3 +- .../src/edit_prediction_tests.rs | 72 ++++++++++++++++++- crates/edit_prediction/src/prediction.rs | 46 +++++++----- .../src/zed_edit_prediction_delegate.rs | 11 ++- .../src/edit_prediction_types.rs | 2 +- 6 files changed, 112 insertions(+), 24 deletions(-) diff --git a/crates/cloud_llm_client/src/cloud_llm_client.rs b/crates/cloud_llm_client/src/cloud_llm_client.rs index a33a2e8cc1..2bc734cc10 100644 --- a/crates/cloud_llm_client/src/cloud_llm_client.rs +++ b/crates/cloud_llm_client/src/cloud_llm_client.rs @@ -196,6 +196,8 @@ pub enum EditPredictionRejectReason { Empty, /// Edits returned, but none remained after interpolation InterpolatedEmpty, + /// Edits returned, but could not be interpolated after buffer changes + InterpolateFailed, /// The new prediction was preferred over the current one Replaced, /// The current prediction was preferred over the new one diff --git a/crates/codestral/src/codestral.rs b/crates/codestral/src/codestral.rs index 64de772aec..c5889ae666 100644 --- a/crates/codestral/src/codestral.rs +++ b/crates/codestral/src/codestral.rs @@ -82,9 +82,10 @@ struct CurrentCompletion { impl CurrentCompletion { /// Attempts to adjust the edits based on changes made to the buffer since the completion was generated. - /// Returns None if the user's edits conflict with the predicted edits. + /// Returns None if no predicted edits remain or the user's edits conflict with the predicted edits. fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option, Arc)>> { edit_prediction_types::interpolate_edits(&self.snapshot, new_snapshot, &self.edits) + .filter(|edits| !edits.is_empty()) } } diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index a2a196d135..dd8b866a6b 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -1389,7 +1389,7 @@ async fn test_interpolated_empty(cx: &mut TestAppContext) { let (request, respond_tx) = requests.predict.next().await.unwrap(); buffer.update(cx, |buffer, cx| { - buffer.set_text("Hello!\nHow are you?\nBye", cx); + buffer.edit([(10..10, " are you?")], None, cx); }); let mut response = model_response(&request, SIMPLE_DIFF); @@ -1412,7 +1412,6 @@ async fn test_interpolated_empty(cx: &mut TestAppContext) { assert!(shown_predictions[0].editable_range.is_some()); }); - // prediction is reported as rejected let (reject_request, _) = requests.reject.next().await.unwrap(); assert_eq!( @@ -1427,6 +1426,75 @@ async fn test_interpolated_empty(cx: &mut TestAppContext) { ); } +#[gpui::test] +async fn test_interpolate_failed(cx: &mut TestAppContext) { + let (ep_store, mut requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "foo.md": "Hello!\nHow\nBye\n" + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.md"), cx).unwrap(); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let position = snapshot.anchor_before(language::Point::new(1, 3)); + + ep_store.update(cx, |ep_store, cx| { + ep_store.refresh_prediction_from_buffer( + project.clone(), + buffer.clone(), + position, + EditPredictionRequestTrigger::Other, + cx, + ); + }); + + let (request, respond_tx) = requests.predict.next().await.unwrap(); + + buffer.update(cx, |buffer, cx| { + buffer.edit([(10..10, " is it?")], None, cx); + }); + + let mut response = model_response(&request, SIMPLE_DIFF); + response.model_version = Some("zeta2:test-interpolate-failed".to_string()); + let id = response.request_id.clone(); + respond_tx.send(response).unwrap(); + + cx.run_until_parked(); + + ep_store.update(cx, |ep_store, cx| { + assert!( + ep_store + .prediction_at(&buffer, None, &project, cx) + .is_none() + ); + assert!(ep_store.rateable_predictions().next().is_none()); + }); + + let (reject_request, _) = requests.reject.next().await.unwrap(); + + assert_eq!( + &reject_request.rejections, + &[EditPredictionRejection { + request_id: id, + reason: EditPredictionRejectReason::InterpolateFailed, + was_shown: false, + model_version: Some("zeta2:test-interpolate-failed".to_string()), + e2e_latency_ms: Some(0), + }] + ); +} + const SIMPLE_DIFF: &str = indoc! { r" --- a/root/foo.md +++ b/root/foo.md diff --git a/crates/edit_prediction/src/prediction.rs b/crates/edit_prediction/src/prediction.rs index ca6379fa6e..0f7896bc88 100644 --- a/crates/edit_prediction/src/prediction.rs +++ b/crates/edit_prediction/src/prediction.rs @@ -49,27 +49,35 @@ impl EditPredictionResult { e2e_latency: std::time::Duration, cx: &mut AsyncApp, ) -> Self { - let (edits, new_snapshot) = (!edits.is_empty()) - .then(|| { - edited_buffer.read_with(cx, |buffer, _cx| { - let new_snapshot = buffer.snapshot(); - let edits: Arc<[(Range, Arc)]> = - interpolate_edits(&edited_buffer_snapshot, &new_snapshot, &edits) - .map(Arc::from) - .unwrap_or_default(); - let snapshot = (!edits.is_empty()).then_some(new_snapshot); - (Some(edits), snapshot) - }) + let (edits, reject_reason, new_snapshot): ( + Arc<[(Range, Arc)]>, + Option, + Option, + ) = if edits.is_empty() { + ( + Arc::default(), + Some(EditPredictionRejectReason::Empty), + None, + ) + } else { + edited_buffer.read_with(cx, |buffer, _cx| { + let new_snapshot = buffer.snapshot(); + match interpolate_edits(&edited_buffer_snapshot, &new_snapshot, &edits) { + Some(edits) if edits.is_empty() => ( + Arc::default(), + Some(EditPredictionRejectReason::InterpolatedEmpty), + None, + ), + Some(edits) => (Arc::from(edits), None, Some(new_snapshot)), + None => ( + Arc::default(), + Some(EditPredictionRejectReason::InterpolateFailed), + None, + ), + } }) - .unwrap_or_default(); - let snapshot = new_snapshot.unwrap_or_else(|| edited_buffer_snapshot.clone()); - - let reject_reason = match edits.as_ref() { - None => Some(EditPredictionRejectReason::Empty), - Some(edits) if edits.is_empty() => Some(EditPredictionRejectReason::InterpolatedEmpty), - Some(_) => None, }; - let edits = edits.unwrap_or_default(); + let snapshot = new_snapshot.unwrap_or_else(|| edited_buffer_snapshot.clone()); let edit_preview = if !edits.is_empty() { edited_buffer diff --git a/crates/edit_prediction/src/zed_edit_prediction_delegate.rs b/crates/edit_prediction/src/zed_edit_prediction_delegate.rs index c3cb556c7b..b7ae7d629e 100644 --- a/crates/edit_prediction/src/zed_edit_prediction_delegate.rs +++ b/crates/edit_prediction/src/zed_edit_prediction_delegate.rs @@ -222,13 +222,22 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate { let Some(edits) = prediction.interpolate(&snapshot) else { store.reject_current_prediction( - EditPredictionRejectReason::InterpolatedEmpty, + EditPredictionRejectReason::InterpolateFailed, &self.project, cx, ); return None; }; + if edits.is_empty() { + store.reject_current_prediction( + EditPredictionRejectReason::InterpolatedEmpty, + &self.project, + cx, + ); + return None; + } + let cursor_row = cursor_position.to_point(&snapshot).row; let (closest_edit_ix, (closest_edit_range, _)) = edits.iter().enumerate().min_by_key(|(_, (range, _))| { diff --git a/crates/edit_prediction_types/src/edit_prediction_types.rs b/crates/edit_prediction_types/src/edit_prediction_types.rs index a285e8aa70..bcd1f5d8e6 100644 --- a/crates/edit_prediction_types/src/edit_prediction_types.rs +++ b/crates/edit_prediction_types/src/edit_prediction_types.rs @@ -392,5 +392,5 @@ pub fn interpolate_edits( edits.extend(model_edits.cloned()); - if edits.is_empty() { None } else { Some(edits) } + Some(edits) }