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

ep: Jump example capture (#58236)

Browse Source 
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 ...
This commit is contained in:
Ben Kunkle authored and GitHub committed 2026-06-01 20:18:48 +00:00
1 parent 64ab35c4b6
commit e36459aead
14 files changed
+1219 -436
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+3
View File
@@ -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",
crates/cloud_api_types/Cargo.toml
+2
View File
@@ -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
crates/cloud_api_types/src/cloud_api_types.rs
+37
View File
@@ -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<String>,
pub revision: Option<String>,
pub uncommitted_diff: Option<String>,
pub recently_opened_files: Vec<JumpExampleRecentFile>,
pub recently_viewed_files: Vec<JumpExampleRecentFile>,
pub cursor_path: Arc<Path>,
pub cursor_position: String,
pub edit_history: Vec<Arc<zeta_prompt::Event>>,
pub diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
pub future_edit_history: String,
pub navigation_history: Vec<JumpExampleRecentFile>,
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<Path>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cursor_position: Option<usize>,
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct EditPredictionSettledKeptChars {
#[serde(rename = "edit_bytes_candidate_new")]
crates/edit_prediction/Cargo.toml
+1
View File
@@ -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
crates/edit_prediction/src/capture_example.rs
+42 -203
View File
@@ -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<Project>,
@@ -18,7 +19,7 @@ pub fn capture_example(
events: Vec<StoredEvent>,
recently_opened_files: Vec<RecentFile>,
recently_viewed_files: Vec<RecentFile>,
uncommitted_diffs_by_path: HashMap<Arc<Path>, Entity<BufferDiff>>,
uncommitted_diff_snapshot: UncommittedDiffSnapshot,
populate_expected_patch: bool,
cx: &mut App,
) -> Option<Task<Result<ExampleSpec>>> {
@@ -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<Arc<Path>, (Entity<Buffer>, Entity<BufferDiff>)> =
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<Path> = 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<Point>) {
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<Arc<Path>, (language::BufferSnapshot, BufferDiffSnapshot)>,
) -> String {
let mut uncommitted_diff = String::new();
let mut snapshots_by_path = snapshots_by_path.into_iter().collect::<Vec<_>>();
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<u32>, Range<u32>)> = 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::<String>(),
&buffer_snapshot
.text_for_range(
Point::new(new_range.start, 0)
..row_start_or_max(&buffer_snapshot, new_range.end),
)
.collect::<String>(),
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::<Vec<_>>();
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,
crates/edit_prediction/src/data_collection.rs
+290
View File
@@ -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<Path>, BufferSnapshot, BufferDiffSnapshot)>;
pub type UncommittedDiffResult = std::result::Result<UncommittedDiffSnapshot, Arc<anyhow::Error>>;
pub use zeta_prompt::udiff::CURSOR_POSITION_MARKER;
pub fn uncommitted_diffs_for_events(
project: Entity<Project>,
worktree_id: WorktreeId,
events: Vec<StoredEvent>,
cx: &Context<'_, EditPredictionStore>,
) -> Task<UncommittedDiffResult> {
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<Path> = 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::<Result<Vec<_>>>()
.map_err(Arc::new)?;
let mut snapshots_by_path: HashMap<Arc<Path>, (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<Point>) {
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<u32>, Range<u32>)> = 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::<String>(),
&buffer_snapshot
.text_for_range(
Point::new(new_range.start, 0)
..row_start_or_max(&buffer_snapshot, new_range.end),
)
.collect::<String>(),
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
}
crates/edit_prediction/src/edit_prediction.rs
+278 -167
View File
@@ -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<Arc<zeta_prompt::Event>>,
stored_events: Vec<StoredEvent>,
related_files: Vec<RelatedFile>,
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<Anchor>,
pub uncommitted_diff: Option<Entity<BufferDiff>>,
pub(crate) file_context: Option<Entity<StoredFileContext>>,
}
pub(crate) struct StoredFileContext {
pub(crate) uncommitted_diff: Option<Entity<BufferDiff>>,
pub(crate) git_changed_file_sets: Option<Arc<FileHistoryChangedFileSets>>,
pub(crate) git_changed_file_sets_task: Option<Task<()>>,
}
impl StoredEvent {
@@ -318,15 +335,30 @@ fn lines_between_ranges(left: &Range<Point>, right: &Range<Point>) -> u32 {
0
}
fn push_recent_file(files: &mut VecDeque<RecentFile>, 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<StoredEvent>,
last_event: Option<LastEvent>,
recent_paths: VecDeque<ProjectPath>,
recently_viewed_files: VecDeque<RecentFile>,
recently_opened_files: VecDeque<RecentFile>,
registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
file_contexts: HashMap<ProjectPath, WeakEntity<StoredFileContext>>,
current_prediction: Option<CurrentEditPrediction>,
last_edit_source: Option<BufferEditSource>,
next_pending_prediction_id: usize,
pending_predictions: ArrayVec<PendingPrediction, 2, u8>,
pending_jump_example_captures: Vec<PendingJumpExampleCapture>,
starting_jump_example_captures: Vec<PendingJumpExampleCaptureKey>,
debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,
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<EditPredictionStore>,
) -> Entity<StoredFileContext> {
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<EditPredictionStore>) {
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<TextBufferSnapshot>,
last_edit_time: Option<Instant>,
file_context: Option<Entity<StoredFileContext>>,
}
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<Project>,
worktree_id: WorktreeId,
events: Vec<StoredEvent>,
cx: &mut Context<Self>,
) -> Task<Result<HashMap<Arc<Path>, Entity<BufferDiff>>>> {
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<Path> = 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<Project>,
@@ -1186,10 +1191,77 @@ impl EditPredictionStore {
project: &Entity<Project>,
cx: &mut Context<Self>,
) {
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<StoredFileContext>,
project: &Entity<Project>,
project_path: &ProjectPath,
cx: &mut Context<Self>,
) {
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<Project>,
@@ -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<BufferEditPrediction<'_>> {
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<Self>,
) {
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<Project>,
cx: &App,
) -> (Vec<RecentFile>, Vec<RecentFile>) {
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<Project>) -> Vec<RecentFile> {
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<Project>) -> Vec<RecentFile> {
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);
crates/edit_prediction/src/example_spec.rs
+3 -45
View File
@@ -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
crates/edit_prediction/src/jump_example.rs
+372
View File
@@ -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<dyn File>,
edit_history: Vec<Arc<zeta_prompt::Event>>,
recently_opened_files: Vec<RecentFile>,
recently_viewed_files: Vec<RecentFile>,
worktree_root_name: String,
cursor_position: String,
started_at: Instant,
uncommitted_diff: Option<String>,
pub future_events: Vec<Arc<zeta_prompt::Event>>,
pub navigation_history: Vec<RecentFile>,
diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
repository_url: Option<String>,
revision: Option<String>,
can_collect_data: bool,
is_in_open_source_repo: bool,
}
#[derive(Eq, PartialEq, Hash, Clone)]
pub struct PendingJumpExampleCaptureKey {
worktree_id: WorktreeId,
file_path: Arc<RelPath>,
row_bucket: u32,
}
pub fn try_start_jump_example_capture(
project_state: &ProjectState,
uncommitted_diffs: Shared<Task<UncommittedDiffResult>>,
project: Entity<Project>,
snapshot: BufferSnapshot,
position: language::Anchor,
trigger: JumpExampleTrigger,
stored_events: Vec<StoredEvent>,
diagnostic_search_range: Range<Point>,
can_collect_data: bool,
is_in_open_source_repo: bool,
cx: &mut Context<EditPredictionStore>,
) {
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::<Vec<_>>();
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<EditPredictionStore>,
) {
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<EditPredictionStore>,
capture: PendingJumpExampleCapture,
cx: &mut AsyncApp,
) -> Task<Result<()>> {
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::<SubmitEditPredictionJumpExampleResponse>(
|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<RecentFile>) -> Vec<JumpExampleRecentFile> {
files
.into_iter()
.map(|file| JumpExampleRecentFile {
path: file.path,
cursor_position: file.cursor_position,
})
.collect()
}
fn render_jump_example_events(events: &[Arc<zeta_prompt::Event>], 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
}
crates/edit_prediction/src/zeta.rs
+17 -10
View File
@@ -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<crate::StoredEvent>,
Task<Result<collections::HashMap<Arc<Path>, Entity<BufferDiff>>>>,
)>,
capture_data: Option<Shared<Task<UncommittedDiffResult>>>,
repo_url: Option<String>,
cx: &mut Context<EditPredictionStore>,
) -> Task<Result<Option<EditPredictionResult>>> {
@@ -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,
crates/edit_prediction_cli/src/synthesize.rs
+3 -4
View File
@@ -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)
}
crates/fs/src/fake_git_repo.rs
+11 -3
View File
@@ -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<RepoPath>,
_commit_limit: usize,
) -> BoxFuture<'_, Result<Vec<FileHistoryChangedFileSets>>> {
async move { Ok(vec![FileHistoryChangedFileSets::default(); paths.len()]) }.boxed()
}
fn commit_data_reader(&self) -> Result<CommitDataReader> {
let fs = self.fs.clone();
let dot_git_path = self.dot_git_path.clone();
crates/git/src/repository.rs
+133
View File
@@ -512,6 +512,11 @@ pub struct CommitDiff {
pub files: Vec<CommitFile>,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct FileHistoryChangedFileSets {
pub file_sets: Vec<Vec<RepoPath>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum CommitFileStatus {
Added,
@@ -1047,6 +1052,12 @@ pub trait GitRepository: Send + Sync {
request_tx: Sender<Oid>,
) -> BoxFuture<'_, Result<()>>;
fn file_history_changed_files(
&self,
paths: Vec<RepoPath>,
commit_limit: usize,
) -> BoxFuture<'_, Result<Vec<FileHistoryChangedFileSets>>>;
fn commit_data_reader(&self) -> Result<CommitDataReader>;
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<RepoPath>,
commit_limit: usize,
) -> BoxFuture<'_, Result<Vec<FileHistoryChangedFileSets>>> {
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<CommitDataReader> {
let git_binary = self.git_binary();
@@ -3209,6 +3264,39 @@ async fn read_single_commit_response<R: smol::io::AsyncBufRead + Unpin>(
.ok_or_else(|| anyhow!("failed to parse commit {}", sha))
}
fn parse_file_history_changed_files_output(
output: &str,
queried_paths: &[RepoPath],
) -> Vec<FileHistoryChangedFileSets> {
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::<std::collections::BTreeSet<_>>();
if changed_files.is_empty() {
continue;
}
let file_set = changed_files.iter().cloned().collect::<Vec<_>>();
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<Item = &'a str>,
) -> Vec<Arc<InitialGraphCommitData>> {
@@ -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();
crates/project/src/git_store.rs
+27 -4
View File
@@ -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<RepoPath>,
commit_limit: usize,
) -> oneshot::Receiver<Result<Vec<FileHistoryChangedFileSets>>> {
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,