diff --git a/Cargo.lock b/Cargo.lock index 9bb127b905..c08e373793 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3265,6 +3265,8 @@ dependencies = [ "serde", "serde_json", "strum 0.27.2", + "uuid", + "zeta_prompt", ] [[package]] @@ -5569,6 +5571,7 @@ dependencies = [ "feature_flags", "fs", "futures 0.3.32", + "git", "gpui", "heapless", "indoc", diff --git a/crates/cloud_api_types/Cargo.toml b/crates/cloud_api_types/Cargo.toml index 13f6f0833f..66080cf9ec 100644 --- a/crates/cloud_api_types/Cargo.toml +++ b/crates/cloud_api_types/Cargo.toml @@ -19,6 +19,8 @@ cloud_llm_client.workspace = true serde.workspace = true serde_json.workspace = true strum.workspace = true +uuid.workspace = true +zeta_prompt.workspace = true [dev-dependencies] pretty_assertions.workspace = true diff --git a/crates/cloud_api_types/src/cloud_api_types.rs b/crates/cloud_api_types/src/cloud_api_types.rs index 0422c4fa03..e5a67e68ef 100644 --- a/crates/cloud_api_types/src/cloud_api_types.rs +++ b/crates/cloud_api_types/src/cloud_api_types.rs @@ -6,6 +6,7 @@ mod timestamp; pub mod websocket_protocol; use std::collections::BTreeMap; +use std::path::Path; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -148,6 +149,42 @@ pub struct SubmitEditPredictionSettledBody { #[derive(Debug, PartialEq, Serialize, Deserialize)] pub struct SubmitEditPredictionSettledResponse {} +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub struct SubmitEditPredictionJumpExampleBody { + pub request_id: uuid::Uuid, + pub trigger: JumpExampleTrigger, + pub repository_url: Option, + pub revision: Option, + pub uncommitted_diff: Option, + pub recently_opened_files: Vec, + pub recently_viewed_files: Vec, + pub cursor_path: Arc, + pub cursor_position: String, + pub edit_history: Vec>, + pub diagnostics: Vec, + pub future_edit_history: String, + pub navigation_history: Vec, + pub is_in_open_source_repo: bool, + pub can_collect_data: bool, +} + +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub struct SubmitEditPredictionJumpExampleResponse {} + +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum JumpExampleTrigger { + Prediction, + Diagnostic, +} + +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub struct JumpExampleRecentFile { + pub path: Arc, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cursor_position: Option, +} + #[derive(Debug, PartialEq, Serialize, Deserialize)] pub struct EditPredictionSettledKeptChars { #[serde(rename = "edit_bytes_candidate_new")] diff --git a/crates/edit_prediction/Cargo.toml b/crates/edit_prediction/Cargo.toml index ea11fac3cd..49b4da1c1e 100644 --- a/crates/edit_prediction/Cargo.toml +++ b/crates/edit_prediction/Cargo.toml @@ -35,6 +35,7 @@ edit_prediction_metrics = { workspace = true, features = ["tree-sitter"] } feature_flags.workspace = true fs.workspace = true futures.workspace = true +git.workspace = true gpui.workspace = true indoc.workspace = true itertools.workspace = true diff --git a/crates/edit_prediction/src/capture_example.rs b/crates/edit_prediction/src/capture_example.rs index c01cd93578..1f5ccf8ed5 100644 --- a/crates/edit_prediction/src/capture_example.rs +++ b/crates/edit_prediction/src/capture_example.rs @@ -1,15 +1,16 @@ use crate::{ StoredEvent, + data_collection::{ + UncommittedDiffSnapshot, compute_cursor_excerpt, compute_uncommitted_diff, + format_cursor_excerpt, + }, example_spec::{ExampleSpec, RecentFile}, }; use anyhow::Result; -use buffer_diff::{BufferDiff, BufferDiffSnapshot}; -use collections::HashMap; use gpui::{App, Entity, Task}; use language::Buffer; use project::Project; -use std::{collections::hash_map, fmt::Write as _, ops::Range, path::Path, sync::Arc}; -use text::Point; +use std::{fmt::Write, path::Path, sync::Arc}; pub fn capture_example( project: Entity, @@ -18,7 +19,7 @@ pub fn capture_example( events: Vec, recently_opened_files: Vec, recently_viewed_files: Vec, - uncommitted_diffs_by_path: HashMap, Entity>, + uncommitted_diff_snapshot: UncommittedDiffSnapshot, populate_expected_patch: bool, cx: &mut App, ) -> Option>> { @@ -41,39 +42,10 @@ pub fn capture_example( let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string(); 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; - }; - - 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); - diff_buffers_by_path.contains_key(relative_path) - }); + let uncommitted_diff = cx + .background_executor() + .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) + .await; let line_comment_prefix = snapshot .language() @@ -85,19 +57,6 @@ 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(uncommitted_diff_snapshots) }) - .await; - let mut edit_history = String::new(); for stored_event in &events { write_event_with_relative_paths(&mut edit_history, &stored_event.event, &root_name); @@ -131,7 +90,7 @@ pub fn capture_example( rejected_patch = Some(empty_patch); } - let mut spec = ExampleSpec { + let spec = ExampleSpec { name: generate_timestamp_name(), repository_url, revision, @@ -142,7 +101,11 @@ pub fn capture_example( recently_viewed_files, uncommitted_diff_contains_edit_history, cursor_path, - cursor_position: String::new(), + cursor_position: format_cursor_excerpt( + &cursor_excerpt, + cursor_offset_in_excerpt, + &line_comment_prefix, + ), edit_history, expected_patches, rejected_patch, @@ -150,26 +113,17 @@ pub fn capture_example( human_feedback: Vec::new(), rating: None, }; - spec.set_cursor_excerpt( - &cursor_excerpt, - cursor_offset_in_excerpt, - &line_comment_prefix, - ); Ok(spec) })) } -fn strip_root_name<'a>(path: &'a Path, root_name: &str) -> &'a Path { - path.strip_prefix(root_name).unwrap_or(path) -} - -fn write_event_with_relative_paths( +pub(crate) fn write_event_with_relative_paths( output: &mut String, event: &zeta_prompt::Event, root_name: &str, ) { fn write_relative_path(output: &mut String, path: &Path, root_name: &str) { - for component in strip_root_name(path, root_name).components() { + for component in path.strip_prefix(root_name).unwrap_or(path).components() { output.push('/'); write!(output, "{}", component.as_os_str().to_string_lossy()).ok(); } @@ -190,142 +144,6 @@ fn write_event_with_relative_paths( output.push_str(diff); } -fn compute_cursor_excerpt( - snapshot: &language::BufferSnapshot, - cursor_anchor: language::Anchor, -) -> (String, usize, Range) { - use text::ToOffset as _; - use text::ToPoint as _; - - let cursor_offset = cursor_anchor.to_offset(snapshot); - let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = - crate::cursor_excerpt::compute_cursor_excerpt(snapshot, cursor_offset); - let syntax_ranges = crate::cursor_excerpt::compute_syntax_ranges( - snapshot, - cursor_offset, - &excerpt_offset_range, - ); - let excerpt_text: String = snapshot.text_for_range(excerpt_point_range).collect(); - let (_, context_range) = zeta_prompt::compute_editable_and_context_ranges( - &excerpt_text, - cursor_offset_in_excerpt, - &syntax_ranges, - 100, - 50, - ); - let context_text = excerpt_text[context_range.clone()].to_string(); - let cursor_in_context = cursor_offset_in_excerpt.saturating_sub(context_range.start); - let context_buffer_start = - (excerpt_offset_range.start + context_range.start).to_point(snapshot); - let context_buffer_end = (excerpt_offset_range.start + context_range.end).to_point(snapshot); - ( - context_text, - cursor_in_context, - context_buffer_start..context_buffer_end, - ) -} - -fn compute_uncommitted_diff( - snapshots_by_path: HashMap, (language::BufferSnapshot, BufferDiffSnapshot)>, -) -> String { - let mut uncommitted_diff = String::new(); - 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 { @@ -343,6 +161,7 @@ fn generate_timestamp_name() -> String { mod tests { use super::*; use crate::EditPredictionStore; + use crate::data_collection::uncommitted_diffs_for_events; use client::RefreshLlmTokenListener; use client::{Client, UserStore}; use clock::FakeSystemClock; @@ -535,19 +354,39 @@ mod tests { ); let worktree_id = buffer.read_with(cx, |buffer, cx| buffer.file().unwrap().worktree_id(cx)); + let failed_capture = ep_store + .update(cx, |_store, cx| { + uncommitted_diffs_for_events(project.clone(), worktree_id, events.clone(), cx) + }) + .await; + assert!(failed_capture.is_err()); + + let project_events = events + .into_iter() + .filter(|event| { + let zeta_prompt::Event::BufferChange { path, .. } = event.event.as_ref(); + path.as_ref() != "/external/external.rs" + }) + .collect::>(); let uncommitted_diffs_by_path = ep_store - .update(cx, |store, cx| { - store.uncommitted_diffs_for_events(project.clone(), worktree_id, events.clone(), cx) + .update(cx, |_store, cx| { + uncommitted_diffs_for_events( + project.clone(), + worktree_id, + project_events.clone(), + cx, + ) }) .await .unwrap(); + let mut example = cx .update(|cx| { capture_example( project.clone(), buffer.clone(), Anchor::min_for_buffer(buffer.read(cx).remote_id()), - events, + project_events, Vec::new(), Vec::new(), uncommitted_diffs_by_path, diff --git a/crates/edit_prediction/src/data_collection.rs b/crates/edit_prediction/src/data_collection.rs new file mode 100644 index 0000000000..ae8a3bca72 --- /dev/null +++ b/crates/edit_prediction/src/data_collection.rs @@ -0,0 +1,290 @@ +use crate::{EditPredictionStore, StoredEvent}; + +use anyhow::{Context as _, Result}; +use buffer_diff::BufferDiffSnapshot; +use collections::HashMap; +use gpui::{Context, Entity, Task}; +use language::BufferSnapshot; +use project::{Project, WorktreeId}; +use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc}; +use text::{OffsetRangeExt, Point}; + +pub type UncommittedDiffSnapshot = Vec<(Arc, BufferSnapshot, BufferDiffSnapshot)>; +pub type UncommittedDiffResult = std::result::Result>; + +pub use zeta_prompt::udiff::CURSOR_POSITION_MARKER; + +pub fn uncommitted_diffs_for_events( + project: Entity, + worktree_id: WorktreeId, + events: Vec, + cx: &Context<'_, EditPredictionStore>, +) -> Task { + let git_store = project.read_with(cx, |project, _| project.git_store().clone()); + + cx.spawn(async move |_store, cx| { + let events_with_paths = events + .into_iter() + .map(|stored_event| { + let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref(); + 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)) + }) + .map(|(project_path, relative_path)| { + (stored_event, project_path, relative_path) + }) + .context("failed to find project path for uncommitted diff capture") + }) + .collect::>>() + .map_err(Arc::new)?; + + let mut snapshots_by_path: HashMap, (BufferSnapshot, BufferDiffSnapshot)> = + HashMap::default(); + for (stored_event, project_path, relative_path) in events_with_paths.iter().rev() { + if snapshots_by_path.contains_key(relative_path) { + continue; + } + + let buffer = project + .update(cx, |project, cx| { + project.open_buffer(project_path.clone(), cx) + }) + .await + .context("failed to open buffer for uncommitted diff capture") + .map_err(Arc::new)?; + let file_context = stored_event.file_context.clone(); + let cached_diff = file_context.as_ref().and_then(|file_context| { + file_context.read_with(cx, |file_context, _| file_context.uncommitted_diff.clone()) + }); + let diff = match cached_diff { + Some(diff) => diff, + None => { + let diff = git_store + .update(cx, |git_store, cx| { + git_store.open_uncommitted_diff(buffer.clone(), cx) + }) + .await + .context("failed to open uncommitted diff for capture") + .map_err(Arc::new)?; + if let Some(file_context) = file_context { + file_context.update(cx, |file_context, _| { + file_context.uncommitted_diff = Some(diff.clone()); + }); + } + diff + } + }; + + let buffer_snapshot = buffer.update(cx, |buffer, _| buffer.snapshot()); + let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx)); + snapshots_by_path.insert(relative_path.clone(), (buffer_snapshot, diff_snapshot)); + } + + let uncommitted_diff_snapshots = snapshots_by_path + .into_iter() + .map(|(relative_path, (snapshot, diff_snapshot))| { + (relative_path, snapshot, diff_snapshot) + }) + .collect(); + + Ok(uncommitted_diff_snapshots) + }) +} + +pub fn compute_cursor_excerpt( + snapshot: &language::BufferSnapshot, + cursor_anchor: language::Anchor, +) -> (String, usize, Range) { + use text::ToOffset as _; + use text::ToPoint as _; + + let cursor_offset = cursor_anchor.to_offset(snapshot); + let (excerpt_point_range, excerpt_offset_range, cursor_offset_in_excerpt) = + crate::cursor_excerpt::compute_cursor_excerpt(snapshot, cursor_offset); + let syntax_ranges = crate::cursor_excerpt::compute_syntax_ranges( + snapshot, + cursor_offset, + &excerpt_offset_range, + ); + let excerpt_text: String = snapshot.text_for_range(excerpt_point_range).collect(); + let (_, context_range) = zeta_prompt::compute_editable_and_context_ranges( + &excerpt_text, + cursor_offset_in_excerpt, + &syntax_ranges, + 100, + 50, + ); + let context_text = excerpt_text[context_range.clone()].to_string(); + let cursor_in_context = cursor_offset_in_excerpt.saturating_sub(context_range.start); + let context_buffer_start = + (excerpt_offset_range.start + context_range.start).to_point(snapshot); + let context_buffer_end = (excerpt_offset_range.start + context_range.end).to_point(snapshot); + ( + context_text, + cursor_in_context, + context_buffer_start..context_buffer_end, + ) +} + +pub(crate) fn compute_uncommitted_diff(snapshot: UncommittedDiffSnapshot) -> String { + let mut uncommitted_diff = String::new(); + let mut snapshots_by_path = snapshot; + 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 +} + +pub(crate) fn estimate_uncommitted_diff_byte_size(snapshot: &UncommittedDiffSnapshot) -> usize { + let mut size = 0; + for (_, buffer_snapshot, diff_snapshot) in snapshot { + for hunk in diff_snapshot.hunks(buffer_snapshot) { + size += hunk.diff_base_byte_range.len(); + size += hunk.range.to_offset(buffer_snapshot).len(); + } + } + size +} + +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 + } +} + +pub fn format_cursor_excerpt( + excerpt: &str, + cursor_offset: usize, + line_comment_prefix: &str, +) -> String { + let cursor_line_start = excerpt[..cursor_offset] + .rfind('\n') + .map(|pos| pos + 1) + .unwrap_or(0); + let cursor_line_end = excerpt[cursor_line_start..] + .find('\n') + .map(|pos| cursor_line_start + pos + 1) + .unwrap_or(excerpt.len()); + let cursor_line = &excerpt[cursor_line_start..cursor_line_end]; + let cursor_line_indent = &cursor_line[..cursor_line.len() - cursor_line.trim_start().len()]; + let cursor_column = cursor_offset - cursor_line_start; + + let mut marker_line = String::new(); + if cursor_column < line_comment_prefix.len() { + for _ in 0..cursor_column { + marker_line.push(' '); + } + marker_line.push_str(line_comment_prefix); + write!(marker_line, " <{}", CURSOR_POSITION_MARKER).unwrap(); + } else { + if cursor_column >= cursor_line_indent.len() + line_comment_prefix.len() { + marker_line.push_str(cursor_line_indent); + } + marker_line.push_str(line_comment_prefix); + while marker_line.len() < cursor_column { + marker_line.push(' '); + } + write!(marker_line, "^{}", CURSOR_POSITION_MARKER).unwrap(); + } + + let mut result = String::with_capacity(excerpt.len() + marker_line.len() + 2); + result.push_str(&excerpt[..cursor_line_end]); + if !result.ends_with('\n') { + result.push('\n'); + } + result.push_str(&marker_line); + if cursor_line_end < excerpt.len() { + result.push('\n'); + result.push_str(&excerpt[cursor_line_end..]); + } + result +} diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 53bd4bfbbc..ab88dbf2fa 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -1,4 +1,4 @@ -use anyhow::Result; +use anyhow::{Context as _, Result}; use buffer_diff::BufferDiff; use client::{Client, EditPredictionUsage, UserStore, global_llm_token}; use cloud_api_client::LlmApiToken; @@ -28,11 +28,12 @@ use futures::{ channel::mpsc::{self, UnboundedReceiver}, select_biased, }; +use git::repository::FileHistoryChangedFileSets; use gpui::BackgroundExecutor; use gpui::TaskExt; use gpui::http_client::Url; use gpui::{ - App, AsyncApp, Entity, EntityId, Global, SharedString, Task, WeakEntity, actions, + App, AsyncApp, Context, Entity, EntityId, Global, SharedString, Task, WeakEntity, actions, http_client::{self, AsyncBody, Method}, prelude::*, }; @@ -51,6 +52,7 @@ use settings::{ }; use std::collections::{VecDeque, hash_map}; use std::env; +use std::rc::Rc; use text::{AnchorRangeExt, Edit}; use workspace::{AppState, Workspace}; use zeta_prompt::{ZetaFormat, ZetaPromptInput}; @@ -58,7 +60,6 @@ use zeta_prompt::{ZetaFormat, ZetaPromptInput}; use std::mem; use std::ops::Range; use std::path::Path; -use std::rc::Rc; use std::str::FromStr as _; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -67,8 +68,10 @@ use thiserror::Error; use util::{RangeExt as _, ResultExt as _}; pub mod cursor_excerpt; +pub mod data_collection; pub mod example_spec; pub mod fim; +mod jump_example; mod license_detection; pub mod mercury; pub mod metrics; @@ -88,8 +91,13 @@ pub mod zeta; mod edit_prediction_tests; use crate::cursor_excerpt::expand_context_syntactically_then_linewise; +use crate::data_collection::uncommitted_diffs_for_events; use crate::example_spec::ExampleSpec; use crate::example_spec::RecentFile; +use crate::jump_example::{ + JUMP_EXAMPLE_NAVIGATION_COUNT, JumpExampleTrigger, PendingJumpExampleCapture, + PendingJumpExampleCaptureKey, +}; use crate::license_detection::LicenseDetectionWatcher; use crate::mercury::Mercury; pub use crate::metrics::{KeptRateResult, compute_kept_rate}; @@ -113,9 +121,11 @@ actions!( /// Maximum number of events to track. const EVENT_COUNT_MAX: usize = 10; +const RECENT_PATH_COUNT_MAX: usize = 20; const CHANGE_GROUPING_LINE_SPAN: u32 = 8; const EDIT_HISTORY_DIFF_SIZE_LIMIT: usize = 2048 * 3; // ~2048 tokens or ~50% of typical prompt budget const COLLABORATOR_EDIT_LOCALITY_CONTEXT_TOKENS: usize = 512; +const GIT_CHANGED_FILE_SETS_COMMIT_LIMIT: usize = 100; const LAST_CHANGE_GROUPING_TIME: Duration = Duration::from_secs(1); const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice"; const REJECT_REQUEST_DEBOUNCE: Duration = Duration::from_secs(15); @@ -187,6 +197,7 @@ pub struct EditPredictionModelInput { snapshot: BufferSnapshot, position: Anchor, events: Vec>, + stored_events: Vec, related_files: Vec, mode: PredictEditsMode, trigger: PredictEditsRequestTrigger, @@ -239,7 +250,13 @@ pub struct StoredEvent { pub old_snapshot: TextBufferSnapshot, pub new_snapshot_version: clock::Global, pub total_edit_range: Range, - pub uncommitted_diff: Option>, + pub(crate) file_context: Option>, +} + +pub(crate) struct StoredFileContext { + pub(crate) uncommitted_diff: Option>, + pub(crate) git_changed_file_sets: Option>, + pub(crate) git_changed_file_sets_task: Option>, } impl StoredEvent { @@ -318,15 +335,30 @@ fn lines_between_ranges(left: &Range, right: &Range) -> u32 { 0 } +fn push_recent_file(files: &mut VecDeque, mut file: RecentFile) { + if let Some(ix) = files.iter().position(|probe| probe.path == file.path) + && let Some(previous) = files.remove(ix) + && file.cursor_position.is_none() + { + file.cursor_position = previous.cursor_position; + } + files.push_front(file); + files.truncate(RECENT_PATH_COUNT_MAX); +} + struct ProjectState { events: VecDeque, last_event: Option, - recent_paths: VecDeque, + recently_viewed_files: VecDeque, + recently_opened_files: VecDeque, registered_buffers: HashMap, + file_contexts: HashMap>, current_prediction: Option, last_edit_source: Option, next_pending_prediction_id: usize, pending_predictions: ArrayVec, + pending_jump_example_captures: Vec, + starting_jump_example_captures: Vec, debug_tx: Option>, last_edit_prediction_refresh: Option<(EntityId, Instant)>, last_jump_prediction_refresh: Option<(EntityId, Instant)>, @@ -393,6 +425,59 @@ impl ProjectState { let registered_buffer = self.registered_buffers.get(&active_buffer.entity_id())?; Some((active_buffer, registered_buffer.last_position)) } + + fn file_context_for_path( + &mut self, + path: ProjectPath, + cx: &mut Context, + ) -> Entity { + if let Some(context) = self + .file_contexts + .get_mut(&path) + .and_then(|entry| entry.upgrade()) + { + context + } else { + let context = cx.new(|_| StoredFileContext { + uncommitted_diff: None, + git_changed_file_sets: None, + git_changed_file_sets_task: None, + }); + self.file_contexts.insert(path, context.downgrade()); + context + } + } + + fn update_recent_file_cursor(&mut self, path: &Path, cursor_position: usize) { + for file in &mut self.recently_opened_files { + if file.path.as_ref() == path && file.cursor_position.is_none() { + file.cursor_position = Some(cursor_position); + } + } + for file in &mut self.recently_viewed_files { + if file.path.as_ref() == path { + file.cursor_position = Some(cursor_position); + } + } + } + + fn finalize_last_event(&mut self, cx: &mut Context) { + let Some(event) = self.last_event.take() else { + return; + }; + let Some(event) = event.finalize(&self.license_detection_watchers, cx) else { + return; + }; + + for capture in &mut self.pending_jump_example_captures { + capture.future_events.push(event.event.clone()); + } + jump_example::drain_completed_jump_example_captures(self, cx); + if self.events.len() + 1 >= EVENT_COUNT_MAX { + self.events.pop_front(); + } + self.events.push_back(event); + } } #[derive(Debug, Clone)] @@ -532,6 +617,7 @@ struct LastEvent { predicted: bool, snapshot_after_last_editing_pause: Option, last_edit_time: Option, + file_context: Option>, } impl LastEvent { @@ -575,7 +661,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, + file_context: self.file_context.clone(), }) } } @@ -610,6 +696,7 @@ impl LastEvent { predicted: self.predicted, snapshot_after_last_editing_pause: None, last_edit_time: self.last_edit_time, + file_context: self.file_context.clone(), }; let after = LastEvent { @@ -623,6 +710,7 @@ impl LastEvent { predicted: self.predicted, snapshot_after_last_editing_pause: None, last_edit_time: self.last_edit_time, + file_context: self.file_context.clone(), }; (before, Some(after)) @@ -1011,89 +1099,6 @@ 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, @@ -1186,10 +1191,77 @@ impl EditPredictionStore { project: &Entity, cx: &mut Context, ) { + let opened_path = buffer + .read(cx) + .file() + .map(|file| ProjectPath::from_file(file.as_ref(), cx)); let project_state = self.get_or_init_project(project, cx); + if let Some(path) = opened_path { + push_recent_file( + &mut project_state.recently_opened_files, + RecentFile { + path: path.path.as_std_path().into(), + cursor_position: None, + }, + ); + } Self::register_buffer_impl(project_state, buffer, project, cx); } + fn ensure_git_changed_file_sets_loading( + file_context: &Entity, + project: &Entity, + project_path: &ProjectPath, + cx: &mut Context, + ) { + let should_start = file_context.update(cx, |file_context, _| { + file_context.git_changed_file_sets.is_none() + && file_context.git_changed_file_sets_task.is_none() + }); + if !should_start { + return; + } + + let Some((repository, repo_path)) = project + .read(cx) + .git_store() + .read(cx) + .repository_and_path_for_project_path(project_path, cx) + else { + file_context.update(cx, |file_context, _| { + file_context.git_changed_file_sets = Some(Arc::default()); + }); + return; + }; + + let receiver = repository.update(cx, |repository, _| { + repository + .file_history_changed_files(vec![repo_path], GIT_CHANGED_FILE_SETS_COMMIT_LIMIT) + }); + let task = cx.spawn({ + let file_context = file_context.downgrade(); + async move |_, cx| { + let result = receiver.await; + let Some(file_context) = file_context.upgrade() else { + return; + }; + file_context.update(cx, |file_context, _| { + file_context.git_changed_file_sets = result + .context("failed to receive git changed file sets") + .flatten() + .map(|mut file_sets| file_sets.pop().unwrap_or_default()) + .context("failed to load git changed file sets") + .map(Arc::new) + .log_err(); + file_context.git_changed_file_sets_task = None; + }); + } + }); + file_context.update(cx, |file_context, _| { + file_context.git_changed_file_sets_task = Some(task); + }); + } + fn get_or_init_project( &mut self, project: &Entity, @@ -1209,13 +1281,17 @@ impl EditPredictionStore { }, events: VecDeque::new(), last_event: None, - recent_paths: VecDeque::new(), + recently_viewed_files: VecDeque::new(), + recently_opened_files: VecDeque::new(), debug_tx: None, registered_buffers: HashMap::default(), + file_contexts: HashMap::default(), current_prediction: None, last_edit_source: None, cancelled_predictions: HashSet::default(), pending_predictions: ArrayVec::new(), + pending_jump_example_captures: Vec::new(), + starting_jump_example_captures: Vec::new(), next_pending_prediction_id: 0, last_edit_prediction_refresh: None, last_jump_prediction_refresh: None, @@ -1326,14 +1402,31 @@ impl EditPredictionStore { }; let path = project.read(cx).path_for_entry(*active_entry_id, cx); if let Some(path) = path { - if let Some(ix) = project_state - .recent_paths - .iter() - .position(|probe| probe == &path) - { - project_state.recent_paths.remove(ix); + let cursor_position = project + .read(cx) + .buffer_store() + .read(cx) + .get_by_path(&path) + .and_then(|buffer| { + let position = project_state + .registered_buffers + .get(&buffer.entity_id())? + .last_position?; + Some(position.to_offset(&buffer.read(cx).snapshot())) + }); + + let recent_file = RecentFile { + path: path.path.as_std_path().into(), + cursor_position, + }; + for capture in &mut project_state.pending_jump_example_captures { + capture.navigation_history.push(recent_file.clone()); + if capture.navigation_history.len() > JUMP_EXAMPLE_NAVIGATION_COUNT { + capture.navigation_history.remove(0); + } } - project_state.recent_paths.push_front(path); + push_recent_file(&mut project_state.recently_viewed_files, recent_file); + jump_example::drain_completed_jump_example_captures(project_state, cx); } } project::Event::DiagnosticsUpdated { .. } => { @@ -1492,17 +1585,8 @@ impl EditPredictionStore { compute_diff_between_snapshots_in_range(&old_snapshot, &new_snapshot, &edit_range) .is_some(); - let events = &mut project_state.events; - if !is_recordable_history_edit { - if let Some(event) = project_state.last_event.take() { - if let Some(event) = event.finalize(&project_state.license_detection_watchers, cx) { - if events.len() + 1 >= EVENT_COUNT_MAX { - events.pop_front(); - } - events.push_back(event); - } - } + project_state.finalize_last_event(cx); return; } @@ -1541,16 +1625,21 @@ impl EditPredictionStore { } } - if let Some(event) = project_state.last_event.take() { - if let Some(event) = event.finalize(&project_state.license_detection_watchers, cx) { - if events.len() + 1 >= EVENT_COUNT_MAX { - events.pop_front(); - } - events.push_back(event); - } - } + project_state.finalize_last_event(cx); - merge_trailing_events_if_needed(events, &old_snapshot, &new_snapshot, &edit_range); + merge_trailing_events_if_needed( + &mut project_state.events, + &old_snapshot, + &new_snapshot, + &edit_range, + ); + + let file_context = new_file.as_ref().map(|file| { + let project_path = ProjectPath::from_file(file.as_ref(), cx); + let file_context = project_state.file_context_for_path(project_path.clone(), cx); + Self::ensure_git_changed_file_sets_loading(&file_context, project, &project_path, cx); + file_context + }); project_state.last_event = Some(LastEvent { old_file, @@ -1563,6 +1652,7 @@ impl EditPredictionStore { predicted: is_predicted, snapshot_after_last_editing_pause: None, last_edit_time: Some(now), + file_context, }); } @@ -1574,12 +1664,18 @@ impl EditPredictionStore { cx: &App, ) -> Option> { let project_state = self.projects.get_mut(&project.entity_id())?; - if let Some(position) = position - && let Some(buffer) = project_state + if let Some(position) = position { + let snapshot = buffer.read(cx).snapshot(); + let cursor_position = position.to_offset(&snapshot); + if let Some(file) = snapshot.file() { + project_state.update_recent_file_cursor(file.path().as_std_path(), cursor_position); + } + if let Some(buffer) = project_state .registered_buffers .get_mut(&buffer.entity_id()) - { - buffer.last_position = Some(position); + { + buffer.last_position = Some(position); + } } let CurrentEditPrediction { @@ -2610,6 +2706,7 @@ impl EditPredictionStore { snapshot, position, events, + stored_events: stored_events.clone(), related_files, mode, trigger, @@ -2621,22 +2718,40 @@ impl EditPredictionStore { let task = match self.edit_prediction_model { EditPredictionModel::Zeta => { - 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, - ), - ) - }); + let capture_data = if let Some(worktree_id) = capture_worktree_id + && can_collect_data + { + let uncommitted_diff_snapshot = uncommitted_diffs_for_events( + project.clone(), + worktree_id, + stored_events.clone(), + cx, + ) + .shared(); + jump_example::try_start_jump_example_capture( + project_state, + uncommitted_diff_snapshot.clone(), + inputs.project.clone(), + inputs.snapshot.clone(), + inputs.position, + match trigger { + PredictEditsRequestTrigger::Diagnostics => { + JumpExampleTrigger::Diagnostic + } + _ => JumpExampleTrigger::Prediction, + }, + stored_events, + inputs.diagnostic_search_range.clone(), + can_collect_data, + is_open_source, + cx, + ); + rand::random_ratio(1, 10).then(|| uncommitted_diff_snapshot) + } else { + None + }; - zeta::request_prediction_with_zeta(self, inputs, capture_events, repo_url, cx) + zeta::request_prediction_with_zeta(self, inputs, capture_data, repo_url, cx) } EditPredictionModel::Fim { format } => fim::request_prediction(inputs, format, cx), EditPredictionModel::Mercury => { @@ -2954,43 +3069,39 @@ impl EditPredictionStore { cx: &mut Context, ) { let project_state = self.get_or_init_project(project, cx); - 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) + project_state.recently_viewed_files = paths .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, - }) + .map(|path| RecentFile { + path: path.path.as_std_path().into(), + cursor_position: None, }) .collect(); - let recently_viewed_files = self - .projects + } + + pub fn recently_opened_files_for_project(&self, project: &Entity) -> Vec { + self.projects .get(&project.entity_id()) .map(|project_state| { project_state - .recent_paths + .recently_opened_files .iter() - .map(|path| RecentFile { - path: path.path.as_std_path().into(), - cursor_position: None, - }) + .cloned() .collect() }) - .unwrap_or_default(); + .unwrap_or_default() + } - (recently_opened_files, recently_viewed_files) + pub fn recently_viewed_files_for_project(&self, project: &Entity) -> Vec { + self.projects + .get(&project.entity_id()) + .map(|project_state| { + project_state + .recently_viewed_files + .iter() + .cloned() + .collect() + }) + .unwrap_or_default() } fn is_file_open_source( @@ -3171,8 +3282,8 @@ fn merge_trailing_events_if_needed( return; } - let mut events_to_merge = events.range(events.len() - mergeable_count..).peekable(); - let oldest_event = events_to_merge.peek().unwrap(); + let merge_start = events.len() - mergeable_count; + let oldest_event = &events[merge_start]; let oldest_snapshot = oldest_event.old_snapshot.clone(); let newest_snapshot = end_snapshot; let mut merged_edit_range = oldest_event.total_edit_range.clone(); @@ -3199,9 +3310,9 @@ fn merge_trailing_events_if_needed( path: path.clone(), diff, in_open_source_repo: *in_open_source_repo, - predicted: events_to_merge.all(|e| { + predicted: events.range(merge_start..).all(|event| { matches!( - e.event.as_ref(), + event.event.as_ref(), zeta_prompt::Event::BufferChange { predicted: true, .. @@ -3213,7 +3324,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, + file_context: oldest_event.file_context.clone(), }, }; events.truncate(events.len() - mergeable_count); diff --git a/crates/edit_prediction/src/example_spec.rs b/crates/edit_prediction/src/example_spec.rs index d3f4bbd3e8..c610b777dc 100644 --- a/crates/edit_prediction/src/example_spec.rs +++ b/crates/edit_prediction/src/example_spec.rs @@ -6,6 +6,8 @@ use telemetry_events::EditPredictionRating; pub use zeta_prompt::udiff::{ CURSOR_POSITION_MARKER, encode_cursor_in_patch, extract_cursor_from_patch, }; + +use crate::data_collection::format_cursor_excerpt; pub const INLINE_CURSOR_MARKER: &str = "<|user_cursor|>"; /// Maximum cursor file size to capture (64KB). @@ -482,51 +484,7 @@ impl ExampleSpec { cursor_offset: usize, line_comment_prefix: &str, ) { - // Find which line the cursor is on and its column - let cursor_line_start = excerpt[..cursor_offset] - .rfind('\n') - .map(|pos| pos + 1) - .unwrap_or(0); - let cursor_line_end = excerpt[cursor_line_start..] - .find('\n') - .map(|pos| cursor_line_start + pos + 1) - .unwrap_or(excerpt.len()); - let cursor_line = &excerpt[cursor_line_start..cursor_line_end]; - let cursor_line_indent = &cursor_line[..cursor_line.len() - cursor_line.trim_start().len()]; - let cursor_column = cursor_offset - cursor_line_start; - - // Build the marker line - let mut marker_line = String::new(); - if cursor_column < line_comment_prefix.len() { - for _ in 0..cursor_column { - marker_line.push(' '); - } - marker_line.push_str(line_comment_prefix); - write!(marker_line, " <{}", CURSOR_POSITION_MARKER).unwrap(); - } else { - if cursor_column >= cursor_line_indent.len() + line_comment_prefix.len() { - marker_line.push_str(cursor_line_indent); - } - marker_line.push_str(line_comment_prefix); - while marker_line.len() < cursor_column { - marker_line.push(' '); - } - write!(marker_line, "^{}", CURSOR_POSITION_MARKER).unwrap(); - } - - // Build the final cursor_position string - let mut result = String::with_capacity(excerpt.len() + marker_line.len() + 2); - result.push_str(&excerpt[..cursor_line_end]); - if !result.ends_with('\n') { - result.push('\n'); - } - result.push_str(&marker_line); - if cursor_line_end < excerpt.len() { - result.push('\n'); - result.push_str(&excerpt[cursor_line_end..]); - } - - self.cursor_position = result; + self.cursor_position = format_cursor_excerpt(excerpt, cursor_offset, line_comment_prefix); } /// Returns all of the possible expected patches for this example, each with an optional diff --git a/crates/edit_prediction/src/jump_example.rs b/crates/edit_prediction/src/jump_example.rs new file mode 100644 index 0000000000..33991df9c3 --- /dev/null +++ b/crates/edit_prediction/src/jump_example.rs @@ -0,0 +1,372 @@ +use std::{ + ops::Range, + sync::Arc, + time::{Duration, Instant}, +}; + +use anyhow::{Context as _, Result}; +pub use cloud_api_types::JumpExampleTrigger; +use cloud_api_types::{ + JumpExampleRecentFile, SubmitEditPredictionJumpExampleBody, + SubmitEditPredictionJumpExampleResponse, +}; +use futures::future::Shared; +use gpui::{AppContext as _, AsyncApp, Context, Entity, Task, TaskExt as _, WeakEntity}; +use language::{BufferSnapshot, File, Point}; +use project::{Project, WorktreeId}; +use release_channel::AppVersion; + +use text::ToPoint as _; +use util::rel_path::RelPath; + +use crate::{ + EditPredictionStore, ProjectState, StoredEvent, + data_collection::{ + UncommittedDiffResult, compute_cursor_excerpt, compute_uncommitted_diff, + estimate_uncommitted_diff_byte_size, format_cursor_excerpt, + }, + example_spec::RecentFile, + zeta, +}; + +pub const JUMP_EXAMPLE_MAX_PENDING_CAPTURE_COUNT: usize = 10; +pub const JUMP_EXAMPLE_FUTURE_EVENT_COUNT: usize = 2; +pub const JUMP_EXAMPLE_TTL: Duration = Duration::from_secs(60 * 2); +pub const JUMP_EXAMPLE_NAVIGATION_COUNT: usize = 20; +pub const JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024; + +pub struct PendingJumpExampleCapture { + key: PendingJumpExampleCaptureKey, + trigger: JumpExampleTrigger, + file: Arc, + edit_history: Vec>, + recently_opened_files: Vec, + recently_viewed_files: Vec, + worktree_root_name: String, + cursor_position: String, + started_at: Instant, + uncommitted_diff: Option, + pub future_events: Vec>, + pub navigation_history: Vec, + diagnostics: Vec, + repository_url: Option, + revision: Option, + can_collect_data: bool, + is_in_open_source_repo: bool, +} + +#[derive(Eq, PartialEq, Hash, Clone)] +pub struct PendingJumpExampleCaptureKey { + worktree_id: WorktreeId, + file_path: Arc, + row_bucket: u32, +} + +pub fn try_start_jump_example_capture( + project_state: &ProjectState, + uncommitted_diffs: Shared>, + project: Entity, + snapshot: BufferSnapshot, + position: language::Anchor, + trigger: JumpExampleTrigger, + stored_events: Vec, + diagnostic_search_range: Range, + can_collect_data: bool, + is_in_open_source_repo: bool, + cx: &mut Context, +) { + let Some(file) = snapshot.file().cloned() else { + return; + }; + + let example_key = PendingJumpExampleCaptureKey { + worktree_id: file.worktree_id(cx), + file_path: file.path().clone(), + row_bucket: position.to_point(&snapshot).row / 10, + }; + let should_capture_example = project_state.pending_jump_example_captures.len() + < JUMP_EXAMPLE_MAX_PENDING_CAPTURE_COUNT + && !project_state + .starting_jump_example_captures + .contains(&example_key) + && !project_state + .pending_jump_example_captures + .iter() + .any(|capture| &capture.key == &example_key); + + if !should_capture_example { + return; + } + + let _project = project.clone(); + let _example_key = example_key.clone(); + let task = cx.spawn(async move |ep_store, cx| { + let project = _project; + let example_key = _example_key; + let Some(ep_store) = ep_store.upgrade() else { + return anyhow::Ok(()); + }; + ep_store.update(cx, |ep_store, cx| { + let project_state = ep_store.get_or_init_project(&project, cx); + project_state + .starting_jump_example_captures + .push(example_key.clone()); + }); + + let (repository, worktree) = project.read_with(cx, |project, cx| { + let repository = project.active_repository(cx); + let worktree_id = file.worktree_id(cx); + let worktree = project.worktree_for_id(worktree_id, cx); + (repository, worktree) + }); + let Some(worktree) = worktree else { + return Ok(()); + }; + + let diagnostics = zeta::active_buffer_diagnostics( + &snapshot, + diagnostic_search_range.clone(), + position.to_point(&snapshot).row, + 100, + ); + + let uncommitted_diff = 'uncommitted_diff: { + if repository.is_none() { + break 'uncommitted_diff None; + } + let uncommitted_diff_snapshot = uncommitted_diffs + .await + .map_err(|error| anyhow::anyhow!("{error:?}")) + .context("failed to capture uncommitted diff")?; + let estimated_byte_size = + estimate_uncommitted_diff_byte_size(&uncommitted_diff_snapshot); + if estimated_byte_size > JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE { + break 'uncommitted_diff None; + } + + let uncommitted_diff = cx + .background_executor() + .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) + .await; + if uncommitted_diff.len() > JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE { + break 'uncommitted_diff None; + } + Some(uncommitted_diff) + }; + + let edit_history = stored_events + .iter() + .map(|e| e.event.clone()) + .collect::>(); + let (repository_url, revision) = if let Some(repository) = &repository { + repository.read_with(cx, |repository, _| { + let snapshot = repository.snapshot(); + ( + snapshot + .remote_origin_url + .clone() + .or_else(|| snapshot.remote_upstream_url.clone()), + snapshot + .head_commit + .as_ref() + .map(|commit| commit.sha.to_string()), + ) + }) + } else { + (None, None) + }; + let line_comment_prefix = snapshot + .language() + .and_then(|language| language.config().line_comments.first()) + .map(|prefix| prefix.to_string()) + .unwrap_or_default(); + let (cursor_excerpt, cursor_offset_in_excerpt, _) = cx + .background_executor() + .spawn(async move { compute_cursor_excerpt(&snapshot, position) }) + .await; + let cursor_position = format_cursor_excerpt( + &cursor_excerpt, + cursor_offset_in_excerpt, + &line_comment_prefix, + ); + let now = cx.background_executor().now(); + ep_store.update(cx, |ep_store, cx| { + let recently_opened_files = ep_store.recently_opened_files_for_project(&project); + let recently_viewed_files = ep_store.recently_viewed_files_for_project(&project); + let project_state = ep_store.get_or_init_project(&project, cx); + project_state + .pending_jump_example_captures + .push(PendingJumpExampleCapture { + key: example_key, + trigger, + file, + uncommitted_diff, + edit_history, + recently_opened_files, + recently_viewed_files, + repository_url, + revision, + diagnostics, + worktree_root_name: worktree.read(cx).root_name_str().to_owned(), + cursor_position, + started_at: now, + future_events: Vec::new(), + navigation_history: Vec::new(), + is_in_open_source_repo, + can_collect_data, + }); + drain_completed_jump_example_captures(project_state, cx); + }); + Ok(()) + }); + cx.spawn(async move |ep_store, cx| { + let result = task.await; + ep_store + .update(cx, |ep_store, cx| { + ep_store + .get_or_init_project(&project, cx) + .starting_jump_example_captures + .retain(|key| key != &example_key); + }) + .ok(); + result + }) + .detach_and_log_err(cx); +} + +pub fn drain_completed_jump_example_captures( + project_state: &mut ProjectState, + cx: &mut Context, +) { + let now = cx.background_executor().now(); + + let mut capture_index = 0; + while capture_index < project_state.pending_jump_example_captures.len() { + let capture = &project_state.pending_jump_example_captures[capture_index]; + let finished = capture.future_events.len() >= JUMP_EXAMPLE_FUTURE_EVENT_COUNT + || now.saturating_duration_since(capture.started_at) >= JUMP_EXAMPLE_TTL; + if !finished { + capture_index += 1; + continue; + } + + let capture = project_state + .pending_jump_example_captures + .remove(capture_index); + cx.spawn(async move |this, cx| { + let result = submit_jump_example_capture_task(this, capture, cx).await; + if let Err(error) = result { + log::error!("failed to submit jump opportunity capture: {error:?}"); + } + }) + .detach(); + } +} + +fn submit_jump_example_capture_task( + this: WeakEntity, + capture: PendingJumpExampleCapture, + cx: &mut AsyncApp, +) -> Task> { + let Some((organization_id, client, llm_token, app_version)) = this + .update(cx, |this, cx| { + ( + this.user_store + .read(cx) + .current_organization() + .map(|organization| organization.id.clone()), + this.client.clone(), + this.llm_token.clone(), + AppVersion::global(cx), + ) + }) + .ok() + else { + return Task::ready(Ok(())); + }; + cx.background_spawn(async move { + let PendingJumpExampleCapture { + key: _, + trigger, + file, + edit_history, + recently_opened_files, + recently_viewed_files, + worktree_root_name, + cursor_position, + started_at: _, + uncommitted_diff, + future_events, + navigation_history, + diagnostics, + repository_url, + revision, + is_in_open_source_repo, + can_collect_data, + } = capture; + let future_edit_history = render_jump_example_events(&future_events, &worktree_root_name); + + let cursor_path = file.path().as_std_path().into(); + let example = SubmitEditPredictionJumpExampleBody { + request_id: uuid::Uuid::new_v4(), + trigger, + repository_url, + revision, + uncommitted_diff, + recently_opened_files: jump_example_recent_files(recently_opened_files), + recently_viewed_files: jump_example_recent_files(recently_viewed_files), + cursor_path, + cursor_position, + edit_history, + diagnostics, + future_edit_history, + navigation_history: jump_example_recent_files(navigation_history), + is_in_open_source_repo, + can_collect_data, + }; + let json_bytes = serde_json::to_vec(&example)?; + let compressed = zstd::encode_all(&json_bytes[..], 3)?; + let url = client + .http_client() + .build_zed_llm_url("/predict_edits/jump_example", &[])?; + EditPredictionStore::send_api_request::( + |builder| { + Ok(builder + .uri(url.as_ref()) + .header("Content-Encoding", "zstd") + .body(compressed.clone().into())?) + }, + client, + llm_token, + organization_id, + app_version, + ) + .await?; + Ok(()) + }) +} + +fn jump_example_recent_files(files: Vec) -> Vec { + files + .into_iter() + .map(|file| JumpExampleRecentFile { + path: file.path, + cursor_position: file.cursor_position, + }) + .collect() +} + +fn render_jump_example_events(events: &[Arc], root_name: &str) -> String { + let mut edit_history = String::new(); + for event in events { + crate::capture_example::write_event_with_relative_paths( + &mut edit_history, + event, + root_name, + ); + if !edit_history.ends_with('\n') { + edit_history.push('\n'); + } + } + edit_history +} diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 28aa122951..ee2bcd62f0 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -3,14 +3,15 @@ use crate::{ EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, ZedUpdateRequiredError, buffer_path_with_id_fallback, cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, + data_collection::UncommittedDiffResult, prediction::EditPredictionResult, }; -use anyhow::Result; -use buffer_diff::BufferDiff; +use anyhow::{Context as _, Result}; use cloud_llm_client::{ AcceptEditPredictionBody, EditPredictionRejectReason, predict_edits_v3::RawCompletionRequest, }; use edit_prediction_types::PredictedCursorPosition; +use futures::future::Shared; use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*}; use language::{ Buffer, BufferSnapshot, DiagnosticSeverity, EditPredictionPromptFormat, OffsetRangeExt as _, @@ -40,6 +41,7 @@ pub fn request_prediction_with_zeta( position, related_files, events, + stored_events, debug_tx, mode, trigger, @@ -49,10 +51,7 @@ pub fn request_prediction_with_zeta( is_open_source, .. }: EditPredictionModelInput, - capture_data: Option<( - Vec, - Task, Entity>>>, - )>, + capture_data: Option>>, repo_url: Option, cx: &mut Context, ) -> Task>> { @@ -407,21 +406,29 @@ 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(|(events, uncommitted_diffs)| { + let example_task = capture_data.and_then(|uncommitted_diffs| { let (recently_opened_files, recently_viewed_files) = this - .read_with(cx, |this, cx| this.recent_paths_for_project(&project, cx)) + .read_with(cx, |this, _| { + ( + this.recently_opened_files_for_project(&project), + this.recently_viewed_files_for_project(&project), + ) + }) .ok()?; Some(cx.spawn({ let project = project.clone(); let edited_buffer = edited_buffer.clone(); async move |cx| { - let uncommitted_diffs = uncommitted_diffs.await?; + let uncommitted_diffs = uncommitted_diffs + .await + .map_err(|error| anyhow::anyhow!("{error:?}")) + .context("failed to capture uncommitted diff")?; let Some(task) = cx.update(|cx| { crate::capture_example::capture_example( project.clone(), edited_buffer.clone(), position, - events, + stored_events, recently_opened_files, recently_viewed_files, uncommitted_diffs, diff --git a/crates/edit_prediction_cli/src/synthesize.rs b/crates/edit_prediction_cli/src/synthesize.rs index 76c8e0fee3..71d691bf9f 100644 --- a/crates/edit_prediction_cli/src/synthesize.rs +++ b/crates/edit_prediction_cli/src/synthesize.rs @@ -9,6 +9,7 @@ use anyhow::{Context as _, Result}; use chrono::Local; use collections::{HashMap, HashSet}; use edit_prediction::{ + data_collection::format_cursor_excerpt, example_spec::ExampleSpec, udiff::{apply_diff_to_string, edits_for_diff}, }; @@ -783,7 +784,7 @@ async fn build_example( truncate_message(&commit.message, 60), response.reasoning ); - let mut spec = ExampleSpec { + let spec = ExampleSpec { name: response.name.clone(), repository_url: repo_url.to_string(), revision: commit.parent_sha.clone(), @@ -794,16 +795,14 @@ async fn build_example( recently_viewed_files: Vec::new(), uncommitted_diff_contains_edit_history: false, cursor_path: Arc::from(Path::new(&cursor_file)), - cursor_position: String::new(), + cursor_position: format_cursor_excerpt(&excerpt, cursor_offset, comment_prefix), edit_history, expected_patches: vec![expected_patch_with_header], rejected_patch: None, - telemetry: None, human_feedback: Vec::new(), rating: None, }; - spec.set_cursor_excerpt(&excerpt, cursor_offset, comment_prefix); Ok(spec) } diff --git a/crates/fs/src/fake_git_repo.rs b/crates/fs/src/fake_git_repo.rs index 256ca70db3..7635317aed 100644 --- a/crates/fs/src/fake_git_repo.rs +++ b/crates/fs/src/fake_git_repo.rs @@ -12,9 +12,9 @@ use git::{ blame::Blame, repository::{ AskPassDelegate, Branch, CommitData, CommitDataReader, CommitDetails, CommitOptions, - CreateWorktreeTarget, FetchOptions, GRAPH_CHUNK_SIZE, GitRepository, - GitRepositoryCheckpoint, InitialGraphCommitData, LogOrder, LogSource, PushOptions, RefEdit, - Remote, RepoPath, ResetMode, SearchCommitArgs, Worktree, + CreateWorktreeTarget, FetchOptions, FileHistoryChangedFileSets, GRAPH_CHUNK_SIZE, + GitRepository, GitRepositoryCheckpoint, InitialGraphCommitData, LogOrder, LogSource, + PushOptions, RefEdit, Remote, RepoPath, ResetMode, SearchCommitArgs, Worktree, }, stash::GitStash, status::{ @@ -1504,6 +1504,14 @@ impl GitRepository for FakeGitRepository { .boxed() } + fn file_history_changed_files( + &self, + paths: Vec, + _commit_limit: usize, + ) -> BoxFuture<'_, Result>> { + async move { Ok(vec![FileHistoryChangedFileSets::default(); paths.len()]) }.boxed() + } + fn commit_data_reader(&self) -> Result { let fs = self.fs.clone(); let dot_git_path = self.dot_git_path.clone(); diff --git a/crates/git/src/repository.rs b/crates/git/src/repository.rs index 6cfa48282b..7da5ae4490 100644 --- a/crates/git/src/repository.rs +++ b/crates/git/src/repository.rs @@ -512,6 +512,11 @@ pub struct CommitDiff { pub files: Vec, } +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct FileHistoryChangedFileSets { + pub file_sets: Vec>, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum CommitFileStatus { Added, @@ -1047,6 +1052,12 @@ pub trait GitRepository: Send + Sync { request_tx: Sender, ) -> BoxFuture<'_, Result<()>>; + fn file_history_changed_files( + &self, + paths: Vec, + commit_limit: usize, + ) -> BoxFuture<'_, Result>>; + fn commit_data_reader(&self) -> Result; fn update_ref(&self, ref_name: String, commit: String) -> BoxFuture<'_, Result<()>>; @@ -3100,6 +3111,50 @@ impl GitRepository for RealGitRepository { .boxed() } + fn file_history_changed_files( + &self, + paths: Vec, + commit_limit: usize, + ) -> BoxFuture<'_, Result>> { + let git = self.git_binary(); + + async move { + if paths.is_empty() { + return Ok(Vec::new()); + } + + if commit_limit == 0 { + return Ok(vec![FileHistoryChangedFileSets::default(); paths.len()]); + } + + let max_count_arg = format!("--max-count={commit_limit}"); + let mut args = [ + "log", + max_count_arg.as_str(), + "--full-diff", + "--no-renames", + "--name-only", + "-z", + "--format=%x1e", + "--", + ] + .map(OsString::from) + .to_vec(); + args.extend(paths.iter().map(|path| OsString::from(path.as_unix_str()))); + + let output = git.build_command(&args).output().await?; + anyhow::ensure!( + output.status.success(), + "git log failed:\n{}", + String::from_utf8_lossy(&output.stderr) + ); + + let stdout = String::from_utf8_lossy(&output.stdout); + Ok(parse_file_history_changed_files_output(&stdout, &paths)) + } + .boxed() + } + fn commit_data_reader(&self) -> Result { let git_binary = self.git_binary(); @@ -3209,6 +3264,39 @@ async fn read_single_commit_response( .ok_or_else(|| anyhow!("failed to parse commit {}", sha)) } +fn parse_file_history_changed_files_output( + output: &str, + queried_paths: &[RepoPath], +) -> Vec { + let mut histories = vec![FileHistoryChangedFileSets::default(); queried_paths.len()]; + + for record in output.split('\x1e') { + let changed_files = record + .split('\0') + .filter_map(|field| { + let path = field.trim_start_matches('\n'); + if path.is_empty() { + return None; + } + RepoPath::new(path).ok() + }) + .collect::>(); + + if changed_files.is_empty() { + continue; + } + + let file_set = changed_files.iter().cloned().collect::>(); + for (index, queried_path) in queried_paths.iter().enumerate() { + if changed_files.contains(queried_path) { + histories[index].file_sets.push(file_set.clone()); + } + } + } + + histories +} + fn parse_initial_graph_output<'a>( lines: impl Iterator, ) -> Vec> { @@ -3742,6 +3830,51 @@ mod tests { assert_eq!(commit.tag_names(), ["v1.0.0", "v1.1.0"]); } + #[test] + fn test_parse_file_history_changed_files_output() { + let queried_paths = vec![ + RepoPath::new("src/a.rs").unwrap(), + RepoPath::new("src/b.rs").unwrap(), + ]; + let output = concat!( + "\x1e\0\nsrc/a.rs\0src/shared.rs\0", + "\x1e\0\nsrc/b.rs\0src/shared.rs\0", + "\x1e\0\nsrc/a.rs\0src/b.rs\0src/shared.rs\0", + ); + + let histories = parse_file_history_changed_files_output(output, &queried_paths); + + assert_eq!(histories.len(), 2); + assert_eq!( + histories[0].file_sets, + vec![ + vec![ + RepoPath::new("src/a.rs").unwrap(), + RepoPath::new("src/shared.rs").unwrap(), + ], + vec![ + RepoPath::new("src/a.rs").unwrap(), + RepoPath::new("src/b.rs").unwrap(), + RepoPath::new("src/shared.rs").unwrap(), + ], + ] + ); + assert_eq!( + histories[1].file_sets, + vec![ + vec![ + RepoPath::new("src/b.rs").unwrap(), + RepoPath::new("src/shared.rs").unwrap(), + ], + vec![ + RepoPath::new("src/a.rs").unwrap(), + RepoPath::new("src/b.rs").unwrap(), + RepoPath::new("src/shared.rs").unwrap(), + ], + ] + ); + } + #[gpui::test] async fn test_build_command_untrusted_includes_both_safety_args(cx: &mut TestAppContext) { cx.executor().allow_parking(); diff --git a/crates/project/src/git_store.rs b/crates/project/src/git_store.rs index 5df59cbd18..64353d0fc3 100644 --- a/crates/project/src/git_store.rs +++ b/crates/project/src/git_store.rs @@ -35,10 +35,10 @@ use git::{ parse_git_remote_url, repository::{ Branch, BranchesScanResult, CommitData, CommitDetails, CommitDiff, CommitFile, - CommitOptions, CreateWorktreeTarget, DiffType, FetchOptions, GitCommitTemplate, - GitRepository, GitRepositoryCheckpoint, InitialGraphCommitData, LogOrder, LogSource, - PushOptions, Remote, RemoteCommandOutput, RepoPath, ResetMode, SearchCommitArgs, - UpstreamTrackingStatus, Worktree as GitWorktree, delete_branch_flag, + CommitOptions, CreateWorktreeTarget, DiffType, FetchOptions, FileHistoryChangedFileSets, + GitCommitTemplate, GitRepository, GitRepositoryCheckpoint, InitialGraphCommitData, + LogOrder, LogSource, PushOptions, Remote, RemoteCommandOutput, RepoPath, ResetMode, + SearchCommitArgs, UpstreamTrackingStatus, Worktree as GitWorktree, delete_branch_flag, }, stash::{GitStash, StashEntry}, status::{ @@ -5271,6 +5271,29 @@ impl Repository { }) } + pub fn file_history_changed_files( + &mut self, + paths: Vec, + commit_limit: usize, + ) -> oneshot::Receiver>> { + self.send_job( + "file_history_changed_files", + None, + move |git_repo, _cx| async move { + match git_repo { + RepositoryState::Local(LocalRepositoryState { backend, .. }) => { + backend + .file_history_changed_files(paths, commit_limit) + .await + } + RepositoryState::Remote(_) => { + anyhow::bail!("file history changed files is only supported locally") + } + } + }, + ) + } + pub fn get_graph_data( &self, log_source: LogSource,