From bb807e1e3ad0c65c2b34703d2bdfd0b09e1f88b2 Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Sat, 20 Jun 2026 17:40:15 -0500 Subject: [PATCH] ep: Track settled telemetry for empty predictions (#59645) # Objective ## Solution - Describe the solution used to achieve the objective above. ## Testing - Did you test these changes? If so, how? - Are there any parts that need more testing? - How can other people (reviewers) test your changes? Is there anything specific they need to know? - If relevant, what platforms did you test these changes on, and are there any important ones you can't test? ## Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards ([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) and [icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md) guidelines) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable ## Showcase > This section is optional. If this PR does not include a visual change or does not add a new user-facing feature, you can delete this section. - Help others understand the result of this PR by showcasing your awesome work! - If this PR includes a visual change, consider adding a screenshot, GIF, or video - A before/after comparison is very useful for changes to existing features! While a showcase should aim to be brief and digestible, you can use a toggleable section to save space on longer showcases:
Click to view showcase My super cool demos here
--- Release Notes: - N/A or Added/Fixed/Improved ... --- crates/edit_prediction/src/edit_prediction.rs | 163 +++++++++--------- .../src/edit_prediction_tests.rs | 12 +- crates/edit_prediction/src/prediction.rs | 94 ++++------ crates/edit_prediction/src/zeta.rs | 118 ++++++------- crates/edit_prediction_cli/src/predict.rs | 2 +- .../src/rate_prediction_modal.rs | 18 +- 6 files changed, 176 insertions(+), 231 deletions(-) diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index ac95aac0df..1136a482a8 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -180,7 +180,7 @@ pub struct EditPredictionStore { legacy_data_collection_enabled: bool, reject_predictions_tx: mpsc::UnboundedSender, settled_predictions_tx: mpsc::UnboundedSender, - shown_predictions: VecDeque, + rateable_predictions: VecDeque, rated_predictions: HashSet, #[cfg(test)] settled_event_callback: Option>, @@ -965,7 +965,7 @@ impl EditPredictionStore { reject_predictions_tx: reject_tx, settled_predictions_tx, rated_predictions: Default::default(), - shown_predictions: Default::default(), + rateable_predictions: Default::default(), #[cfg(test)] settled_event_callback: None, @@ -1034,7 +1034,7 @@ impl EditPredictionStore { pub fn active_experiment(&self) -> Option<&str> { self.preferred_experiment.as_deref().or_else(|| { - self.shown_predictions + self.rateable_predictions .iter() .find_map(|p| p.model_version.as_ref()) .and_then(|model_version| model_version.strip_prefix("zeta2:")) @@ -2174,10 +2174,10 @@ impl EditPredictionStore { } if is_first_non_jump_show { - self.shown_predictions + self.rateable_predictions .push_front(current_prediction.prediction.clone()); - if self.shown_predictions.len() > 50 { - let completion = self.shown_predictions.pop_back().unwrap(); + if self.rateable_predictions.len() > 50 { + let completion = self.rateable_predictions.pop_back().unwrap(); self.rated_predictions.remove(&completion.id); } } @@ -2390,11 +2390,10 @@ impl EditPredictionStore { prediction_result } else { EditPredictionResult { - id: prediction_result.id, - prediction: Err(EditPredictionRejectReason::CurrentPreferred), - display_prediction: None, - model_version: prediction_result.model_version, - e2e_latency: prediction_result.e2e_latency, + reject_reason: Some( + EditPredictionRejectReason::CurrentPreferred, + ), + ..prediction_result } }, PredictionRequestedBy::DiagnosticsUpdate, @@ -2566,9 +2565,12 @@ impl EditPredictionStore { } let new_prediction_result = do_refresh(this.clone(), cx).await.log_err().flatten(); - let new_prediction_metadata = new_prediction_result - .as_ref() - .map(|(prediction, _)| (prediction.id.clone(), prediction.model_version.clone())); + let new_prediction_metadata = new_prediction_result.as_ref().map(|(result, _)| { + ( + result.prediction.id.clone(), + result.prediction.model_version.clone(), + ) + }); // When a prediction completes, remove it from the pending list, and cancel // any pending predictions that were enqueued before it. @@ -2582,81 +2584,72 @@ impl EditPredictionStore { let new_current_prediction = if !is_cancelled && let Some((prediction_result, requested_by)) = new_prediction_result { - match prediction_result { - EditPredictionResult { - prediction: Ok(prediction), - e2e_latency, - .. - } => { - let new_prediction = CurrentEditPrediction { - requested_by, - prediction, - was_shown: false, - shown_with: None, - e2e_latency, - }; + let EditPredictionResult { + prediction, + reject_reason, + e2e_latency, + } = prediction_result; - if let Some(current_prediction) = - project_state.current_prediction.as_ref() + if let Some(reject_reason) = reject_reason { + let should_allow_rating_prediction = matches!( + reject_reason, + EditPredictionRejectReason::Empty + | EditPredictionRejectReason::InterpolatedEmpty + ); + let prediction_id = prediction.id.clone(); + let model_version = prediction.model_version.clone(); + + this.reject_prediction( + prediction_id, + reject_reason, + false, + model_version, + Some(e2e_latency), + cx, + ); + + if should_allow_rating_prediction { + this.rateable_predictions.push_front(prediction); + if this.rateable_predictions.len() > 50 + && let Some(completion) = this.rateable_predictions.pop_back() { - if new_prediction.should_replace_prediction(¤t_prediction, cx) - { - this.reject_current_prediction( - EditPredictionRejectReason::Replaced, - &project, - cx, - ); - - Some(new_prediction) - } else { - this.reject_prediction( - new_prediction.prediction.id, - EditPredictionRejectReason::CurrentPreferred, - false, - new_prediction.prediction.model_version, - Some(new_prediction.e2e_latency), - cx, - ); - None - } - } else { - Some(new_prediction) + this.rated_predictions.remove(&completion.id); } } - EditPredictionResult { - id, - prediction: Err(reject_reason), - display_prediction, - model_version, + + None + } else { + let new_prediction = CurrentEditPrediction { + requested_by, + prediction, + was_shown: false, + shown_with: None, e2e_latency, - } => { - let should_show_rejected_prediction = matches!( - reject_reason, - EditPredictionRejectReason::Empty - | EditPredictionRejectReason::InterpolatedEmpty - ); + }; - this.reject_prediction( - id, - reject_reason, - false, - model_version, - Some(e2e_latency), - cx, - ); + if let Some(current_prediction) = project_state.current_prediction.as_ref() + { + if new_prediction.should_replace_prediction(¤t_prediction, cx) { + this.reject_current_prediction( + EditPredictionRejectReason::Replaced, + &project, + cx, + ); - if should_show_rejected_prediction - && let Some(display_prediction) = display_prediction - { - this.shown_predictions.push_front(display_prediction); - if this.shown_predictions.len() > 50 - && let Some(completion) = this.shown_predictions.pop_back() - { - this.rated_predictions.remove(&completion.id); - } + Some(new_prediction) + } else { + this.reject_prediction( + new_prediction.prediction.id, + EditPredictionRejectReason::CurrentPreferred, + false, + new_prediction.prediction.model_version, + Some(new_prediction.e2e_latency), + cx, + ); + None } - - None + } else { + Some(new_prediction) } } } else { @@ -3319,12 +3312,12 @@ impl EditPredictionStore { }) } - pub fn shown_predictions(&self) -> impl DoubleEndedIterator { - self.shown_predictions.iter() + pub fn rateable_predictions(&self) -> impl DoubleEndedIterator { + self.rateable_predictions.iter() } - pub fn shown_completions_len(&self) -> usize { - self.shown_predictions.len() + pub fn rateable_predictions_count(&self) -> usize { + self.rateable_predictions.len() } pub fn is_prediction_rated(&self, id: &EditPredictionId) -> bool { diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 5195393f2f..3dc02f8eb5 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -482,7 +482,7 @@ async fn test_simple_request(cx: &mut TestAppContext) { )) .unwrap(); - let prediction = prediction_task.await.unwrap().unwrap().prediction.unwrap(); + let prediction = prediction_task.await.unwrap().unwrap().prediction; assert_eq!(prediction.edits.len(), 1); assert_eq!( @@ -565,7 +565,7 @@ async fn test_request_events(cx: &mut TestAppContext) { )) .unwrap(); - let prediction = prediction_task.await.unwrap().unwrap().prediction.unwrap(); + let prediction = prediction_task.await.unwrap().unwrap().prediction; assert_eq!(prediction.edits.len(), 1); assert_eq!(prediction.edits[0].1.as_ref(), " are you?"); @@ -1478,7 +1478,7 @@ async fn test_empty_prediction(cx: &mut TestAppContext) { .prediction_at(&buffer, None, &project, cx) .is_none() ); - let shown_predictions = ep_store.shown_predictions().collect::>(); + let shown_predictions = ep_store.rateable_predictions().collect::>(); assert_eq!(shown_predictions.len(), 1); assert_eq!(shown_predictions[0].id.to_string(), id); assert!(shown_predictions[0].edits.is_empty()); @@ -1556,7 +1556,7 @@ async fn test_interpolated_empty(cx: &mut TestAppContext) { .prediction_at(&buffer, None, &project, cx) .is_none() ); - let shown_predictions = ep_store.shown_predictions().collect::>(); + let shown_predictions = ep_store.rateable_predictions().collect::>(); assert_eq!(shown_predictions.len(), 1); assert_eq!(shown_predictions[0].id.to_string(), id); assert!(shown_predictions[0].edits.is_empty()); @@ -1777,7 +1777,7 @@ async fn test_current_preferred(cx: &mut TestAppContext) { first_id ); let shown_prediction_ids = ep_store - .shown_predictions() + .rateable_predictions() .map(|prediction| prediction.id.to_string()) .collect::>(); assert!(shown_prediction_ids.is_empty()); @@ -3351,7 +3351,7 @@ async fn run_edit_prediction( cx, ) }); - prediction_task.await.unwrap().unwrap().prediction.unwrap() + prediction_task.await.unwrap().unwrap().prediction } async fn make_test_ep_store( diff --git a/crates/edit_prediction/src/prediction.rs b/crates/edit_prediction/src/prediction.rs index c01c52e2bd..b9d41c44ea 100644 --- a/crates/edit_prediction/src/prediction.rs +++ b/crates/edit_prediction/src/prediction.rs @@ -23,10 +23,8 @@ 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, - pub display_prediction: Option, - pub model_version: Option, + pub prediction: EditPrediction, + pub reject_reason: Option, pub e2e_latency: std::time::Duration, } @@ -44,65 +42,38 @@ impl EditPredictionResult { e2e_latency: std::time::Duration, cx: &mut AsyncApp, ) -> Self { - if edits.is_empty() { - let empty_edits = Arc::new([]); - return Self { - id: id.clone(), - prediction: Err(EditPredictionRejectReason::Empty), - display_prediction: Some(EditPrediction { - id, - edits: empty_edits, - cursor_position: None, - editable_range, - snapshot: edited_buffer_snapshot.clone(), - edit_preview: EditPreview::unchanged(edited_buffer_snapshot), - buffer: edited_buffer.clone(), - inputs, - model_version: model_version.clone(), - trigger, - }), - model_version, - e2e_latency, - }; - } + 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) + }) + }) + .unwrap_or_default(); + let snapshot = new_snapshot.unwrap_or_else(|| edited_buffer_snapshot.clone()); - let (edits, snapshot) = edited_buffer.read_with(cx, |buffer, _cx| { - let new_snapshot = buffer.snapshot(); - let edits: Option, Arc)]>> = - interpolate_edits(&edited_buffer_snapshot, &new_snapshot, &edits).map(Arc::from); + 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(); - (edits, new_snapshot) - }); - - let Some(edits) = edits else { - let empty_edits: Arc<[(Range, Arc)]> = Vec::new().into(); - return Self { - id: id.clone(), - prediction: Err(EditPredictionRejectReason::InterpolatedEmpty), - display_prediction: Some(EditPrediction { - id, - edits: empty_edits, - cursor_position: None, - editable_range, - snapshot: edited_buffer_snapshot.clone(), - edit_preview: EditPreview::unchanged(edited_buffer_snapshot), - inputs, - buffer: edited_buffer.clone(), - model_version: model_version.clone(), - trigger, - }), - model_version, - e2e_latency, - }; + let edit_preview = if !edits.is_empty() { + edited_buffer + .read_with(cx, |buffer, cx| buffer.preview_edits(edits.clone(), cx)) + .await + } else { + EditPreview::unchanged(edited_buffer_snapshot) }; - let edit_preview = edited_buffer - .read_with(cx, |buffer, cx| buffer.preview_edits(edits.clone(), cx)) - .await; - Self { - id: id.clone(), - prediction: Ok(EditPrediction { + prediction: EditPrediction { id, edits, cursor_position, @@ -111,11 +82,10 @@ impl EditPredictionResult { edit_preview, inputs, buffer: edited_buffer.clone(), - model_version: model_version.clone(), + model_version, trigger, - }), - display_prediction: None, - model_version, + }, + reject_reason, e2e_latency, } } diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 013929b391..2631d0fd22 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -7,9 +7,7 @@ use crate::{ prediction::EditPredictionResult, }; use anyhow::{Context as _, Result}; -use cloud_llm_client::{ - AcceptEditPredictionBody, EditPredictionRejectReason, predict_edits_v3::RawCompletionRequest, -}; +use cloud_llm_client::{AcceptEditPredictionBody, predict_edits_v3::RawCompletionRequest}; use edit_prediction_types::PredictedCursorPosition; use futures::future::Shared; use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*}; @@ -23,12 +21,11 @@ use ui::SharedString; use workspace::notifications::simple_message_notification::MessageNotification; use workspace::notifications::{NotificationId, show_app_notification}; use workspace::workspace_error::{ErrorAction, ErrorSeverity, WorkspaceError}; -use zeta_prompt::{ParsedOutput, ZetaPromptInput}; use std::{ops::Range, path::Path, sync::Arc}; use zeta_prompt::{ - ZetaFormat, excerpt_range_for_format, format_zeta_prompt, get_prefill, - parse_zeta2_model_output, stop_tokens_for_format, + ParsedOutput, ZetaFormat, ZetaPromptInput, excerpt_ranges_for_format, format_zeta_prompt, + get_prefill, parse_zeta2_model_output, stop_tokens_for_format, zeta1::{self, EDITABLE_REGION_END_MARKER}, }; @@ -304,78 +301,70 @@ pub fn request_prediction_with_zeta( log::trace!("Got edit prediction response"); - let Some(ParsedOutput { - new_editable_region: mut output_text, - range_in_excerpt: editable_range_in_excerpt, - cursor_offset_in_new_editable_region: cursor_offset_in_output, - }) = output - else { - let editable_range_in_excerpt = - excerpt_range_for_format(zeta_format, &prompt_input.excerpt_ranges).0; - let editable_range_in_buffer = editable_range_in_excerpt.start - + full_context_offset_range.start - ..editable_range_in_excerpt.end + full_context_offset_range.start; - - return Ok(( - Some(( - request_id, - Some(Prediction { - prompt_input, - buffer, - snapshot: snapshot.clone(), - edits: Vec::new(), - cursor_position: None, - editable_range_in_buffer, - }), - model_version, - )), - usage, - )); - }; + let (new_editable_region, cursor_offset_in_output, editable_range_in_excerpt) = + if let Some(output) = output { + let ParsedOutput { + new_editable_region: output_text, + range_in_excerpt: editable_range_in_excerpt, + cursor_offset_in_new_editable_region: cursor_offset_in_output, + } = output; + ( + Some(output_text), + cursor_offset_in_output, + editable_range_in_excerpt, + ) + } else { + let (editable_range, _) = + excerpt_ranges_for_format(zeta_format, &prompt_input.excerpt_ranges); + (None, None, editable_range) + }; let editable_range_in_buffer = editable_range_in_excerpt.start + full_context_offset_range.start ..editable_range_in_excerpt.end + full_context_offset_range.start; - let mut old_text = snapshot - .text_for_range(editable_range_in_buffer.clone()) - .collect::(); - if let Some(debug_tx) = &debug_tx { debug_tx .unbounded_send(DebugEvent::EditPredictionFinished( EditPredictionFinishedDebugEvent { buffer: buffer.downgrade(), position, - model_output: Some(output_text.clone()), + model_output: new_editable_region.clone(), }, )) .ok(); } + let (edits, cursor_position) = new_editable_region + .map(|mut output_text| { + let mut old_text = snapshot + .text_for_range(editable_range_in_buffer.clone()) + .collect::(); - if !output_text.is_empty() && !output_text.ends_with('\n') { - output_text.push('\n'); - } - if !old_text.is_empty() && !old_text.ends_with('\n') { - old_text.push('\n'); - } + if !output_text.is_empty() && !output_text.ends_with('\n') { + output_text.push('\n'); + } + if !old_text.is_empty() && !old_text.ends_with('\n') { + old_text.push('\n'); + } - let (edits, cursor_position) = compute_edits_and_cursor_position( - old_text, - &output_text, - editable_range_in_buffer.start, - cursor_offset_in_output, - &snapshot, - ); + compute_edits_and_cursor_position( + old_text, + &output_text, + editable_range_in_buffer.start, + cursor_offset_in_output, + &snapshot, + ) + }) + .unwrap_or_default(); - let prediction = Some(Prediction { + let prediction = Prediction { prompt_input, buffer, snapshot: snapshot.clone(), edits, cursor_position, editable_range_in_buffer, - }); + }; anyhow::Ok((Some((request_id, prediction, model_version)), usage)) } @@ -389,32 +378,24 @@ pub fn request_prediction_with_zeta( }; let request_duration = cx.background_executor().now() - request_start; - let Some(Prediction { + let Prediction { prompt_input: inputs, buffer: edited_buffer, snapshot: edited_buffer_snapshot, edits, cursor_position, editable_range_in_buffer, - .. - }) = prediction - else { - return Ok(Some(EditPredictionResult { - id, - prediction: Err(EditPredictionRejectReason::Empty), - display_prediction: None, - model_version, - e2e_latency: request_duration, - })); - }; + } = prediction; + let editable_anchor_range = + edited_buffer_snapshot.anchor_range_inside(editable_range_in_buffer.clone()); let result = EditPredictionResult::new( id, &edited_buffer, &edited_buffer_snapshot, edits.into(), cursor_position, - Some(edited_buffer_snapshot.anchor_range_inside(editable_range_in_buffer.clone())), + Some(editable_anchor_range), inputs, model_version, trigger, @@ -423,7 +404,8 @@ pub fn request_prediction_with_zeta( ) .await; - if can_collect_data && let Ok(prediction) = &result.prediction { + if can_collect_data { + let prediction = &result.prediction; let weak_this = this.clone(); let request_id = prediction.id.clone(); let edited_buffer = edited_buffer.clone(); diff --git a/crates/edit_prediction_cli/src/predict.rs b/crates/edit_prediction_cli/src/predict.rs index 5ba7efb5bf..a0e1ffdcaf 100644 --- a/crates/edit_prediction_cli/src/predict.rs +++ b/crates/edit_prediction_cli/src/predict.rs @@ -302,7 +302,7 @@ pub async fn run_prediction( .await?; let actual_patch = prediction.and_then(|prediction| { - let prediction = prediction.prediction.ok()?; + let prediction = prediction.prediction; prediction .edit_preview .as_unified_diff(prediction.snapshot.file(), &prediction.edits) diff --git a/crates/edit_prediction_ui/src/rate_prediction_modal.rs b/crates/edit_prediction_ui/src/rate_prediction_modal.rs index deae41c21e..be19ac5a74 100644 --- a/crates/edit_prediction_ui/src/rate_prediction_modal.rs +++ b/crates/edit_prediction_ui/src/rate_prediction_modal.rs @@ -138,7 +138,7 @@ impl RatePredictionsModal { self.selected_index += 1; self.selected_index = usize::min( self.selected_index, - self.ep_store.read(cx).shown_predictions().count(), + self.ep_store.read(cx).rateable_predictions().count(), ); cx.notify(); } @@ -157,7 +157,7 @@ impl RatePredictionsModal { let next_index = self .ep_store .read(cx) - .shown_predictions() + .rateable_predictions() .skip(self.selected_index) .enumerate() .skip(1) // Skip straight to the next item @@ -172,12 +172,12 @@ impl RatePredictionsModal { fn select_prev_edit(&mut self, _: &PreviousEdit, _: &mut Window, cx: &mut Context) { let ep_store = self.ep_store.read(cx); - let completions_len = ep_store.shown_completions_len(); + let completions_len = ep_store.rateable_predictions_count(); let prev_index = self .ep_store .read(cx) - .shown_predictions() + .rateable_predictions() .rev() .skip((completions_len - 1) - self.selected_index) .enumerate() @@ -198,7 +198,7 @@ impl RatePredictionsModal { } fn select_last(&mut self, _: &menu::SelectLast, _window: &mut Window, cx: &mut Context) { - self.selected_index = self.ep_store.read(cx).shown_completions_len() - 1; + self.selected_index = self.ep_store.read(cx).rateable_predictions_count() - 1; cx.notify(); } @@ -283,7 +283,7 @@ impl RatePredictionsModal { let completion = self .ep_store .read(cx) - .shown_predictions() + .rateable_predictions() .skip(self.selected_index) .take(1) .next() @@ -296,7 +296,7 @@ impl RatePredictionsModal { let completion = self .ep_store .read(cx) - .shown_predictions() + .rateable_predictions() .skip(self.selected_index) .take(1) .next() @@ -421,7 +421,7 @@ impl RatePredictionsModal { self.selected_index = self .ep_store .read(cx) - .shown_predictions() + .rateable_predictions() .enumerate() .find(|(_, completion_b)| prediction.id == completion_b.id) .map(|(ix, _)| ix) @@ -1127,7 +1127,7 @@ impl RatePredictionsModal { fn render_shown_completions(&self, cx: &Context) -> impl Iterator { self.ep_store .read(cx) - .shown_predictions() + .rateable_predictions() .cloned() .enumerate() .map(|(index, completion)| {