use crate::{ StoredEvent, data_collection::{ UncommittedDiffSnapshot, compute_cursor_excerpt, compute_uncommitted_diff, format_cursor_excerpt, }, example_spec::{ExampleSpec, RecentFile},};use anyhow::Result;use gpui::{App, Entity, Task};use language::Buffer;use project::Project;use std::{fmt::Write, path::Path, sync::Arc}; pub fn capture_example( project: Entity<Project>, buffer: Entity<Buffer>, cursor_anchor: language::Anchor, events: Vec<StoredEvent>, recently_opened_files: Vec<RecentFile>, recently_viewed_files: Vec<RecentFile>, uncommitted_diff_snapshot: UncommittedDiffSnapshot, populate_expected_patch: bool, cx: &mut App,) -> Option<Task<Result<ExampleSpec>>> { let snapshot = buffer.read(cx).snapshot(); let file = snapshot.file()?; let worktree_id = file.worktree_id(cx); let repository = project.read(cx).active_repository(cx)?; let repository_snapshot = repository.read(cx).snapshot(); let worktree = project.read(cx).worktree_for_id(worktree_id, cx)?; let root_name = worktree.read(cx).root_name_str().to_owned(); let cursor_path: Arc<Path> = file.path().as_std_path().into(); if worktree.read(cx).abs_path() != repository_snapshot.work_directory_abs_path { return None; } let repository_url = repository_snapshot .remote_origin_url .clone() .or_else(|| repository_snapshot.remote_upstream_url.clone())?; let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string(); Some(cx.spawn(async move |cx| { let uncommitted_diff = cx .background_executor() .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) .await; let line_comment_prefix = snapshot .language() .and_then(|lang| lang.config().line_comments.first()) .map(|s| s.to_string()) .unwrap_or_default(); let (cursor_excerpt, cursor_offset_in_excerpt, cursor_excerpt_range) = cx .background_executor() .spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) }) .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); if !edit_history.ends_with('\n') { edit_history.push('\n'); } } let uncommitted_diff_contains_edit_history = !edit_history.is_empty(); // Initialize an empty patch with context lines, to make it easy // to write the expected patch by hand. let mut expected_patches = Vec::new(); let mut rejected_patch = None; if populate_expected_patch { let mut empty_patch = String::new(); let start_row = cursor_excerpt_range.start.row + 1; let row_count = cursor_excerpt_range.end.row - cursor_excerpt_range.start.row + 1; writeln!(&mut empty_patch, "--- a/{}", cursor_path.display()).ok(); writeln!(&mut empty_patch, "+++ b/{}", cursor_path.display()).ok(); writeln!( &mut empty_patch, "@@ -{},{} +{},{} @@", start_row, row_count, start_row, row_count, ) .ok(); for line in cursor_excerpt.lines() { writeln!(&mut empty_patch, " {}", line).ok(); } expected_patches.push(empty_patch.clone()); rejected_patch = Some(empty_patch); } let spec = ExampleSpec { name: generate_timestamp_name(), repository_url, revision, tags: Vec::new(), reasoning: None, uncommitted_diff, recently_opened_files, recently_viewed_files, uncommitted_diff_contains_edit_history, cursor_path, cursor_position: format_cursor_excerpt( &cursor_excerpt, cursor_offset_in_excerpt, &line_comment_prefix, ), edit_history, expected_patches, rejected_patch, telemetry: None, human_feedback: Vec::new(), rating: None, }; Ok(spec) }))} 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 path.strip_prefix(root_name).unwrap_or(path).components() { output.push('/'); write!(output, "{}", component.as_os_str().to_string_lossy()).ok(); } } let zeta_prompt::Event::BufferChange { path, old_path, diff, .. } = event; output.push_str("--- a"); write_relative_path(output, old_path.as_ref(), root_name); output.push_str("\n+++ b"); write_relative_path(output, path.as_ref(), root_name); output.push('\n'); output.push_str(diff);} fn generate_timestamp_name() -> String { let format = time::format_description::parse("[year]-[month]-[day] [hour]:[minute]:[second]"); match format { Ok(format) => { let now = time::OffsetDateTime::now_local() .unwrap_or_else(|_| time::OffsetDateTime::now_utc()); now.format(&format) .unwrap_or_else(|_| "unknown-time".to_string()) } Err(_) => "unknown-time".to_string(), }} #[cfg(test)]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; use gpui::{AppContext as _, TestAppContext, http_client::FakeHttpClient}; use indoc::indoc; use language::{Anchor, Point}; use project::{FakeFs, Project}; use serde_json::json; use settings::SettingsStore; use std::path::Path; #[gpui::test] async fn test_capture_example(cx: &mut TestAppContext) { init_test(cx); let fs = FakeFs::new(cx.executor()); let committed_contents = indoc! {" fn main() { one(); two(); three(); four(); five(); six(); seven(); eight(); nine(); } "}; let disk_contents = indoc! {" fn main() { // comment 1 one(); two(); three(); four(); five(); six(); seven(); eight(); // comment 2 nine(); } "}; fs.insert_tree( "/project", json!({ ".git": {}, "src": { "deleted.rs": "pub fn deleted_file() {\n deleted();\n}\n", "main.rs": disk_contents, "new.rs": "pub fn new_file() {\n}\n", } }), ) .await; // Create an external file outside the main project fs.insert_tree( "/external", json!({ "external.rs": "fn external() {}\n", }), ) .await; fs.set_head_for_repo( Path::new("/project/.git"), &[ ( "src/deleted.rs", "pub fn deleted_file() {\n deleted();\n}\n".to_string(), ), ("src/main.rs", committed_contents.to_string()), ], "abc123def456", ); fs.set_remote_for_repo( Path::new("/project/.git"), "origin", "https://github.com/test/repo.git", ); let project = Project::test(fs.clone(), ["/project".as_ref()], cx).await; let buffer = project .update(cx, |project, cx| { project.open_local_buffer("/project/src/main.rs", cx) }) .await .unwrap(); let ep_store = cx.read(|cx| EditPredictionStore::try_global(cx).unwrap()); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&buffer, &project, cx) }); cx.run_until_parked(); let deleted_file_buffer = project .update(cx, |project, cx| { project.open_local_buffer("/project/src/deleted.rs", cx) }) .await .unwrap(); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&deleted_file_buffer, &project, cx) }); cx.run_until_parked(); deleted_file_buffer.update(cx, |buffer, cx| { buffer.edit([(0..buffer.len(), "")], None, cx); }); cx.run_until_parked(); buffer.update(cx, |buffer, cx| { let point = Point::new(6, 0); buffer.edit([(point..point, " // comment 3\n")], None, cx); let point = Point::new(4, 0); buffer.edit([(point..point, " // comment 4\n")], None, cx); pretty_assertions::assert_eq!( buffer.text(), indoc! {" fn main() { // comment 1 one(); two(); // comment 4 three(); four(); // comment 3 five(); six(); seven(); eight(); // comment 2 nine(); } "} ); }); cx.run_until_parked(); let new_file_buffer = project .update(cx, |project, cx| { project.open_local_buffer("/project/src/new.rs", cx) }) .await .unwrap(); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&new_file_buffer, &project, cx) }); cx.run_until_parked(); new_file_buffer.update(cx, |buffer, cx| { let point = Point::new(1, 0); buffer.edit([(point..point, " created();\n")], None, cx); }); cx.run_until_parked(); // Open and edit an external file (outside the main project's worktree) let external_buffer = project .update(cx, |project, cx| { project.open_local_buffer("/external/external.rs", cx) }) .await .unwrap(); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&external_buffer, &project, cx) }); cx.run_until_parked(); external_buffer.update(cx, |buffer, cx| { let point = Point::new(0, 0); buffer.edit([(point..point, "// external edit\n")], None, cx); }); cx.run_until_parked(); // Verify the external edit was recorded in events let events = ep_store.update(cx, |store, cx| store.edit_history_for_project(&project, cx)); assert!( matches!( events .last() .unwrap() .event .as_ref(), zeta_prompt::Event::BufferChange { path, .. } if path.as_ref() == "/external/external.rs" ), "external file edit should be in events" ); 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| { 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()), project_events, Vec::new(), Vec::new(), uncommitted_diffs_by_path, true, cx, ) .unwrap() }) .await .unwrap(); example.name = "test".to_string(); pretty_assertions::assert_eq!( example, ExampleSpec { name: "test".to_string(), repository_url: "https://github.com/test/repo.git".to_string(), revision: "abc123def456".to_string(), tags: Vec::new(), reasoning: None, uncommitted_diff: indoc! {" --- a/src/deleted.rs +++ b/src/deleted.rs @@ -1,3 +1,0 @@ -pub fn deleted_file() { - deleted(); -} --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,15 @@ fn main() { + // comment 1 one(); two(); + // comment 4 three(); four(); + // comment 3 five(); six(); seven(); eight(); + // comment 2 nine(); } --- /dev/null +++ b/src/new.rs @@ -0,0 +1,3 @@ +pub fn new_file() { + created(); +} "} .to_string(), recently_opened_files: Vec::new(), recently_viewed_files: Vec::new(), uncommitted_diff_contains_edit_history: true, cursor_path: Path::new("src/main.rs").into(), cursor_position: indoc! {" fn main() { ^[CURSOR_POSITION] // comment 1 one(); two(); // comment 4 three(); four(); // comment 3 five(); six(); seven(); eight(); // comment 2 nine(); } "} .to_string(), edit_history: indoc! {" --- a/src/deleted.rs +++ b/src/deleted.rs @@ -1,3 +1,0 @@ -pub fn deleted_file() { - deleted(); -} --- a/src/main.rs +++ b/src/main.rs @@ -2,8 +2,10 @@ // comment 1 one(); two(); + // comment 4 three(); four(); + // comment 3 five(); six(); seven(); --- a/src/new.rs +++ b/src/new.rs @@ -1,2 +1,3 @@ pub fn new_file() { + created(); } "} .to_string(), expected_patches: vec![ indoc! {" --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,16 @@ fn main() { // comment 1Lines 1–500 of 560. Load more
use crate::{ StoredEvent, data_collection::{ UncommittedDiffSnapshot, compute_cursor_excerpt, compute_uncommitted_diff, format_cursor_excerpt, }, example_spec::{ExampleSpec, RecentFile}, }; use anyhow::Result; use gpui::{App, Entity, Task}; use language::Buffer; use project::Project; use std::{fmt::Write, path::Path, sync::Arc}; pub fn capture_example( project: Entity<Project>, buffer: Entity<Buffer>, cursor_anchor: language::Anchor, events: Vec<StoredEvent>, recently_opened_files: Vec<RecentFile>, recently_viewed_files: Vec<RecentFile>, uncommitted_diff_snapshot: UncommittedDiffSnapshot, populate_expected_patch: bool, cx: &mut App, ) -> Option<Task<Result<ExampleSpec>>> { let snapshot = buffer.read(cx).snapshot(); let file = snapshot.file()?; let worktree_id = file.worktree_id(cx); let repository = project.read(cx).active_repository(cx)?; let repository_snapshot = repository.read(cx).snapshot(); let worktree = project.read(cx).worktree_for_id(worktree_id, cx)?; let root_name = worktree.read(cx).root_name_str().to_owned(); let cursor_path: Arc<Path> = file.path().as_std_path().into(); if worktree.read(cx).abs_path() != repository_snapshot.work_directory_abs_path { return None; } let repository_url = repository_snapshot .remote_origin_url .clone() .or_else(|| repository_snapshot.remote_upstream_url.clone())?; let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string(); Some(cx.spawn(async move |cx| { let uncommitted_diff = cx .background_executor() .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) .await; let line_comment_prefix = snapshot .language() .and_then(|lang| lang.config().line_comments.first()) .map(|s| s.to_string()) .unwrap_or_default(); let (cursor_excerpt, cursor_offset_in_excerpt, cursor_excerpt_range) = cx .background_executor() .spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) }) .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); if !edit_history.ends_with('\n') { edit_history.push('\n'); } } let uncommitted_diff_contains_edit_history = !edit_history.is_empty(); // Initialize an empty patch with context lines, to make it easy // to write the expected patch by hand. let mut expected_patches = Vec::new(); let mut rejected_patch = None; if populate_expected_patch { let mut empty_patch = String::new(); let start_row = cursor_excerpt_range.start.row + 1; let row_count = cursor_excerpt_range.end.row - cursor_excerpt_range.start.row + 1; writeln!(&mut empty_patch, "--- a/{}", cursor_path.display()).ok(); writeln!(&mut empty_patch, "+++ b/{}", cursor_path.display()).ok(); writeln!( &mut empty_patch, "@@ -{},{} +{},{} @@", start_row, row_count, start_row, row_count, ) .ok(); for line in cursor_excerpt.lines() { writeln!(&mut empty_patch, " {}", line).ok(); } expected_patches.push(empty_patch.clone()); rejected_patch = Some(empty_patch); } let spec = ExampleSpec { name: generate_timestamp_name(), repository_url, revision, tags: Vec::new(), reasoning: None, uncommitted_diff, recently_opened_files, recently_viewed_files, uncommitted_diff_contains_edit_history, cursor_path, cursor_position: format_cursor_excerpt( &cursor_excerpt, cursor_offset_in_excerpt, &line_comment_prefix, ), edit_history, expected_patches, rejected_patch, telemetry: None, human_feedback: Vec::new(), rating: None, }; Ok(spec) })) } 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 path.strip_prefix(root_name).unwrap_or(path).components() { output.push('/'); write!(output, "{}", component.as_os_str().to_string_lossy()).ok(); } } let zeta_prompt::Event::BufferChange { path, old_path, diff, .. } = event; output.push_str("--- a"); write_relative_path(output, old_path.as_ref(), root_name); output.push_str("\n+++ b"); write_relative_path(output, path.as_ref(), root_name); output.push('\n'); output.push_str(diff); } fn generate_timestamp_name() -> String { let format = time::format_description::parse("[year]-[month]-[day] [hour]:[minute]:[second]"); match format { Ok(format) => { let now = time::OffsetDateTime::now_local() .unwrap_or_else(|_| time::OffsetDateTime::now_utc()); now.format(&format) .unwrap_or_else(|_| "unknown-time".to_string()) } Err(_) => "unknown-time".to_string(), } } #[cfg(test)] 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; use gpui::{AppContext as _, TestAppContext, http_client::FakeHttpClient}; use indoc::indoc; use language::{Anchor, Point}; use project::{FakeFs, Project}; use serde_json::json; use settings::SettingsStore; use std::path::Path; #[gpui::test] async fn test_capture_example(cx: &mut TestAppContext) { init_test(cx); let fs = FakeFs::new(cx.executor()); let committed_contents = indoc! {" fn main() { one(); two(); three(); four(); five(); six(); seven(); eight(); nine(); } "}; let disk_contents = indoc! {" fn main() { // comment 1 one(); two(); three(); four(); five(); six(); seven(); eight(); // comment 2 nine(); } "}; fs.insert_tree( "/project", json!({ ".git": {}, "src": { "deleted.rs": "pub fn deleted_file() {\n deleted();\n}\n", "main.rs": disk_contents, "new.rs": "pub fn new_file() {\n}\n", } }), ) .await; // Create an external file outside the main project fs.insert_tree( "/external", json!({ "external.rs": "fn external() {}\n", }), ) .await; fs.set_head_for_repo( Path::new("/project/.git"), &[ ( "src/deleted.rs", "pub fn deleted_file() {\n deleted();\n}\n".to_string(), ), ("src/main.rs", committed_contents.to_string()), ], "abc123def456", ); fs.set_remote_for_repo( Path::new("/project/.git"), "origin", "https://github.com/test/repo.git", ); let project = Project::test(fs.clone(), ["/project".as_ref()], cx).await; let buffer = project .update(cx, |project, cx| { project.open_local_buffer("/project/src/main.rs", cx) }) .await .unwrap(); let ep_store = cx.read(|cx| EditPredictionStore::try_global(cx).unwrap()); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&buffer, &project, cx) }); cx.run_until_parked(); let deleted_file_buffer = project .update(cx, |project, cx| { project.open_local_buffer("/project/src/deleted.rs", cx) }) .await .unwrap(); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&deleted_file_buffer, &project, cx) }); cx.run_until_parked(); deleted_file_buffer.update(cx, |buffer, cx| { buffer.edit([(0..buffer.len(), "")], None, cx); }); cx.run_until_parked(); buffer.update(cx, |buffer, cx| { let point = Point::new(6, 0); buffer.edit([(point..point, " // comment 3\n")], None, cx); let point = Point::new(4, 0); buffer.edit([(point..point, " // comment 4\n")], None, cx); pretty_assertions::assert_eq!( buffer.text(), indoc! {" fn main() { // comment 1 one(); two(); // comment 4 three(); four(); // comment 3 five(); six(); seven(); eight(); // comment 2 nine(); } "} ); }); cx.run_until_parked(); let new_file_buffer = project .update(cx, |project, cx| { project.open_local_buffer("/project/src/new.rs", cx) }) .await .unwrap(); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&new_file_buffer, &project, cx) }); cx.run_until_parked(); new_file_buffer.update(cx, |buffer, cx| { let point = Point::new(1, 0); buffer.edit([(point..point, " created();\n")], None, cx); }); cx.run_until_parked(); // Open and edit an external file (outside the main project's worktree) let external_buffer = project .update(cx, |project, cx| { project.open_local_buffer("/external/external.rs", cx) }) .await .unwrap(); ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&external_buffer, &project, cx) }); cx.run_until_parked(); external_buffer.update(cx, |buffer, cx| { let point = Point::new(0, 0); buffer.edit([(point..point, "// external edit\n")], None, cx); }); cx.run_until_parked(); // Verify the external edit was recorded in events let events = ep_store.update(cx, |store, cx| store.edit_history_for_project(&project, cx)); assert!( matches!( events .last() .unwrap() .event .as_ref(), zeta_prompt::Event::BufferChange { path, .. } if path.as_ref() == "/external/external.rs" ), "external file edit should be in events" ); 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| { 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()), project_events, Vec::new(), Vec::new(), uncommitted_diffs_by_path, true, cx, ) .unwrap() }) .await .unwrap(); example.name = "test".to_string(); pretty_assertions::assert_eq!( example, ExampleSpec { name: "test".to_string(), repository_url: "https://github.com/test/repo.git".to_string(), revision: "abc123def456".to_string(), tags: Vec::new(), reasoning: None, uncommitted_diff: indoc! {" --- a/src/deleted.rs +++ b/src/deleted.rs @@ -1,3 +1,0 @@ -pub fn deleted_file() { - deleted(); -} --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,15 @@ fn main() { + // comment 1 one(); two(); + // comment 4 three(); four(); + // comment 3 five(); six(); seven(); eight(); + // comment 2 nine(); } --- /dev/null +++ b/src/new.rs @@ -0,0 +1,3 @@ +pub fn new_file() { + created(); +} "} .to_string(), recently_opened_files: Vec::new(), recently_viewed_files: Vec::new(), uncommitted_diff_contains_edit_history: true, cursor_path: Path::new("src/main.rs").into(), cursor_position: indoc! {" fn main() { ^[CURSOR_POSITION] // comment 1 one(); two(); // comment 4 three(); four(); // comment 3 five(); six(); seven(); eight(); // comment 2 nine(); } "} .to_string(), edit_history: indoc! {" --- a/src/deleted.rs +++ b/src/deleted.rs @@ -1,3 +1,0 @@ -pub fn deleted_file() { - deleted(); -} --- a/src/main.rs +++ b/src/main.rs @@ -2,8 +2,10 @@ // comment 1 one(); two(); + // comment 4 three(); four(); + // comment 3 five(); six(); seven(); --- a/src/new.rs +++ b/src/new.rs @@ -1,2 +1,3 @@ pub fn new_file() { + created(); } "} .to_string(), expected_patches: vec![ indoc! {" --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,16 @@ fn main() { // comment 1