yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ep: Track settled telemetry for empty predictions (#59645)

Browse Source 
# 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:

<details>
  <summary>Click to view showcase</summary>

My super cool demos here

</details>

---

Release Notes:

- N/A or Added/Fixed/Improved ...
This commit is contained in:
Ben Kunkle authored and GitHub committed 2026-06-20 22:40:15 +00:00
1 parent d8e04a57e0
commit bb807e1e3a
6 files changed
+176 -231
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/edit_prediction/src/edit_prediction.rs
+78 -85
View File
@@ -180,7 +180,7 @@ pub struct EditPredictionStore {
legacy_data_collection_enabled: bool,
reject_predictions_tx: mpsc::UnboundedSender<EditPredictionRejectionPayload>,
settled_predictions_tx: mpsc::UnboundedSender<Instant>,
shown_predictions: VecDeque<EditPrediction>,
rateable_predictions: VecDeque<EditPrediction>,
rated_predictions: HashSet<EditPredictionId>,
#[cfg(test)]
settled_event_callback: Option<Box<dyn Fn(EditPredictionId, String)>>,
@@ -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(&current_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(&current_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<Item = &EditPrediction> {
self.shown_predictions.iter()
pub fn rateable_predictions(&self) -> impl DoubleEndedIterator<Item = &EditPrediction> {
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 {
crates/edit_prediction/src/edit_prediction_tests.rs
+6 -6
View File
@@ -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::<Vec<_>>();
let shown_predictions = ep_store.rateable_predictions().collect::<Vec<_>>();
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::<Vec<_>>();
let shown_predictions = ep_store.rateable_predictions().collect::<Vec<_>>();
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::<Vec<_>>();
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(
crates/edit_prediction/src/prediction.rs
+32 -62
View File
@@ -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<EditPrediction, EditPredictionRejectReason>,
pub display_prediction: Option<EditPrediction>,
pub model_version: Option<String>,
pub prediction: EditPrediction,
pub reject_reason: Option<EditPredictionRejectReason>,
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<Anchor>, Arc<str>)]> =
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<[(Range<Anchor>, Arc<str>)]>> =
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<Anchor>, Arc<str>)]> = 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,
}
}
crates/edit_prediction/src/zeta.rs
+50 -68
View File
@@ -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::<String>();
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::<String>();
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();
crates/edit_prediction_cli/src/predict.rs
+1 -1
View File
@@ -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)
crates/edit_prediction_ui/src/rate_prediction_modal.rs
+9 -9
View File
@@ -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<Self>) {
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>) {
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<Self>) -> impl Iterator<Item = ListItem> {
self.ep_store
.read(cx)
.shown_predictions()
.rateable_predictions()
.cloned()
.enumerate()
.map(|(index, completion)| {