From 560f784776b4662577ac6b393ccf47ffbba0cd08 Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Thu, 21 May 2026 11:23:20 -0400 Subject: [PATCH] ep: Increase example capture rate (#57208) 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 is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes #ISSUE Release Notes: - N/A or Added/Fixed/Improved ... --- crates/buffer_diff/src/buffer_diff.rs | 4 + crates/edit_prediction/src/capture_example.rs | 288 ++++++++++++++---- crates/edit_prediction/src/edit_prediction.rs | 145 ++++++++- .../src/edit_prediction_tests.rs | 3 + crates/edit_prediction/src/example_spec.rs | 98 ++++++ crates/edit_prediction/src/zeta.rs | 46 ++- .../edit_prediction_cli/src/format_prompt.rs | 6 + .../edit_prediction_cli/src/pull_examples.rs | 3 + crates/edit_prediction_cli/src/repair.rs | 3 + .../edit_prediction_cli/src/split_commit.rs | 6 + crates/edit_prediction_cli/src/synthesize.rs | 3 + .../src/edit_prediction_button.rs | 9 +- .../src/edit_prediction_ui.rs | 76 +---- crates/language/src/language_settings.rs | 2 - crates/settings_content/src/language.rs | 4 +- 15 files changed, 530 insertions(+), 166 deletions(-) diff --git a/crates/buffer_diff/src/buffer_diff.rs b/crates/buffer_diff/src/buffer_diff.rs index 56c3fe9f51..ccf8dfd341 100644 --- a/crates/buffer_diff/src/buffer_diff.rs +++ b/crates/buffer_diff/src/buffer_diff.rs @@ -267,6 +267,10 @@ impl BufferDiffSnapshot { .then(|| self.inner.base_text.text()) } + pub fn base_text_exists(&self) -> bool { + self.inner.base_text_exists + } + pub fn secondary_diff(&self) -> Option<&BufferDiffSnapshot> { self.secondary_diff.as_deref() } diff --git a/crates/edit_prediction/src/capture_example.rs b/crates/edit_prediction/src/capture_example.rs index 9463456132..c01cd93578 100644 --- a/crates/edit_prediction/src/capture_example.rs +++ b/crates/edit_prediction/src/capture_example.rs @@ -1,18 +1,24 @@ -use crate::{StoredEvent, example_spec::ExampleSpec}; +use crate::{ + StoredEvent, + example_spec::{ExampleSpec, RecentFile}, +}; use anyhow::Result; -use buffer_diff::BufferDiffSnapshot; +use buffer_diff::{BufferDiff, BufferDiffSnapshot}; use collections::HashMap; use gpui::{App, Entity, Task}; use language::Buffer; -use project::{Project, WorktreeId}; +use project::Project; use std::{collections::hash_map, fmt::Write as _, ops::Range, path::Path, sync::Arc}; -use text::{BufferSnapshot as TextBufferSnapshot, Point}; +use text::Point; pub fn capture_example( project: Entity, buffer: Entity, cursor_anchor: language::Anchor, - mut events: Vec, + events: Vec, + recently_opened_files: Vec, + recently_viewed_files: Vec, + uncommitted_diffs_by_path: HashMap, Entity>, populate_expected_patch: bool, cx: &mut App, ) -> Option>> { @@ -34,16 +40,39 @@ pub fn capture_example( .or_else(|| repository_snapshot.remote_upstream_url.clone())?; let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string(); - let git_store = project.read(cx).git_store().clone(); + Some(cx.spawn(async move |cx| { + let mut events = events; + let mut diff_buffers_by_path: HashMap, (Entity, Entity)> = + HashMap::default(); + for stored_event in &events { + let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref(); + let Some((project_path, relative_path)) = project.read_with(cx, |project, cx| { + let project_path = project + .find_project_path(path, cx) + .filter(|path| path.worktree_id == worktree_id)?; + let relative_path: Arc = project_path.path.as_std_path().into(); + Some((project_path, relative_path)) + }) else { + continue; + }; - Some(cx.spawn(async move |mut cx| { - let snapshots_by_path = - collect_snapshots(&project, &git_store, worktree_id, &events, &mut cx).await?; + if let hash_map::Entry::Vacant(entry) = diff_buffers_by_path.entry(relative_path) { + let Some(diff) = uncommitted_diffs_by_path.get(entry.key()).cloned() else { + continue; + }; + let buffer = project + .update(cx, |project, cx| { + project.open_buffer(project_path.clone(), cx) + }) + .await?; + entry.insert((buffer, diff)); + } + } events.retain(|stored_event| { let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref(); let relative_path = strip_root_name(path, &root_name); - snapshots_by_path.contains_key(relative_path) + diff_buffers_by_path.contains_key(relative_path) }); let line_comment_prefix = snapshot @@ -56,9 +85,17 @@ pub fn capture_example( .background_executor() .spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) }) .await; + let uncommitted_diff_snapshots = diff_buffers_by_path + .into_iter() + .map(|(relative_path, (buffer, diff))| { + let snapshot = buffer.update(cx, |buffer, _| buffer.snapshot()); + let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx)); + (relative_path, (snapshot, diff_snapshot)) + }) + .collect(); let uncommitted_diff = cx .background_executor() - .spawn(async move { compute_uncommitted_diff(snapshots_by_path) }) + .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshots) }) .await; let mut edit_history = String::new(); @@ -68,6 +105,7 @@ pub fn capture_example( edit_history.push('\n'); } } + let uncommitted_diff_contains_edit_history = !edit_history.is_empty(); // Initialize an empty patch with context lines, to make it easy // to write the expected patch by hand. @@ -100,6 +138,9 @@ pub fn capture_example( tags: Vec::new(), reasoning: None, uncommitted_diff, + recently_opened_files, + recently_viewed_files, + uncommitted_diff_contains_edit_history, cursor_path, cursor_position: String::new(), edit_history, @@ -184,63 +225,107 @@ fn compute_cursor_excerpt( ) } -async fn collect_snapshots( - project: &Entity, - git_store: &Entity, - worktree_id: WorktreeId, - events: &[StoredEvent], - cx: &mut gpui::AsyncApp, -) -> Result, (TextBufferSnapshot, BufferDiffSnapshot)>> { - let mut snapshots_by_path = HashMap::default(); - for stored_event in events { - let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref(); - if let Some((project_path, relative_path)) = project.read_with(cx, |project, cx| { - let project_path = project - .find_project_path(path, cx) - .filter(|path| path.worktree_id == worktree_id)?; - let relative_path: Arc = project_path.path.as_std_path().into(); - Some((project_path, relative_path)) - }) { - if let hash_map::Entry::Vacant(entry) = snapshots_by_path.entry(relative_path) { - let buffer = project - .update(cx, |project, cx| { - project.open_buffer(project_path.clone(), cx) - }) - .await?; - let diff = git_store - .update(cx, |git_store, cx| { - git_store.open_uncommitted_diff(buffer.clone(), cx) - }) - .await?; - let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx)); - entry.insert((stored_event.old_snapshot.clone(), diff_snapshot)); - } - } - } - Ok(snapshots_by_path) -} - fn compute_uncommitted_diff( - snapshots_by_path: HashMap, (TextBufferSnapshot, BufferDiffSnapshot)>, + snapshots_by_path: HashMap, (language::BufferSnapshot, BufferDiffSnapshot)>, ) -> String { let mut uncommitted_diff = String::new(); - for (relative_path, (before_text, diff_snapshot)) in snapshots_by_path { - if let Some(head_text) = &diff_snapshot.base_text_string() { - let file_diff = language::unified_diff(head_text, &before_text.text()); - if !file_diff.is_empty() { - let path_str = relative_path.to_string_lossy(); - writeln!(uncommitted_diff, "--- a/{path_str}").ok(); - writeln!(uncommitted_diff, "+++ b/{path_str}").ok(); - uncommitted_diff.push_str(&file_diff); - if !uncommitted_diff.ends_with('\n') { - uncommitted_diff.push('\n'); - } + let mut snapshots_by_path = snapshots_by_path.into_iter().collect::>(); + snapshots_by_path.sort_by(|(left_path, _), (right_path, _)| left_path.cmp(right_path)); + for (relative_path, (buffer_snapshot, diff_snapshot)) in snapshots_by_path { + let base_snapshot = diff_snapshot.base_text(); + let is_existing_file = diff_snapshot.base_text_exists(); + + let new_path_str = relative_path.to_string_lossy(); + let old_path_str = if is_existing_file { + new_path_str.as_ref() + } else { + "/dev/null" + }; + writeln!( + uncommitted_diff, + "--- {}{old_path_str}", + if is_existing_file { "a/" } else { "" } + ) + .ok(); + writeln!(uncommitted_diff, "+++ b/{new_path_str}").ok(); + + if !is_existing_file { + let new_text = buffer_snapshot.text(); + writeln!( + uncommitted_diff, + "@@ -0,0 +1,{} @@", + new_text.lines().count() + ) + .ok(); + for line in new_text.lines() { + writeln!(uncommitted_diff, "+{line}").ok(); } + continue; + } + + let mut ranges: Vec<(Range, Range)> = Vec::new(); + for hunk in (&diff_snapshot).hunks(&buffer_snapshot) { + let old_start = base_snapshot + .offset_to_point(hunk.diff_base_byte_range.start) + .row; + let old_end = + exclusive_end_row(base_snapshot.offset_to_point(hunk.diff_base_byte_range.end)); + let new_start = hunk.range.start.row; + let new_end = exclusive_end_row(hunk.range.end); + let old_range = old_start.saturating_sub(3)..old_end + 3; + let new_range = new_start.saturating_sub(3)..new_end + 3; + + if let Some((last_old_range, last_new_range)) = ranges.last_mut() + && (old_range.start <= last_old_range.end || new_range.start <= last_new_range.end) + { + last_old_range.end = last_old_range.end.max(old_range.end); + last_new_range.end = last_new_range.end.max(new_range.end); + continue; + } + ranges.push((old_range, new_range)); + } + + for (old_range, new_range) in ranges { + uncommitted_diff.push_str(&language::unified_diff_with_offsets( + &base_snapshot + .text_for_range( + Point::new(old_range.start, 0) + ..row_start_or_max(base_snapshot, old_range.end), + ) + .collect::(), + &buffer_snapshot + .text_for_range( + Point::new(new_range.start, 0) + ..row_start_or_max(&buffer_snapshot, new_range.end), + ) + .collect::(), + old_range.start, + new_range.start, + )); + } + if !uncommitted_diff.ends_with('\n') { + uncommitted_diff.push('\n'); } } uncommitted_diff } +fn row_start_or_max(snapshot: &language::BufferSnapshot, row: u32) -> Point { + if row >= snapshot.max_point().row { + snapshot.max_point() + } else { + Point::new(row, 0) + } +} + +fn exclusive_end_row(point: Point) -> u32 { + if point.column == 0 { + point.row + } else { + point.row + 1 + } +} + fn generate_timestamp_name() -> String { let format = time::format_description::parse("[year]-[month]-[day] [hour]:[minute]:[second]"); match format { @@ -309,7 +394,9 @@ mod tests { json!({ ".git": {}, "src": { + "deleted.rs": "pub fn deleted_file() {\n deleted();\n}\n", "main.rs": disk_contents, + "new.rs": "pub fn new_file() {\n}\n", } }), ) @@ -326,7 +413,13 @@ mod tests { fs.set_head_for_repo( Path::new("/project/.git"), - &[("src/main.rs", committed_contents.to_string())], + &[ + ( + "src/deleted.rs", + "pub fn deleted_file() {\n deleted();\n}\n".to_string(), + ), + ("src/main.rs", committed_contents.to_string()), + ], "abc123def456", ); fs.set_remote_for_repo( @@ -350,6 +443,21 @@ mod tests { }); cx.run_until_parked(); + let deleted_file_buffer = project + .update(cx, |project, cx| { + project.open_local_buffer("/project/src/deleted.rs", cx) + }) + .await + .unwrap(); + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&deleted_file_buffer, &project, cx) + }); + cx.run_until_parked(); + deleted_file_buffer.update(cx, |buffer, cx| { + buffer.edit([(0..buffer.len(), "")], None, cx); + }); + cx.run_until_parked(); + buffer.update(cx, |buffer, cx| { let point = Point::new(6, 0); buffer.edit([(point..point, " // comment 3\n")], None, cx); @@ -379,6 +487,22 @@ mod tests { }); cx.run_until_parked(); + let new_file_buffer = project + .update(cx, |project, cx| { + project.open_local_buffer("/project/src/new.rs", cx) + }) + .await + .unwrap(); + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&new_file_buffer, &project, cx) + }); + cx.run_until_parked(); + new_file_buffer.update(cx, |buffer, cx| { + let point = Point::new(1, 0); + buffer.edit([(point..point, " created();\n")], None, cx); + }); + cx.run_until_parked(); + // Open and edit an external file (outside the main project's worktree) let external_buffer = project .update(cx, |project, cx| { @@ -410,6 +534,13 @@ mod tests { "external file edit should be in events" ); + let worktree_id = buffer.read_with(cx, |buffer, cx| buffer.file().unwrap().worktree_id(cx)); + let uncommitted_diffs_by_path = ep_store + .update(cx, |store, cx| { + store.uncommitted_diffs_for_events(project.clone(), worktree_id, events.clone(), cx) + }) + .await + .unwrap(); let mut example = cx .update(|cx| { capture_example( @@ -417,6 +548,9 @@ mod tests { buffer.clone(), Anchor::min_for_buffer(buffer.read(cx).remote_id()), events, + Vec::new(), + Vec::new(), + uncommitted_diffs_by_path, true, cx, ) @@ -435,23 +569,41 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: indoc! {" + --- a/src/deleted.rs + +++ b/src/deleted.rs + @@ -1,3 +1,0 @@ + -pub fn deleted_file() { + - deleted(); + -} --- a/src/main.rs +++ b/src/main.rs - @@ -1,4 +1,5 @@ + @@ -1,11 +1,15 @@ fn main() { + // comment 1 one(); two(); + + // comment 4 three(); - @@ -7,5 +8,6 @@ + four(); + + // comment 3 + five(); six(); seven(); eight(); + // comment 2 nine(); } + --- /dev/null + +++ b/src/new.rs + @@ -0,0 +1,3 @@ + +pub fn new_file() { + + created(); + +} "} .to_string(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: true, cursor_path: Path::new("src/main.rs").into(), cursor_position: indoc! {" fn main() { @@ -473,6 +625,12 @@ mod tests { "} .to_string(), edit_history: indoc! {" + --- a/src/deleted.rs + +++ b/src/deleted.rs + @@ -1,3 +1,0 @@ + -pub fn deleted_file() { + - deleted(); + -} --- a/src/main.rs +++ b/src/main.rs @@ -2,8 +2,10 @@ @@ -486,6 +644,12 @@ mod tests { five(); six(); seven(); + --- a/src/new.rs + +++ b/src/new.rs + @@ -1,2 +1,3 @@ + pub fn new_file() { + + created(); + } "} .to_string(), expected_patches: vec![ diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 8b57c442c2..f6faa7d712 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use buffer_diff::BufferDiff; use client::{Client, EditPredictionUsage, UserStore, global_llm_token}; use cloud_api_client::LlmApiToken; use cloud_api_types::{ @@ -88,6 +89,7 @@ mod edit_prediction_tests; use crate::cursor_excerpt::expand_context_syntactically_then_linewise; use crate::example_spec::ExampleSpec; +use crate::example_spec::RecentFile; use crate::license_detection::LicenseDetectionWatcher; use crate::mercury::Mercury; pub use crate::metrics::{KeptRateResult, compute_kept_rate}; @@ -95,7 +97,6 @@ use crate::onboarding_modal::ZedPredictModal; pub use crate::prediction::EditPrediction; pub use crate::prediction::EditPredictionId; use crate::prediction::EditPredictionResult; -pub use capture_example::capture_example; pub use language_model::ApiKeyState; pub use telemetry_events::EditPredictionRating; pub use zed_edit_prediction_delegate::ZedEditPredictionDelegate; @@ -238,6 +239,7 @@ pub struct StoredEvent { pub old_snapshot: TextBufferSnapshot, pub new_snapshot_version: clock::Global, pub total_edit_range: Range, + pub uncommitted_diff: Option>, } impl StoredEvent { @@ -572,6 +574,7 @@ impl LastEvent { new_snapshot_version: self.new_snapshot.version.clone(), total_edit_range: self.new_snapshot.anchor_before(edit_range.start) ..self.new_snapshot.anchor_before(edit_range.end), + uncommitted_diff: None, }) } } @@ -1007,6 +1010,89 @@ impl EditPredictionStore { .unwrap_or_default() } + pub fn uncommitted_diffs_for_events( + &mut self, + project: Entity, + worktree_id: WorktreeId, + events: Vec, + cx: &mut Context, + ) -> Task, Entity>>> { + let project_id = project.entity_id(); + let git_store = project.read(cx).git_store().clone(); + cx.spawn(async move |this, cx| { + let mut diffs_by_path = HashMap::default(); + for stored_event in events.into_iter().rev() { + let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref(); + let Some(project_path) = project.read_with(cx, |project, cx| { + project + .find_project_path(path, cx) + .filter(|path| path.worktree_id == worktree_id) + }) else { + continue; + }; + let relative_path: Arc = project_path.path.as_std_path().into(); + if diffs_by_path.contains_key(&relative_path) { + continue; + } + + let old_snapshot_remote_id = stored_event.old_snapshot.remote_id(); + let new_snapshot_version = stored_event.new_snapshot_version.clone(); + let event_path = path.clone(); + if let Some(diff) = stored_event.uncommitted_diff.or_else(|| { + this.update(cx, |store, _| { + store + .projects + .get(&project_id)? + .events + .iter() + .find(|event| { + event.old_snapshot.remote_id() == old_snapshot_remote_id + && event.new_snapshot_version == new_snapshot_version + })? + .uncommitted_diff + .clone() + }) + .log_err() + .flatten() + }) { + diffs_by_path.insert(relative_path, diff); + continue; + } + + let buffer = project + .update(cx, |project, cx| project.open_buffer(project_path, cx)) + .await?; + let diff = git_store + .update(cx, |git_store, cx| { + git_store.open_uncommitted_diff(buffer.clone(), cx) + }) + .await?; + this.update(cx, |store, _| { + let Some(project) = store.projects.get_mut(&project_id) else { + return; + }; + let Some(target_index) = project.events.iter().position(|event| { + event.old_snapshot.remote_id() == old_snapshot_remote_id + && event.new_snapshot_version == new_snapshot_version + }) else { + return; + }; + for event in project.events.iter_mut().take(target_index + 1) { + let zeta_prompt::Event::BufferChange { path, .. } = event.event.as_ref(); + if path == &event_path { + event.uncommitted_diff = None; + } + } + project.events[target_index].uncommitted_diff = Some(diff.clone()); + }) + .log_err(); + diffs_by_path.insert(relative_path, diff); + } + + Ok(diffs_by_path) + }) + } + pub fn context_for_project<'a>( &'a self, project: &Entity, @@ -2475,6 +2561,7 @@ impl EditPredictionStore { && is_open_source && self.is_data_collection_enabled(cx) && matches!(self.edit_prediction_model, EditPredictionModel::Zeta); + let capture_worktree_id = snapshot.file().map(|file| file.worktree_id(cx)); let inputs = EditPredictionModelInput { project: project.clone(), buffer: active_buffer, @@ -2490,11 +2577,24 @@ impl EditPredictionStore { is_open_source, }; - let capture_data = (can_collect_data && rand::random_ratio(1, 1000)).then(|| stored_events); - let task = match self.edit_prediction_model { EditPredictionModel::Zeta => { - zeta::request_prediction_with_zeta(self, inputs, capture_data, cx) + let capture_events = (can_collect_data && rand::random_ratio(1, 10)) + .then(|| capture_worktree_id) + .flatten() + .map(|worktree_id| { + ( + stored_events.clone(), + self.uncommitted_diffs_for_events( + project.clone(), + worktree_id, + stored_events, + cx, + ), + ) + }); + + zeta::request_prediction_with_zeta(self, inputs, capture_events, cx) } EditPredictionModel::Fim { format } => fim::request_prediction(inputs, format, cx), EditPredictionModel::Mercury => { @@ -2815,6 +2915,42 @@ impl EditPredictionStore { project_state.recent_paths = paths.into_iter().collect(); } + pub fn recent_paths_for_project( + &self, + project: &Entity, + cx: &App, + ) -> (Vec, Vec) { + let recently_opened_files = project + .read(cx) + .opened_buffers(cx) + .into_iter() + .filter_map(|buffer| { + let snapshot = buffer.read(cx).snapshot(); + let file = snapshot.file()?; + Some(RecentFile { + path: file.path().as_std_path().into(), + cursor_position: None, + }) + }) + .collect(); + let recently_viewed_files = self + .projects + .get(&project.entity_id()) + .map(|project_state| { + project_state + .recent_paths + .iter() + .map(|path| RecentFile { + path: path.path.as_std_path().into(), + cursor_position: None, + }) + .collect() + }) + .unwrap_or_default(); + + (recently_opened_files, recently_viewed_files) + } + fn is_file_open_source( &self, project: &Entity, @@ -3035,6 +3171,7 @@ fn merge_trailing_events_if_needed( new_snapshot_version: newest_snapshot.version.clone(), total_edit_range: newest_snapshot.anchor_before(edit_range.start) ..newest_snapshot.anchor_before(edit_range.end), + uncommitted_diff: None, }, }; events.truncate(events.len() - mergeable_count); diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 86f218b401..140b551946 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -3603,6 +3603,9 @@ async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disable tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Path::new("foo.md").into(), cursor_position: "0".to_string(), edit_history: "sensitive edit history".to_string(), diff --git a/crates/edit_prediction/src/example_spec.rs b/crates/edit_prediction/src/example_spec.rs index a7da51173e..d3f4bbd3e8 100644 --- a/crates/edit_prediction/src/example_spec.rs +++ b/crates/edit_prediction/src/example_spec.rs @@ -13,6 +13,13 @@ pub const INLINE_CURSOR_MARKER: &str = "<|user_cursor|>"; /// falling back to git-based loading. pub const MAX_CURSOR_FILE_SIZE: usize = 64 * 1024; +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct RecentFile { + pub path: Arc, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cursor_position: Option, +} + #[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)] pub struct ExampleSpec { #[serde(default)] @@ -25,6 +32,12 @@ pub struct ExampleSpec { pub reasoning: Option, #[serde(default)] pub uncommitted_diff: String, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub recently_opened_files: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub recently_viewed_files: Vec, + #[serde(default, skip_serializing_if = "is_false")] + pub uncommitted_diff_contains_edit_history: bool, pub cursor_path: Arc, pub cursor_position: String, pub edit_history: String, @@ -56,18 +69,62 @@ pub struct TelemetrySource { const REASONING_HEADING: &str = "Reasoning"; const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff"; +const RECENTLY_OPENED_FILES_HEADING: &str = "Recently Opened Files"; +const RECENTLY_VIEWED_FILES_HEADING: &str = "Recently Viewed Files"; const EDIT_HISTORY_HEADING: &str = "Edit History"; const CURSOR_POSITION_HEADING: &str = "Cursor Position"; const EXPECTED_PATCH_HEADING: &str = "Expected Patch"; const REJECTED_PATCH_HEADING: &str = "Rejected Patch"; const ACCEPTED_PREDICTION_MARKER: &str = "// User accepted prediction:"; +fn write_path_list(markdown: &mut String, heading: &str, files: &[RecentFile]) { + if files.is_empty() { + return; + } + + _ = writeln!(markdown, "## {heading}"); + _ = writeln!(markdown); + _ = writeln!(markdown, "```"); + for file in files { + _ = write!(markdown, "{}", file.path.display()); + if let Some(position) = file.cursor_position { + _ = write!(markdown, "\t{position}"); + } + _ = writeln!(markdown); + } + _ = writeln!(markdown, "```"); + markdown.push('\n'); +} + +fn parse_path_list(text: &str) -> Vec { + text.lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(|line| { + let (path, cursor_position) = line + .rsplit_once('\t') + .map(|(path, position)| (path, position.parse().ok())) + .unwrap_or((line, None)); + RecentFile { + path: Path::new(path).into(), + cursor_position, + } + }) + .collect() +} + #[derive(Serialize, Deserialize)] struct FrontMatter<'a> { repository_url: Cow<'a, str>, revision: Cow<'a, str>, #[serde(default, skip_serializing_if = "Vec::is_empty")] tags: Vec, + #[serde(default, skip_serializing_if = "is_false")] + uncommitted_diff_requires_edit_history_rollback: bool, +} + +fn is_false(value: &bool) -> bool { + !*value } impl ExampleSpec { @@ -91,6 +148,8 @@ impl ExampleSpec { repository_url: Cow::Borrowed(&self.repository_url), revision: Cow::Borrowed(&self.revision), tags: self.tags.clone(), + uncommitted_diff_requires_edit_history_rollback: self + .uncommitted_diff_contains_edit_history, }; let front_matter_toml = toml::to_string_pretty(&front_matter).unwrap_or_else(|_| String::new()); @@ -130,6 +189,17 @@ impl ExampleSpec { markdown.push('\n'); } + write_path_list( + &mut markdown, + RECENTLY_OPENED_FILES_HEADING, + &self.recently_opened_files, + ); + write_path_list( + &mut markdown, + RECENTLY_VIEWED_FILES_HEADING, + &self.recently_viewed_files, + ); + _ = writeln!(markdown, "## {}", EDIT_HISTORY_HEADING); _ = writeln!(markdown); @@ -194,6 +264,9 @@ impl ExampleSpec { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Path::new("").into(), cursor_position: String::new(), edit_history: String::new(), @@ -211,6 +284,8 @@ impl ExampleSpec { spec.repository_url = data.repository_url.into_owned(); spec.revision = data.revision.into_owned(); spec.tags = data.tags; + spec.uncommitted_diff_contains_edit_history = + data.uncommitted_diff_requires_edit_history_rollback; } input = rest.trim_start(); } @@ -223,6 +298,8 @@ impl ExampleSpec { enum Section { Start, UncommittedDiff, + RecentlyOpenedFiles, + RecentlyViewedFiles, EditHistory, CursorPosition, ExpectedPatch, @@ -245,6 +322,10 @@ impl ExampleSpec { let title = mem::take(&mut text); current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) { Section::UncommittedDiff + } else if title.eq_ignore_ascii_case(RECENTLY_OPENED_FILES_HEADING) { + Section::RecentlyOpenedFiles + } else if title.eq_ignore_ascii_case(RECENTLY_VIEWED_FILES_HEADING) { + Section::RecentlyViewedFiles } else if title.eq_ignore_ascii_case(EDIT_HISTORY_HEADING) { Section::EditHistory } else if title.eq_ignore_ascii_case(CURSOR_POSITION_HEADING) { @@ -292,6 +373,14 @@ impl ExampleSpec { Section::UncommittedDiff => { spec.uncommitted_diff = mem::take(&mut text); } + Section::RecentlyOpenedFiles => { + spec.recently_opened_files = parse_path_list(&text); + text.clear(); + } + Section::RecentlyViewedFiles => { + spec.recently_viewed_files = parse_path_list(&text); + text.clear(); + } Section::EditHistory => { if next_edit_predicted { spec.edit_history @@ -481,6 +570,9 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Path::new("test.rs").into(), cursor_position: String::new(), edit_history: String::new(), @@ -617,6 +709,9 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Path::new("test.rs").into(), cursor_position: String::new(), edit_history: String::new(), @@ -689,6 +784,9 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Path::new("test.rs").into(), cursor_position: String::new(), edit_history: String::new(), diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index efa7552504..f20e5786de 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -1,11 +1,12 @@ use crate::{ CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, - EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, StoredEvent, + EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, ZedUpdateRequiredError, buffer_path_with_id_fallback, cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, prediction::EditPredictionResult, }; use anyhow::Result; +use buffer_diff::BufferDiff; use cloud_llm_client::{ AcceptEditPredictionBody, EditPredictionRejectReason, predict_edits_v3::RawCompletionRequest, }; @@ -48,7 +49,10 @@ pub fn request_prediction_with_zeta( is_open_source, .. }: EditPredictionModelInput, - capture_data: Option>, + capture_data: Option<( + Vec, + Task, Entity>>>, + )>, cx: &mut Context, ) -> Task>> { let settings = &all_language_settings(None, cx).edit_predictions; @@ -412,17 +416,33 @@ pub fn request_prediction_with_zeta( let editable_range_in_buffer = editable_range_in_buffer.clone(); let edit_preview = prediction.edit_preview.clone(); let model_version = prediction.model_version.clone(); - let example_task = capture_data.and_then(|stored_events| { - cx.update(|cx| { - crate::capture_example( - project.clone(), - edited_buffer.clone(), - position, - stored_events, - false, - cx, - ) - }) + let example_task = capture_data.and_then(|(events, uncommitted_diffs)| { + let (recently_opened_files, recently_viewed_files) = this + .read_with(cx, |this, cx| this.recent_paths_for_project(&project, cx)) + .ok()?; + Some(cx.spawn({ + let project = project.clone(); + let edited_buffer = edited_buffer.clone(); + async move |cx| { + let uncommitted_diffs = uncommitted_diffs.await?; + let Some(task) = cx.update(|cx| { + crate::capture_example::capture_example( + project.clone(), + edited_buffer.clone(), + position, + events, + recently_opened_files, + recently_viewed_files, + uncommitted_diffs, + false, + cx, + ) + }) else { + return Err(anyhow::anyhow!("failed to capture example")); + }; + task.await + } + })) }); cx.spawn(async move |cx| { let example_spec = if let Some(task) = example_task { diff --git a/crates/edit_prediction_cli/src/format_prompt.rs b/crates/edit_prediction_cli/src/format_prompt.rs index f04b263059..e0354a78ac 100644 --- a/crates/edit_prediction_cli/src/format_prompt.rs +++ b/crates/edit_prediction_cli/src/format_prompt.rs @@ -854,6 +854,9 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")), cursor_position: "0:0".to_string(), edit_history: String::new(), @@ -933,6 +936,9 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: std::sync::Arc::from(std::path::Path::new("src/main.rs")), cursor_position: "0:0".to_string(), edit_history: String::new(), diff --git a/crates/edit_prediction_cli/src/pull_examples.rs b/crates/edit_prediction_cli/src/pull_examples.rs index 9062d973ff..88de3053d8 100644 --- a/crates/edit_prediction_cli/src/pull_examples.rs +++ b/crates/edit_prediction_cli/src/pull_examples.rs @@ -1754,6 +1754,9 @@ fn build_example_from_snowflake( tags, reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: input.cursor_path.clone(), cursor_position: build_cursor_position(cursor_excerpt, cursor_offset), edit_history, diff --git a/crates/edit_prediction_cli/src/repair.rs b/crates/edit_prediction_cli/src/repair.rs index 2ae62fd70f..751e55a8b1 100644 --- a/crates/edit_prediction_cli/src/repair.rs +++ b/crates/edit_prediction_cli/src/repair.rs @@ -549,6 +549,9 @@ mod tests { tags: Vec::new(), reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Arc::from(Path::new("src/main.rs")), cursor_position: "0:0".to_string(), edit_history: String::new(), diff --git a/crates/edit_prediction_cli/src/split_commit.rs b/crates/edit_prediction_cli/src/split_commit.rs index b70ac354b5..2421a81f3c 100644 --- a/crates/edit_prediction_cli/src/split_commit.rs +++ b/crates/edit_prediction_cli/src/split_commit.rs @@ -370,6 +370,9 @@ pub fn generate_evaluation_example_from_ordered_commit( tags: vec![], reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, rejected_patch: None, telemetry: None, @@ -1369,6 +1372,9 @@ Date: Mon Jan 1 00:00:00 2024 tags: vec![], reasoning: None, uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, rejected_patch: None, telemetry: None, diff --git a/crates/edit_prediction_cli/src/synthesize.rs b/crates/edit_prediction_cli/src/synthesize.rs index 228690ae49..76c8e0fee3 100644 --- a/crates/edit_prediction_cli/src/synthesize.rs +++ b/crates/edit_prediction_cli/src/synthesize.rs @@ -790,6 +790,9 @@ async fn build_example( tags: Vec::new(), reasoning: Some(reasoning_with_source), uncommitted_diff: String::new(), + recently_opened_files: Vec::new(), + recently_viewed_files: Vec::new(), + uncommitted_diff_contains_edit_history: false, cursor_path: Arc::from(Path::new(&cursor_file)), cursor_position: String::new(), edit_history, diff --git a/crates/edit_prediction_ui/src/edit_prediction_button.rs b/crates/edit_prediction_ui/src/edit_prediction_button.rs index 2246f106e7..cd3164d2ee 100644 --- a/crates/edit_prediction_ui/src/edit_prediction_button.rs +++ b/crates/edit_prediction_ui/src/edit_prediction_button.rs @@ -42,9 +42,7 @@ use workspace::{ }; use zed_actions::{OpenBrowser, OpenSettingsAt}; -use crate::{ - CaptureExample, RatePredictions, rate_prediction_modal::PredictEditsRatePredictionsFeatureFlag, -}; +use crate::{RatePredictions, rate_prediction_modal::PredictEditsRatePredictionsFeatureFlag}; actions!( edit_prediction, @@ -982,10 +980,7 @@ impl EditPredictionButton { .context(editor_focus_handle) .when( cx.has_flag::(), - |this| { - this.action("Capture Prediction Example", CaptureExample.boxed_clone()) - .action("Rate Predictions", RatePredictions.boxed_clone()) - }, + |this| this.action("Rate Predictions", RatePredictions.boxed_clone()), ); } diff --git a/crates/edit_prediction_ui/src/edit_prediction_ui.rs b/crates/edit_prediction_ui/src/edit_prediction_ui.rs index 05f1224f50..2656375457 100644 --- a/crates/edit_prediction_ui/src/edit_prediction_ui.rs +++ b/crates/edit_prediction_ui/src/edit_prediction_ui.rs @@ -3,13 +3,10 @@ mod edit_prediction_context_view; mod rate_prediction_modal; use command_palette_hooks::CommandPaletteFilter; -use edit_prediction::{EditPredictionStore, ResetOnboarding, capture_example}; +use edit_prediction::ResetOnboarding; use edit_prediction_context_view::EditPredictionContextView; -use editor::Editor; use feature_flags::FeatureFlagAppExt as _; -use gpui::TaskExt; use gpui::actions; -use language::language_settings::AllLanguageSettings; use project::DisableAiSettings; use rate_prediction_modal::RatePredictionsModal; use settings::{Settings as _, SettingsStore}; @@ -36,8 +33,6 @@ actions!( [ /// Opens the rate completions modal. RatePredictions, - /// Captures an ExampleSpec from the current editing session and opens it as Markdown. - CaptureExample, ] ); @@ -51,9 +46,6 @@ pub fn init(cx: &mut App) { } }); - workspace.register_action(|workspace, _: &CaptureExample, window, cx| { - capture_example_as_markdown(workspace, window, cx); - }); workspace.register_action_renderer(|div, _, _, cx| { div.on_action(cx.listener( move |workspace, _: &OpenEditPredictionContextView, window, cx| { @@ -84,7 +76,6 @@ fn feature_gate_predict_edits_actions(cx: &mut App) { let reset_onboarding_action_types = [TypeId::of::()]; let all_action_types = [ TypeId::of::(), - TypeId::of::(), TypeId::of::(), zed_actions::OpenZedPredictOnboarding.type_id(), TypeId::of::(), @@ -131,68 +122,3 @@ fn feature_gate_predict_edits_actions(cx: &mut App) { }) .detach(); } - -fn capture_example_as_markdown( - workspace: &mut Workspace, - window: &mut Window, - cx: &mut Context, -) -> Option<()> { - let markdown_language = workspace - .app_state() - .languages - .language_for_name("Markdown"); - - let fs = workspace.app_state().fs.clone(); - let project = workspace.project().clone(); - let editor = workspace.active_item_as::(cx)?; - let editor = editor.read(cx); - let (buffer, cursor_anchor) = editor - .buffer() - .read(cx) - .text_anchor_for_position(editor.selections.newest_anchor().head(), cx)?; - let ep_store = EditPredictionStore::try_global(cx)?; - let events = ep_store.update(cx, |store, cx| store.edit_history_for_project(&project, cx)); - let example = capture_example(project.clone(), buffer, cursor_anchor, events, true, cx)?; - - let examples_dir = AllLanguageSettings::get_global(cx) - .edit_predictions - .examples_dir - .clone(); - - cx.spawn_in(window, async move |workspace_entity, cx| { - let markdown_language = markdown_language.await?; - let example_spec = example.await?; - let buffer = if let Some(dir) = examples_dir { - fs.create_dir(&dir).await.ok(); - let mut path = dir.join(&example_spec.name.replace(' ', "--").replace(':', "-")); - path.set_extension("md"); - project - .update(cx, |project, cx| project.open_local_buffer(&path, cx)) - .await? - } else { - project - .update(cx, |project, cx| { - project.create_buffer(Some(markdown_language.clone()), false, cx) - }) - .await? - }; - - buffer.update(cx, |buffer, cx| { - buffer.set_text(example_spec.to_markdown(), cx); - buffer.set_language(Some(markdown_language), cx); - }); - workspace_entity.update_in(cx, |workspace, window, cx| { - workspace.add_item_to_active_pane( - Box::new( - cx.new(|cx| Editor::for_buffer(buffer, Some(project.clone()), window, cx)), - ), - None, - true, - window, - cx, - ); - }) - }) - .detach_and_log_err(cx); - None -} diff --git a/crates/language/src/language_settings.rs b/crates/language/src/language_settings.rs index 3d90d8d06e..de67375ccc 100644 --- a/crates/language/src/language_settings.rs +++ b/crates/language/src/language_settings.rs @@ -477,7 +477,6 @@ pub struct EditPredictionSettings { /// Settings specific to Ollama. pub ollama: Option, pub open_ai_compatible_api: Option, - pub examples_dir: Option>, /// Controls whether training data collection is enabled. /// /// `Default` means the value stored in the legacy KV store is used as a fallback, @@ -911,7 +910,6 @@ impl settings::Settings for AllLanguageSettings { codestral: codestral_settings, ollama: ollama_settings, open_ai_compatible_api: openai_compatible_settings, - examples_dir: edit_predictions.examples_dir, allow_data_collection: edit_predictions.allow_data_collection.unwrap_or_default(), }, defaults: default_language_settings, diff --git a/crates/settings_content/src/language.rs b/crates/settings_content/src/language.rs index 081406a684..2e5ef23875 100644 --- a/crates/settings_content/src/language.rs +++ b/crates/settings_content/src/language.rs @@ -1,4 +1,4 @@ -use std::{num::NonZeroU32, path::Path}; +use std::num::NonZeroU32; use collections::{HashMap, HashSet}; use schemars::JsonSchema; @@ -137,8 +137,6 @@ pub struct EditPredictionSettingsContent { pub ollama: Option, /// Settings specific to using custom OpenAI-compatible servers for edit prediction. pub open_ai_compatible_api: Option, - /// The directory where manually captured edit prediction examples are stored. - pub examples_dir: Option>, /// Controls whether Zed may collect training data when using Zed's Edit Predictions. /// Data is only ever captured for files in projects that are detected as open source. ///