use crate::{    EditPredictionExampleCaptureFeatureFlag, StoredEvent,    cursor_excerpt::editable_and_context_ranges_for_cursor_position,    example_spec::{        CapturedEvent, CapturedPromptInput, CapturedRelatedExcerpt, CapturedRelatedFile,        ExampleSpec, MAX_CURSOR_FILE_SIZE,    },};use anyhow::Result;use buffer_diff::BufferDiffSnapshot;use collections::HashMap;use feature_flags::FeatureFlagAppExt as _;use gpui::{App, Entity, Task};use language::{Buffer, ToPoint as _};use project::{Project, WorktreeId};use std::{collections::hash_map, fmt::Write as _, ops::Range, path::Path, sync::Arc};use text::{BufferSnapshot as TextBufferSnapshot, Point, ToOffset as _}; pub(crate) const DEFAULT_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS: u16 = 10;pub(crate) const DEFAULT_STAFF_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS: u16 = 100;pub(crate) const ZETA2_TESTING_RATE_PER_10K_PREDICTION: u16 = 500; pub fn capture_example(    project: Entity<Project>,    buffer: Entity<Buffer>,    cursor_anchor: language::Anchor,    mut events: Vec<StoredEvent>,    related_files: Vec<zeta_prompt::RelatedFile>,    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();     let git_store = project.read(cx).git_store().clone();     Some(cx.spawn(async move |mut cx| {        let snapshots_by_path =            collect_snapshots(&project, &git_store, worktree_id, &events, &mut cx).await?;         events.retain(|stored_event| {            let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref();            let relative_path = strip_root_name(path, &root_name);            snapshots_by_path.contains_key(relative_path)        });         let line_comment_prefix = snapshot            .language()            .and_then(|lang| lang.config().line_comments.first())            .map(|s| s.to_string())            .unwrap_or_default();         let full_cursor_offset = cursor_anchor.to_offset(&snapshot);        let cursor_point = cursor_anchor.to_point(&snapshot);        let cursor_file_content = if snapshot.len() <= MAX_CURSOR_FILE_SIZE {            Some(snapshot.text())        } else {            None        };         let (cursor_excerpt, cursor_offset_in_excerpt, cursor_excerpt_range) = cx            .background_executor()            .spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) })            .await;        let uncommitted_diff = cx            .background_executor()            .spawn(async move { compute_uncommitted_diff(snapshots_by_path) })            .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');            }        }         // 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 prompt_input = cursor_file_content.map(|content| {            let captured_events: Vec<CapturedEvent> = events                .iter()                .map(|stored_event| {                    let zeta_prompt::Event::BufferChange {                        path,                        old_path,                        diff,                        predicted,                        in_open_source_repo,                    } = stored_event.event.as_ref();                    CapturedEvent {                        path: strip_root_name(path, &root_name).into(),                        old_path: strip_root_name(old_path, &root_name).into(),                        diff: diff.clone(),                        predicted: *predicted,                        in_open_source_repo: *in_open_source_repo,                    }                })                .collect();             let captured_related_files: Vec<CapturedRelatedFile> = related_files                .iter()                .map(|rf| CapturedRelatedFile {                    path: strip_root_name(&rf.path, &root_name).into(),                    max_row: rf.max_row,                    excerpts: rf                        .excerpts                        .iter()                        .map(|e| CapturedRelatedExcerpt {                            row_range: e.row_range.clone(),                            text: e.text.to_string(),                        })                        .collect(),                })                .collect();             CapturedPromptInput {                cursor_file_content: content,                cursor_offset: full_cursor_offset,                cursor_row: cursor_point.row,                cursor_column: cursor_point.column,                excerpt_start_row: Some(0),                events: captured_events,                related_files: captured_related_files,            }        });         let mut spec = ExampleSpec {            name: generate_timestamp_name(),            repository_url,            revision,            tags: Vec::new(),            reasoning: None,            uncommitted_diff,            cursor_path,            cursor_position: String::new(),            edit_history,            expected_patches,            rejected_patch,            captured_prompt_input: prompt_input,            telemetry: None,            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(    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() {            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 compute_cursor_excerpt(    snapshot: &language::BufferSnapshot,    cursor_anchor: language::Anchor,) -> (String, usize, Range<Point>) {    use text::ToOffset as _;     let cursor_point = cursor_anchor.to_point(snapshot);    let (_editable_range, context_range) =        editable_and_context_ranges_for_cursor_position(cursor_point, snapshot, 100, 50);    let context_start_offset = context_range.start.to_offset(snapshot);    let cursor_offset = cursor_anchor.to_offset(snapshot);    let cursor_offset_in_excerpt = cursor_offset.saturating_sub(context_start_offset);    let excerpt = snapshot        .text_for_range(context_range.clone())        .collect::<String>();    (excerpt, cursor_offset_in_excerpt, context_range)} async fn collect_snapshots(    project: &Entity<Project>,    git_store: &Entity<project::git_store::GitStore>,    worktree_id: WorktreeId,    events: &[StoredEvent],    cx: &mut gpui::AsyncApp,) -> Result<HashMap<Arc<Path>, (TextBufferSnapshot, BufferDiffSnapshot)>> {    let mut snapshots_by_path = HashMap::default();    for stored_event in events {        let zeta_prompt::Event::BufferChange { path, .. } = stored_event.event.as_ref();        if let Some((project_path, relative_path)) = project.read_with(cx, |project, cx| {            let project_path = project                .find_project_path(path, cx)                .filter(|path| path.worktree_id == worktree_id)?;            let relative_path: Arc<Path> = project_path.path.as_std_path().into();            Some((project_path, relative_path))        }) {            if let hash_map::Entry::Vacant(entry) = snapshots_by_path.entry(relative_path) {                let buffer = project                    .update(cx, |project, cx| {                        project.open_buffer(project_path.clone(), cx)                    })                    .await?;                let diff = git_store                    .update(cx, |git_store, cx| {                        git_store.open_uncommitted_diff(buffer.clone(), cx)                    })                    .await?;                let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx));                entry.insert((stored_event.old_snapshot.clone(), diff_snapshot));            }        }    }    Ok(snapshots_by_path)} fn compute_uncommitted_diff(    snapshots_by_path: HashMap<Arc<Path>, (TextBufferSnapshot, BufferDiffSnapshot)>,) -> String {    let mut uncommitted_diff = String::new();    for (relative_path, (before_text, diff_snapshot)) in snapshots_by_path {        if let Some(head_text) = &diff_snapshot.base_text_string() {            let file_diff = language::unified_diff(head_text, &before_text.text());            if !file_diff.is_empty() {                let path_str = relative_path.to_string_lossy();                writeln!(uncommitted_diff, "--- a/{path_str}").ok();                writeln!(uncommitted_diff, "+++ b/{path_str}").ok();                uncommitted_diff.push_str(&file_diff);                if !uncommitted_diff.ends_with('\n') {                    uncommitted_diff.push('\n');                }            }        }    }    uncommitted_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(),    }} pub(crate) fn should_sample_edit_prediction_example_capture(cx: &App) -> bool {    let default_rate = if cx.is_staff() {        DEFAULT_STAFF_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS    } else {        DEFAULT_EXAMPLE_CAPTURE_RATE_PER_10K_PREDICTIONS    };    let capture_rate = language::language_settings::all_language_settings(None, cx)        .edit_predictions        .example_capture_rate        .unwrap_or(default_rate);    cx.has_flag::<EditPredictionExampleCaptureFeatureFlag>()        && rand::random::<u16>() % 10_000 < capture_rate} pub(crate) fn should_send_testing_zeta2_request() -> bool {    rand::random::<u16>() % 10_000 < ZETA2_TESTING_RATE_PER_10K_PREDICTION} #[cfg(test)]mod tests {    use super::*;    use crate::EditPredictionStore;    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": {                    "main.rs": disk_contents,                }            }),        )        .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/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();         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();         // 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_with_pause_split_last_event(&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 mut example = cx            .update(|cx| {                capture_example(                    project.clone(),                    buffer.clone(),                    Anchor::MIN,                    events,                    Vec::new(),                    true,                    cx,                )                .unwrap()            })            .await            .unwrap();        example.name = "test".to_string();

Lines 1–500 of 656. Load more