diff --git a/Cargo.lock b/Cargo.lock index 16078d7982..b38b1ffe78 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3265,7 +3265,6 @@ dependencies = [ "serde", "serde_json", "strum 0.27.2", - "uuid", "zeta_prompt", ] diff --git a/crates/cloud_api_types/Cargo.toml b/crates/cloud_api_types/Cargo.toml index 66080cf9ec..a220f688f5 100644 --- a/crates/cloud_api_types/Cargo.toml +++ b/crates/cloud_api_types/Cargo.toml @@ -19,7 +19,6 @@ cloud_llm_client.workspace = true serde.workspace = true serde_json.workspace = true strum.workspace = true -uuid.workspace = true zeta_prompt.workspace = true [dev-dependencies] diff --git a/crates/cloud_api_types/src/cloud_api_types.rs b/crates/cloud_api_types/src/cloud_api_types.rs index 27fd96b966..6fef7c586d 100644 --- a/crates/cloud_api_types/src/cloud_api_types.rs +++ b/crates/cloud_api_types/src/cloud_api_types.rs @@ -6,6 +6,7 @@ mod timestamp; pub mod websocket_protocol; use std::collections::BTreeMap; +use std::ops::Range; use std::path::Path; use std::sync::Arc; @@ -129,58 +130,58 @@ pub struct SubmitEditPredictionFeedbackBody { } #[derive(Debug, PartialEq, Serialize, Deserialize)] -pub struct SubmitEditPredictionSettledBody { +pub struct SettledEditPrediction { pub request_id: String, + #[serde(skip_serializing_if = "Option::is_none")] pub settled_editable_region: Option, pub ts_error_count_before_prediction: usize, pub ts_error_count_after_prediction: usize, pub can_collect_data: bool, pub is_in_open_source_repo: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sample_data: Option, #[serde(flatten)] pub kept_chars: EditPredictionSettledKeptChars, pub example: Option, pub model_version: Option, #[serde(rename = "e2e_latency")] - pub e2e_latency_ms: u128, + pub e2e_latency_ms: u64, +} + +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub struct SettledEditPredictionSampleData { + pub repository_url: Option, + pub revision: Option, + /// Note: this is only the uncommitted diff for files in `edit_history` + /// This is done to avoid excessive memory usage + #[serde(default, skip_serializing_if = "Option::is_none")] + pub uncommitted_diff: Option, + pub editable_path: Arc, + pub editable_offset_range: Range, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub buffer_diagnostics: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub future_edit_history_events: Vec>, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub navigation_history: Vec, + pub edit_events_before_quiescence: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub next_edit_cursor_offset: Option, +} + +pub const MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST: usize = 32; + +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub struct SubmitEditPredictionSettledBatchBody { + pub predictions: Vec, } #[derive(Debug, PartialEq, Serialize, Deserialize)] pub struct SubmitEditPredictionSettledResponse {} #[derive(Debug, PartialEq, Serialize, Deserialize)] -pub struct SubmitEditPredictionJumpExampleBody { - pub request_id: uuid::Uuid, - pub trigger: JumpExampleTrigger, - pub repository_url: Option, - pub revision: Option, - /// Note: this is only the uncommitted diff for files in `edit_history` - /// This is done to avoid excessive memory usage - pub uncommitted_diff: Option, - pub recently_opened_files: Vec, - pub recently_viewed_files: Vec, - pub cursor_path: Arc, - pub cursor_position: String, - pub edit_history: Vec>, - pub diagnostics: Vec, - pub future_edit_history: String, - pub navigation_history: Vec, - pub is_in_open_source_repo: bool, - pub can_collect_data: bool, -} - -#[derive(Debug, PartialEq, Serialize, Deserialize)] -pub struct SubmitEditPredictionJumpExampleResponse {} - -#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum JumpExampleTrigger { - Prediction, - Diagnostic, -} - -#[derive(Debug, PartialEq, Serialize, Deserialize)] -pub struct JumpExampleRecentFile { +pub struct EditPredictionRecentFile { pub path: Arc, #[serde(default, skip_serializing_if = "Option::is_none")] pub cursor_position: Option, diff --git a/crates/edit_prediction/src/capture_example.rs b/crates/edit_prediction/src/capture_example.rs deleted file mode 100644 index 4f6e0f34b3..0000000000 --- a/crates/edit_prediction/src/capture_example.rs +++ /dev/null @@ -1,565 +0,0 @@ -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, - buffer: Entity, - cursor_anchor: language::Anchor, - events: Vec, - recently_opened_files: Vec, - recently_viewed_files: Vec, - uncommitted_diff_snapshot: UncommittedDiffSnapshot, - populate_expected_patch: bool, - cx: &mut App, -) -> Option>> { - 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 = 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 unfiltered_uncommitted_diffs = ep_store - .update(cx, |_store, cx| { - uncommitted_diffs_for_events(project.clone(), worktree_id, events.clone(), cx) - }) - .await - .unwrap(); - assert!( - unfiltered_uncommitted_diffs - .iter() - .all(|(path, _, _)| path.as_ref() != Path::new("external.rs")) - ); - - let project_events = events - .into_iter() - .filter(|event| { - let zeta_prompt::Event::BufferChange { path, .. } = event.event.as_ref(); - path.as_ref() != "/external/external.rs" - }) - .collect::>(); - let uncommitted_diffs_by_path = ep_store - .update(cx, |_store, cx| { - 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 - one(); - two(); - // comment 4 - three(); - four(); - // comment 3 - five(); - six(); - seven(); - eight(); - // comment 2 - nine(); - } - "} - .to_string() - ], - rejected_patch: Some( - indoc! {" - --- a/src/main.rs - +++ b/src/main.rs - @@ -1,16 +1,16 @@ - fn main() { - // comment 1 - one(); - two(); - // comment 4 - three(); - four(); - // comment 3 - five(); - six(); - seven(); - eight(); - // comment 2 - nine(); - } - "} - .to_string() - ), - telemetry: None, - human_feedback: Vec::new(), - rating: None, - } - ); - } - - fn init_test(cx: &mut TestAppContext) { - cx.update(|cx| { - let settings_store = SettingsStore::test(cx); - cx.set_global(settings_store); - zlog::init_test(); - let http_client = FakeHttpClient::with_404_response(); - let client = Client::new(Arc::new(FakeSystemClock::new()), http_client, cx); - let user_store = cx.new(|cx| UserStore::new(client.clone(), cx)); - language_model::init(cx); - RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx); - EditPredictionStore::global(&client, &user_store, cx); - }) - } -} diff --git a/crates/edit_prediction/src/capture_prediction_context.rs b/crates/edit_prediction/src/capture_prediction_context.rs new file mode 100644 index 0000000000..449ce8da70 --- /dev/null +++ b/crates/edit_prediction/src/capture_prediction_context.rs @@ -0,0 +1,85 @@ +use crate::{ + EditPredictionStore, StoredEvent, + data_collection::{compute_uncommitted_diff, uncommitted_diffs_for_events}, + zeta, +}; +use anyhow::Result; +use gpui::{Context, Entity, Task}; +use language::{Buffer, Point, ToPoint as _}; +use project::Project; +use text::OffsetRangeExt as _; + +const MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024; + +pub(crate) struct CapturedPredictionContext { + pub(crate) repository_url: Option, + pub(crate) revision: Option, + pub(crate) uncommitted_diff: Option, + pub(crate) buffer_diagnostics: Vec, +} + +pub(crate) fn capture_prediction_context( + project: Entity, + buffer: Entity, + cursor_anchor: language::Anchor, + stored_events: Vec, + repository_url: Option, + revision: Option, + cx: &mut Context, +) -> Option>> { + let snapshot = buffer.read(cx).snapshot(); + let worktree_id = snapshot.file()?.worktree_id(cx); + let uncommitted_diff_task = + uncommitted_diffs_for_events(project, worktree_id, stored_events, cx); + + Some(cx.spawn(async move |_this, cx| { + let uncommitted_diff_snapshot = match uncommitted_diff_task.await { + Ok(snapshot) => Some(snapshot), + Err(error) => { + log::debug!("failed to capture uncommitted diff: {error:?}"); + None + } + }; + + let uncommitted_diff = if let Some(uncommitted_diff_snapshot) = uncommitted_diff_snapshot { + let estimated_uncommitted_diff_size = uncommitted_diff_snapshot + .iter() + .map(|(_, buffer_snapshot, diff_snapshot)| { + diff_snapshot + .hunks(buffer_snapshot) + .map(|hunk| { + hunk.diff_base_byte_range.len() + + hunk.range.to_offset(buffer_snapshot).len() + }) + .sum::() + }) + .sum::(); + + if estimated_uncommitted_diff_size <= MAX_UNCOMMITTED_DIFF_SIZE { + let uncommitted_diff = cx + .background_executor() + .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) + .await; + (uncommitted_diff.len() <= MAX_UNCOMMITTED_DIFF_SIZE).then_some(uncommitted_diff) + } else { + None + } + } else { + None + }; + + let buffer_diagnostics = zeta::active_buffer_diagnostics( + &snapshot, + Point::new(0, 0)..snapshot.max_point(), + cursor_anchor.to_point(&snapshot).row, + 100, + ); + + Ok(CapturedPredictionContext { + repository_url, + revision, + uncommitted_diff, + buffer_diagnostics, + }) + })) +} diff --git a/crates/edit_prediction/src/data_collection.rs b/crates/edit_prediction/src/data_collection.rs index b56609a9d9..4be37c03b8 100644 --- a/crates/edit_prediction/src/data_collection.rs +++ b/crates/edit_prediction/src/data_collection.rs @@ -7,7 +7,7 @@ use gpui::{Context, Entity, Task}; use language::BufferSnapshot; use project::{Project, ProjectPath, WorktreeId}; use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc}; -use text::{OffsetRangeExt, Point}; +use text::Point; use util::rel_path::RelPath; pub type UncommittedDiffSnapshot = Vec<(Arc, BufferSnapshot, BufferDiffSnapshot)>; @@ -234,17 +234,6 @@ pub(crate) fn compute_uncommitted_diff(snapshot: UncommittedDiffSnapshot) -> Str 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() diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 1136a482a8..85c4f9279c 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -3,8 +3,8 @@ use buffer_diff::BufferDiff; use client::{Client, EditPredictionUsage, UserStore, global_llm_token}; use cloud_api_client::LlmApiToken; use cloud_api_types::{ - EditPredictionSettledKeptChars, OrganizationId, SubmitEditPredictionFeedbackBody, - SubmitEditPredictionSettledBody, + EditPredictionRecentFile, EditPredictionSettledKeptChars, OrganizationId, + SettledEditPrediction, SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody, }; use cloud_llm_client::predict_edits_v3::{ PREDICT_EDITS_MODE_HEADER_NAME, PREDICT_EDITS_REQUEST_ID_HEADER_NAME, @@ -74,7 +74,6 @@ 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; @@ -85,7 +84,7 @@ mod prediction; pub mod udiff; -mod capture_example; +mod capture_prediction_context; pub mod open_ai_compatible; mod zed_edit_prediction_delegate; pub mod zeta; @@ -93,14 +92,9 @@ pub mod zeta; #[cfg(test)] mod edit_prediction_tests; +use crate::capture_prediction_context::{CapturedPredictionContext, capture_prediction_context}; 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}; @@ -136,6 +130,10 @@ const REQUEST_TIMEOUT_BACKOFF: Duration = Duration::from_secs(10); const EDIT_PREDICTION_SETTLED_TTL: Duration = Duration::from_secs(60 * 5); const EDIT_PREDICTION_SETTLED_QUIESCENCE: Duration = Duration::from_secs(10); +const EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS: usize = 4; +/// The server rejects settled bodies larger than 64 KiB (compressed). +const EDIT_PREDICTION_SETTLED_MAX_BODY_BYTES: usize = 63 * 1024; +const EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES: usize = 4 * 1024; pub struct EditPredictionJumpsFeatureFlag; @@ -206,7 +204,6 @@ pub struct EditPredictionModelInput { snapshot: BufferSnapshot, position: Anchor, events: Vec>, - stored_events: Vec, related_files: Vec, mode: PredictEditsMode, trigger: PredictEditsRequestTrigger, @@ -358,6 +355,7 @@ fn push_recent_file(files: &mut VecDeque, mut file: RecentFile) { struct ProjectState { events: VecDeque, last_event: Option, + next_last_event_seq: u64, recently_viewed_files: VecDeque, recently_opened_files: VecDeque, registered_buffers: HashMap, @@ -366,8 +364,7 @@ struct ProjectState { last_edit_source: Option, next_pending_prediction_id: usize, pending_predictions: ArrayVec, - pending_jump_example_captures: Vec, - starting_jump_example_captures: Vec, + pending_prediction_captures: Vec, debug_tx: Option>, last_edit_prediction_refresh: Option<(EntityId, Instant)>, last_jump_prediction_refresh: Option<(EntityId, Instant)>, @@ -471,22 +468,36 @@ impl ProjectState { } fn finalize_last_event(&mut self, cx: &mut Context) { - let Some(event) = self.last_event.take() else { - return; - }; - let Some(event) = event.finalize(&self.license_detection_watchers, cx) else { + let Some(last_event) = self.last_event.take() else { return; }; + let event = last_event.finalize(&self.license_detection_watchers, cx); - for capture in &mut self.pending_jump_example_captures { - capture.future_events.push(event.event.clone()); + for capture in &mut self.pending_prediction_captures { + capture.try_record_future_event( + &last_event, + event.as_ref(), + &self.license_detection_watchers, + cx, + ); } - jump_example::drain_completed_jump_example_captures(self, cx); + + let Some(event) = event else { + return; + }; if self.events.len() + 1 >= EVENT_COUNT_MAX { self.events.pop_front(); } self.events.push_back(event); } + + fn clear_history(&mut self) { + self.events.clear(); + self.last_event.take(); + for capture in &mut self.pending_prediction_captures { + capture.sample_data = None; + } + } } #[derive(Debug, Clone)] @@ -588,9 +599,9 @@ impl std::ops::Deref for BufferEditPrediction<'_> { } } -#[derive(Clone)] -struct PendingSettledPrediction { +struct PendingPredictionCapture { request_id: EditPredictionId, + edited_buffer_id: EntityId, editable_anchor_range: Range, editable_region_before_prediction: String, predicted_editable_region: String, @@ -599,23 +610,104 @@ struct PendingSettledPrediction { organization_id: Option, can_collect_data: bool, is_in_open_source_repo: bool, - example: Option, + sample_data: Option, model_version: Option, enqueued_at: Instant, last_edit_at: Instant, e2e_latency: std::time::Duration, } +struct PendingPredictionCaptureSampleData { + context_task: Task>, + editable_path: Arc, + editable_offset_range: Range, + next_edit_cursor_offset: Option, + future_edit_history_events: Vec>, + navigation_history: VecDeque, + edit_events_before_quiescence: u32, + prompt_history_boundary: Option, +} + +/// Marks where the prompt's edit history ended. Sample data may only include +/// content the user produced after this point. +struct PromptHistoryBoundary { + /// The seq of the first event this capture is expected to observe: the + /// event that was pending when the prediction was requested, or the next + /// event to be created if none was pending. Observing a later seq first + /// means events were lost while the prediction request was in flight. + first_event_seq: u64, + /// The prompt's end snapshot within the event that was pending when the + /// prediction was requested, if any. The first observed event is trimmed + /// to its suffix after this snapshot. + snapshot: Option, +} + +impl PendingPredictionCapture { + /// Records the project's last event (pending or finalizing) into this + /// sample's future edit history. Returns false if the sample must be + /// dropped because its future history can't be captured accurately. + fn try_record_future_event( + &mut self, + last_event: &LastEvent, + finalized_event: Option<&StoredEvent>, + license_detection_watchers: &HashMap>, + cx: &App, + ) { + let Some(sample) = &mut self.sample_data else { + return; + }; + let boundary = sample.prompt_history_boundary.take(); + let suffix_snapshot = match &boundary { + Some(boundary) => { + if last_event.seq != boundary.first_event_seq { + // Events were finalized before this capture was enqueued, + // so events are missing from the future history. + self.sample_data.take(); + return; + } + boundary.snapshot.as_ref() + } + None => None, + }; + + let event = match suffix_snapshot { + Some(snapshot) => { + let suffix = last_event + .suffix_after(snapshot) + .and_then(|suffix| suffix.finalize(license_detection_watchers, cx)); + let Some(suffix) = suffix else { + return; + }; + suffix.event + } + None => match finalized_event { + Some(event) => event.event.clone(), + None => return, + }, + }; + + if !event.in_open_source_repo() { + self.sample_data.take(); + return; + } + sample.edit_events_before_quiescence += 1; + if sample.future_edit_history_events.len() < EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS { + sample.future_edit_history_events.push(event); + } + } +} + struct RegisteredBuffer { file: Option>, snapshot: TextBufferSnapshot, - pending_predictions: Vec, last_position: Option, _subscriptions: [gpui::Subscription; 2], } #[derive(Clone)] struct LastEvent { + /// Project-wide monotonic sequence number identifying this event. + seq: u64, old_snapshot: TextBufferSnapshot, new_snapshot: TextBufferSnapshot, old_file: Option>, @@ -649,7 +741,7 @@ impl LastEvent { }) }); - let (diff, edit_range) = compute_diff_between_snapshots_in_range( + let (diff, old_range, new_range) = compute_diff_between_snapshots_in_range( &self.old_snapshot, &self.new_snapshot, &self.total_edit_range, @@ -663,13 +755,15 @@ impl LastEvent { old_path, path, diff, + old_range, + new_range: new_range.clone(), in_open_source_repo, predicted: self.predicted, }), old_snapshot: self.old_snapshot.clone(), 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), + total_edit_range: self.new_snapshot.anchor_before(new_range.start) + ..self.new_snapshot.anchor_before(new_range.end), file_context: self.file_context.clone(), }) } @@ -680,50 +774,41 @@ impl LastEvent { return (self.clone(), None); }; + let Some(after) = self.suffix_after(boundary_snapshot) else { + return (self.clone(), None); + }; + let total_edit_range_before_pause = self .total_edit_range_at_last_pause_boundary .clone() .unwrap_or_else(|| self.total_edit_range.clone()); - let Some(total_edit_range_after_pause) = - compute_total_edit_range_between_snapshots(boundary_snapshot, &self.new_snapshot) - else { - return (self.clone(), None); - }; - - let latest_edit_range_before_pause = total_edit_range_before_pause.clone(); - let latest_edit_range_after_pause = total_edit_range_after_pause.clone(); - let before = LastEvent { - old_snapshot: self.old_snapshot.clone(), new_snapshot: boundary_snapshot.clone(), - old_file: self.old_file.clone(), - new_file: self.new_file.clone(), - latest_edit_range: latest_edit_range_before_pause, + latest_edit_range: total_edit_range_before_pause.clone(), total_edit_range: total_edit_range_before_pause, total_edit_range_at_last_pause_boundary: None, - 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 { - old_snapshot: boundary_snapshot.clone(), - new_snapshot: self.new_snapshot.clone(), - old_file: self.old_file.clone(), - new_file: self.new_file.clone(), - latest_edit_range: latest_edit_range_after_pause, - total_edit_range: total_edit_range_after_pause, - total_edit_range_at_last_pause_boundary: None, - predicted: self.predicted, - snapshot_after_last_editing_pause: None, - last_edit_time: self.last_edit_time, - file_context: self.file_context.clone(), + ..self.clone() }; (before, Some(after)) } + + /// The portion of this event that happened after `boundary_snapshot`, or + /// None if the buffer hasn't changed since. + pub fn suffix_after(&self, boundary_snapshot: &TextBufferSnapshot) -> Option { + let total_edit_range = + compute_total_edit_range_between_snapshots(boundary_snapshot, &self.new_snapshot)?; + Some(LastEvent { + old_snapshot: boundary_snapshot.clone(), + latest_edit_range: total_edit_range.clone(), + total_edit_range, + total_edit_range_at_last_pause_boundary: None, + snapshot_after_last_editing_pause: None, + ..self.clone() + }) + } } fn compute_total_edit_range_between_snapshots( @@ -791,12 +876,16 @@ fn compute_diff_between_snapshots_in_range( old_snapshot: &TextBufferSnapshot, new_snapshot: &TextBufferSnapshot, total_edit_range: &Range, -) -> Option<(String, Range)> { - let new_start_point = total_edit_range.start.to_point(new_snapshot); - let new_end_point = total_edit_range.end.to_point(new_snapshot); +) -> Option<(String, Range, Range)> { + let new_start_offset = total_edit_range.start.to_offset(new_snapshot); + let new_end_offset = total_edit_range.end.to_offset(new_snapshot); + let new_start_point = new_snapshot.offset_to_point(new_start_offset); + let new_end_point = new_snapshot.offset_to_point(new_end_offset); let old_range = compute_old_range_for_new_range(old_snapshot, new_snapshot, total_edit_range)?; let old_start_point = old_range.start; let old_end_point = old_range.end; + let old_start_offset = old_snapshot.point_to_offset(old_start_point); + let old_end_offset = old_snapshot.point_to_offset(old_end_point); const CONTEXT_LINES: u32 = 3; @@ -832,7 +921,11 @@ fn compute_diff_between_snapshots_in_range( new_context_start_row, ); - Some((diff, new_start_point..new_end_point)) + Some(( + diff, + old_start_offset..old_end_offset, + new_start_offset..new_end_offset, + )) } pub(crate) fn buffer_path_with_id_fallback( @@ -1131,15 +1224,13 @@ impl EditPredictionStore { pub fn clear_history(&mut self) { for project_state in self.projects.values_mut() { - project_state.events.clear(); - project_state.last_event.take(); + project_state.clear_history(); } } pub fn clear_history_for_project(&mut self, project: &Entity) { if let Some(project_state) = self.projects.get_mut(&project.entity_id()) { - project_state.events.clear(); - project_state.last_event.take(); + project_state.clear_history(); } } @@ -1337,6 +1428,7 @@ impl EditPredictionStore { }, events: VecDeque::new(), last_event: None, + next_last_event_seq: 0, recently_viewed_files: VecDeque::new(), recently_opened_files: VecDeque::new(), debug_tx: None, @@ -1346,8 +1438,7 @@ impl EditPredictionStore { last_edit_source: None, cancelled_predictions: HashSet::default(), pending_predictions: ArrayVec::new(), - pending_jump_example_captures: Vec::new(), - starting_jump_example_captures: Vec::new(), + pending_prediction_captures: Vec::new(), next_pending_prediction_id: 0, last_edit_prediction_refresh: None, last_jump_prediction_refresh: None, @@ -1475,14 +1566,23 @@ impl EditPredictionStore { 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); + let can_collect_navigation = project_state + .license_detection_watchers + .get(&path.worktree_id) + .is_some_and(|watcher| watcher.is_project_open_source()); + for capture in &mut project_state.pending_prediction_captures { + if let Some(sample_data) = capture.sample_data.as_mut() { + if can_collect_navigation { + push_recent_file( + &mut sample_data.navigation_history, + recent_file.clone(), + ); + } else { + capture.sample_data = None; + } } } push_recent_file(&mut project_state.recently_viewed_files, recent_file); - jump_example::drain_completed_jump_example_captures(project_state, cx); } } project::Event::DiagnosticsUpdated { .. } => { @@ -1549,7 +1649,6 @@ impl EditPredictionStore { snapshot, file, last_position: None, - pending_predictions: Vec::new(), _subscriptions: [ cx.subscribe(buffer, { let project = project.downgrade(); @@ -1617,9 +1716,19 @@ impl EditPredictionStore { return; }; - for pending_prediction in &mut registered_buffer.pending_predictions { - if edit_range.overlaps(&pending_prediction.editable_anchor_range, &new_snapshot) { - pending_prediction.last_edit_at = now; + for pending_capture in &mut project_state.pending_prediction_captures { + if pending_capture.edited_buffer_id == buffer.entity_id() + && edit_range.overlaps(&pending_capture.editable_anchor_range, &new_snapshot) + { + pending_capture.last_edit_at = now; + if is_local + && !is_predicted + && let Some(sample_data) = pending_capture.sample_data.as_mut() + && sample_data.next_edit_cursor_offset.is_none() + { + sample_data.next_edit_cursor_offset = + Some(edit_range.start.to_offset(&new_snapshot)); + } } } @@ -1697,7 +1806,10 @@ impl EditPredictionStore { file_context }); + let seq = project_state.next_last_event_seq; + project_state.next_last_event_seq += 1; project_state.last_event = Some(LastEvent { + seq, old_file, new_file, old_snapshot, @@ -1900,48 +2012,79 @@ impl EditPredictionStore { let mut oldest_edited_at = None; let mut ready_predictions = Vec::new(); - this.update(cx, |this, _| { - for (_, project_state) in this.projects.iter_mut() { - for (_, registered_buffer) in project_state.registered_buffers.iter_mut() { - let mut pending_index = 0; - while pending_index < registered_buffer.pending_predictions.len() { - let pending_prediction = - ®istered_buffer.pending_predictions[pending_index]; - let age = now.saturating_duration_since(pending_prediction.enqueued_at); - if age >= EDIT_PREDICTION_SETTLED_TTL { - registered_buffer.pending_predictions.remove(pending_index); - continue; - } - - let quiet_for = - now.saturating_duration_since(pending_prediction.last_edit_at); - if quiet_for >= EDIT_PREDICTION_SETTLED_QUIESCENCE { - let pending_prediction = - registered_buffer.pending_predictions.remove(pending_index); - let settled_editable_region = registered_buffer - .snapshot - .text_for_range( - pending_prediction.editable_anchor_range.clone(), - ) - .collect::(); - ready_predictions - .push((pending_prediction, settled_editable_region)); - continue; - } - - if oldest_edited_at - .is_none_or(|time| pending_prediction.last_edit_at < time) - { - oldest_edited_at = Some(pending_prediction.last_edit_at); - } - pending_index += 1; + this.update(cx, |this, cx| { + for project_state in this.projects.values_mut() { + let ProjectState { + last_event, + registered_buffers, + license_detection_watchers, + pending_prediction_captures, + .. + } = project_state; + let pending_last_event = last_event.as_ref().map(|last_event| { + ( + last_event, + last_event.finalize(license_detection_watchers, cx), + ) + }); + let mut pending_index = 0; + while pending_index < pending_prediction_captures.len() { + let pending_capture = &pending_prediction_captures[pending_index]; + let age = now.saturating_duration_since(pending_capture.enqueued_at); + if age >= EDIT_PREDICTION_SETTLED_TTL { + pending_prediction_captures.remove(pending_index); + continue; } + + let quiet_for = now.saturating_duration_since(pending_capture.last_edit_at); + if quiet_for >= EDIT_PREDICTION_SETTLED_QUIESCENCE { + let Some(registered_buffer) = + registered_buffers.get(&pending_capture.edited_buffer_id) + else { + pending_prediction_captures.remove(pending_index); + continue; + }; + let editable_offset_range = pending_capture + .editable_anchor_range + .to_offset(®istered_buffer.snapshot); + if editable_offset_range.len() + > EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES + { + // The prediction was obliterated by a huge edit; + // kept-rate against it would be meaningless and the + // region would blow the body size cap. + pending_prediction_captures.remove(pending_index); + continue; + } + let settled_editable_region = registered_buffer + .snapshot + .text_for_range(editable_offset_range) + .collect::(); + let mut pending_capture = + pending_prediction_captures.remove(pending_index); + if let Some((last_event, finalized_event)) = pending_last_event.as_ref() + { + pending_capture.try_record_future_event( + last_event, + finalized_event.as_ref(), + license_detection_watchers, + cx, + ); + } + ready_predictions.push((pending_capture, settled_editable_region)); + continue; + } + + if oldest_edited_at.is_none_or(|time| pending_capture.last_edit_at < time) { + oldest_edited_at = Some(pending_capture.last_edit_at); + } + pending_index += 1; } } }); - for (pending_prediction, settled_editable_region) in ready_predictions { - let PendingSettledPrediction { + for (pending_capture, settled_editable_region) in ready_predictions { + let PendingPredictionCapture { request_id, editable_region_before_prediction, predicted_editable_region, @@ -1950,11 +2093,11 @@ impl EditPredictionStore { organization_id, can_collect_data, is_in_open_source_repo, - example, + sample_data, model_version, e2e_latency, .. - } = pending_prediction; + } = pending_capture; let settled_editable_region_for_metrics = settled_editable_region.clone(); #[cfg(test)] { @@ -1978,21 +2121,46 @@ impl EditPredictionStore { ); let result: anyhow::Result<()> = async { - let settled_editable_region = - can_collect_data.then_some(settled_editable_region); - let example = if can_collect_data { - example.map(serde_json::to_value).transpose()? + let sample_data = if can_collect_data + && let Some(sample_data) = sample_data + && let Ok(context) = sample_data.context_task.await + { + Some(SettledEditPredictionSampleData { + repository_url: context.repository_url, + revision: context.revision, + uncommitted_diff: context.uncommitted_diff, + editable_path: sample_data.editable_path, + editable_offset_range: sample_data.editable_offset_range, + buffer_diagnostics: context.buffer_diagnostics, + future_edit_history_events: sample_data + .future_edit_history_events, + navigation_history: sample_data + .navigation_history + .into_iter() + .map(|file| EditPredictionRecentFile { + path: file.path, + cursor_position: file.cursor_position, + }) + .collect(), + edit_events_before_quiescence: sample_data + .edit_events_before_quiescence, + next_edit_cursor_offset: sample_data.next_edit_cursor_offset, + }) } else { None }; - let body = SubmitEditPredictionSettledBody { + let settled_editable_region = + can_collect_data.then_some(settled_editable_region); + + let mut body = SettledEditPrediction { request_id: request_id.0.to_string(), settled_editable_region, ts_error_count_before_prediction, ts_error_count_after_prediction, can_collect_data, is_in_open_source_repo, + sample_data, kept_chars: EditPredictionSettledKeptChars { candidate_new: kept_rate_result.candidate_new_chars, reference_new: kept_rate_result.reference_new_chars, @@ -2005,13 +2173,18 @@ impl EditPredictionStore { kept_rate: kept_rate_result.kept_rate, recall_rate: kept_rate_result.recall_rate, }, - example, + example: None, model_version, - e2e_latency_ms: e2e_latency.as_millis(), + e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX)) + as u64, }; - let json_bytes = serde_json::to_vec(&body)?; - let compressed = zstd::encode_all(&json_bytes[..], 3)?; + let mut compressed = + zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?; + if compressed.len() > EDIT_PREDICTION_SETTLED_MAX_BODY_BYTES { + body.sample_data = None; + compressed = zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?; + } let url = client .http_client() @@ -2053,7 +2226,8 @@ impl EditPredictionStore { edited_buffer_snapshot: &BufferSnapshot, editable_offset_range: Range, edit_preview: &EditPreview, - example: Option, + context_task: Option>>, + prompt_history_boundary: Option, model_version: Option, e2e_latency: std::time::Duration, cx: &mut Context, @@ -2073,12 +2247,12 @@ impl EditPredictionStore { .current_organization() .map(|organization| organization.id.clone()); let project_state = this.get_or_init_project(project, cx); - let Some(registered_buffer) = project_state + if !project_state .registered_buffers - .get_mut(&edited_buffer.entity_id()) - else { + .contains_key(&edited_buffer.entity_id()) + { return; - }; + } let editable_region_before_prediction = edited_buffer_snapshot .text_for_range(editable_offset_range.clone()) @@ -2101,12 +2275,30 @@ impl EditPredictionStore { ), ); let editable_anchor_range = - edited_buffer_snapshot.anchor_range_inside(editable_offset_range); + edited_buffer_snapshot.anchor_range_inside(editable_offset_range.clone()); let now = cx.background_executor().now(); - registered_buffer - .pending_predictions - .push(PendingSettledPrediction { + let sample_data = if can_collect_data + && let Some(context_task) = context_task + && let Some(file) = edited_buffer_snapshot.file() + { + Some(PendingPredictionCaptureSampleData { + context_task, + editable_path: file.path().as_std_path().into(), + editable_offset_range, + next_edit_cursor_offset: None, + future_edit_history_events: Vec::new(), + navigation_history: VecDeque::new(), + edit_events_before_quiescence: 0, + prompt_history_boundary, + }) + } else { + None + }; + project_state + .pending_prediction_captures + .push(PendingPredictionCapture { request_id, + edited_buffer_id: edited_buffer.entity_id(), editable_anchor_range, editable_region_before_prediction, predicted_editable_region, @@ -2115,7 +2307,7 @@ impl EditPredictionStore { organization_id, can_collect_data, is_in_open_source_repo, - example, + sample_data, model_version, e2e_latency, enqueued_at: now, @@ -2749,6 +2941,18 @@ impl EditPredictionStore { self.get_or_init_project(&project, cx); let project_state = self.projects.get(&project.entity_id()).unwrap(); let stored_events = project_state.events(cx); + let prompt_history_boundary = Some(PromptHistoryBoundary { + first_event_seq: project_state + .last_event + .as_ref() + .map_or(project_state.next_last_event_seq, |last_event| { + last_event.seq + }), + snapshot: project_state + .last_event + .as_ref() + .map(|last_event| last_event.new_snapshot.clone()), + }); let has_events = !stored_events.is_empty(); let events: Vec> = stored_events.iter().map(|e| e.event.clone()).collect(); @@ -2768,15 +2972,28 @@ impl EditPredictionStore { }; let buffer_id = active_buffer.read(cx).remote_id(); - let repo_url = project + let (repository_url, revision) = project .read(cx) .git_store() .read(cx) .repository_and_path_for_buffer_id(buffer_id, cx) - .and_then(|(repo, _)| repo.read(cx).default_remote_url()); + .map(|(repository, _)| { + let snapshot = repository.read(cx).snapshot(); + ( + snapshot + .remote_origin_url + .clone() + .or_else(|| snapshot.remote_upstream_url.clone()), + snapshot + .head_commit + .as_ref() + .map(|commit| commit.sha.to_string()), + ) + }) + .unwrap_or_default(); let is_staff_zed_repo = cx.is_staff() - && repo_url + && repository_url .as_ref() .is_some_and(|url| is_zed_industries_repo(url)); let is_open_source = is_staff_zed_repo @@ -2790,14 +3007,12 @@ impl EditPredictionStore { && is_open_source && self.is_data_collection_enabled(cx) && matches!(self.edit_prediction_model, EditPredictionModel::Zeta); - let capture_worktree_id = snapshot.file().map(|file| file.worktree_id(cx)); let inputs = EditPredictionModelInput { project: project.clone(), buffer: active_buffer, snapshot, position, events, - stored_events: stored_events.clone(), related_files, mode, trigger, @@ -2809,41 +3024,27 @@ impl EditPredictionStore { let task = match self.edit_prediction_model { EditPredictionModel::Zeta => { - 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 - | PredictEditsRequestTrigger::DiagnosticNavigation => { - 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_data, repo_url, cx) + let context_task = can_collect_data + .then(|| { + capture_prediction_context( + inputs.project.clone(), + inputs.buffer.clone(), + inputs.position, + stored_events, + repository_url.clone(), + revision, + cx, + ) + }) + .flatten(); + zeta::request_prediction_with_zeta( + self, + inputs, + context_task, + prompt_history_boundary, + repository_url, + cx, + ) } EditPredictionModel::Fim { format } => fim::request_prediction(inputs, format, cx), EditPredictionModel::Mercury => { @@ -3440,7 +3641,7 @@ fn merge_trailing_events_if_needed( merge_anchor_ranges(&merged_edit_range, &event.total_edit_range, latest_snapshot); } - if let Some((diff, edit_range)) = compute_diff_between_snapshots_in_range( + if let Some((diff, old_range, new_range)) = compute_diff_between_snapshots_in_range( &oldest_snapshot, newest_snapshot, &merged_edit_range, @@ -3456,6 +3657,8 @@ fn merge_trailing_events_if_needed( old_path: old_path.clone(), path: path.clone(), diff, + old_range, + new_range: new_range.clone(), in_open_source_repo: *in_open_source_repo, predicted: events.range(merge_start..).all(|event| { matches!( @@ -3469,8 +3672,8 @@ fn merge_trailing_events_if_needed( }), old_snapshot: oldest_snapshot.clone(), new_snapshot_version: newest_snapshot.version.clone(), - total_edit_range: newest_snapshot.anchor_before(edit_range.start) - ..newest_snapshot.anchor_before(edit_range.end), + total_edit_range: newest_snapshot.anchor_before(new_range.start) + ..newest_snapshot.anchor_before(new_range.end), file_context: oldest_event.file_context.clone(), }, }; diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 3dc02f8eb5..1ef96d84af 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -3,7 +3,7 @@ use clock::FakeSystemClock; use clock::ReplicaId; use cloud_api_types::{ CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration, - OrganizationEditPredictionConfiguration, OrganizationId, SubmitEditPredictionSettledBody, + OrganizationEditPredictionConfiguration, OrganizationId, SettledEditPrediction, SubmitEditPredictionSettledResponse, }; use cloud_llm_client::{ @@ -1455,6 +1455,7 @@ async fn test_empty_prediction(cx: &mut TestAppContext) { let position = snapshot.anchor_before(language::Point::new(1, 3)); ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&buffer, &project, cx); ep_store.refresh_prediction_from_buffer( project.clone(), buffer.clone(), @@ -1495,7 +1496,7 @@ async fn test_empty_prediction(cx: &mut TestAppContext) { assert_eq!( &reject_request.rejections, &[EditPredictionRejection { - request_id: id, + request_id: id.clone(), reason: EditPredictionRejectReason::Empty, was_shown: false, model_version: Some("zeta2:test-empty".to_string()), @@ -2686,6 +2687,7 @@ fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) { let text = (0..300) .map(|row| format!("line {row} has some diagnostic context\n")) .collect::(); + let long_message = "diagnostic message ".repeat(1000); let buffer = cx.new(|cx| Buffer::local(&text, cx)); buffer.update(cx, |buffer, cx| { @@ -2695,7 +2697,7 @@ fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) { range: text::PointUtf16::new(150, 0)..text::PointUtf16::new(150, 4), diagnostic: Diagnostic { severity: DiagnosticSeverity::ERROR, - message: "long snippet".to_string(), + message: long_message.clone(), group_id: 1, is_primary: true, source_kind: language::DiagnosticSourceKind::Pushed, @@ -2716,7 +2718,15 @@ fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) { ); assert_eq!(active_buffer_diagnostics.len(), 1); - assert!(active_buffer_diagnostics[0].snippet.len() <= 512 * 3 + 2); + assert!( + active_buffer_diagnostics[0].message.len() + <= crate::zeta::MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT * 3 + 2 + ); + assert!(active_buffer_diagnostics[0].message.len() < long_message.len()); + assert!( + active_buffer_diagnostics[0].snippet.len() + <= crate::zeta::MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT * 3 + 2 + ); assert!(active_buffer_diagnostics[0].snippet.len() < text.len()); } @@ -2819,7 +2829,7 @@ struct RequestChannels { oneshot::Sender, )>, reject: mpsc::UnboundedReceiver<(RejectEditPredictionsBody, oneshot::Sender<()>)>, - settled: mpsc::UnboundedReceiver, + settled: mpsc::UnboundedReceiver, } fn init_test_with_fake_client( @@ -3833,6 +3843,7 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) { &edit_preview_a, None, None, + None, Duration::from_secs(0), cx, ); @@ -3900,6 +3911,7 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) { &edit_preview_b, None, None, + None, Duration::from_secs(0), cx, ); @@ -3949,19 +3961,40 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) { } } -#[gpui::test] -async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disabled( +const MIT_LICENSE: &str = indoc! {r#" + MIT License + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. +"#}; + +async fn init_sample_capture_test( + tree: serde_json::Value, cx: &mut TestAppContext, +) -> ( + Entity, + RequestChannels, + Entity, + Entity, ) { - let (ep_store, mut requests) = init_test_with_fake_client(cx); + let (ep_store, requests) = init_test_with_fake_client(cx); let fs = FakeFs::new(cx.executor()); - fs.insert_tree( - "/root", - json!({ - "foo.md": "sensitive source\n" - }), - ) - .await; + fs.insert_tree("/root", tree).await; let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; let buffer = project .update(cx, |project, cx| { @@ -3970,52 +4003,404 @@ async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disable }) .await .unwrap(); - ep_store.update(cx, |ep_store, cx| { ep_store.register_buffer(&buffer, &project, cx); }); + cx.run_until_parked(); + (ep_store, requests, project, buffer) +} +/// Enqueues a settled prediction capture with injected sample data, as the +/// real request path would when data collection is enabled. Returns the +/// editable offset range. +async fn enqueue_sample_capture( + ep_store: &Entity, + project: &Entity, + buffer: &Entity, + id: &str, + editable_point_range: Range, + context: CapturedPredictionContext, + prompt_history_boundary: Option, + navigation_history: VecDeque, + cx: &mut TestAppContext, +) -> Range { let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); - let edits: Arc<[(Range, Arc)]> = - cx.update(|cx| to_completion_edits([(0..9, "replacement".into())], &buffer, cx).into()); + let editable_offset_range = snapshot.point_to_offset(editable_point_range.start) + ..snapshot.point_to_offset(editable_point_range.end); + let empty_edits: Arc<[(Range, Arc)]> = Vec::new().into(); let edit_preview = buffer - .read_with(cx, |buffer, cx| buffer.preview_edits(edits, cx)) + .read_with(cx, |buffer, cx| buffer.preview_edits(empty_edits, cx)) .await; - ep_store.update(cx, |ep_store, cx| { ep_store.enqueue_settled_prediction( - EditPredictionId("prediction-private".into()), - &project, - &buffer, + EditPredictionId(id.to_string().into()), + project, + buffer, &snapshot, - 0..snapshot.len(), + editable_offset_range.clone(), &edit_preview, - Some(ExampleSpec { - name: "test example".to_string(), - repository_url: "https://example.com/repo".to_string(), - revision: "rev".to_string(), - tags: Vec::new(), - reasoning: None, - uncommitted_diff: String::new(), - recently_opened_files: Vec::new(), - recently_viewed_files: Vec::new(), - uncommitted_diff_contains_edit_history: false, - cursor_path: Path::new("foo.md").into(), - cursor_position: "0".to_string(), - edit_history: "sensitive edit history".to_string(), - expected_patches: vec!["sensitive patch".to_string()], - rejected_patch: None, - telemetry: None, - human_feedback: Vec::new(), - rating: None, - }), - Some("test-model".to_string()), - Duration::from_millis(42), + None, + None, + None, + Duration::from_secs(0), cx, ); + let pending_capture = ep_store + .projects + .get_mut(&project.entity_id()) + .unwrap() + .pending_prediction_captures + .last_mut() + .unwrap(); + pending_capture.can_collect_data = true; + pending_capture.sample_data = Some(PendingPredictionCaptureSampleData { + context_task: Task::ready(Ok(context)), + editable_path: snapshot.file().unwrap().path().as_std_path().into(), + editable_offset_range: editable_offset_range.clone(), + next_edit_cursor_offset: None, + future_edit_history_events: Vec::new(), + navigation_history, + edit_events_before_quiescence: 0, + prompt_history_boundary, + }); + }); + editable_offset_range +} + +#[gpui::test] +async fn test_edit_prediction_settled_sends_sample_data_after_quiescence(cx: &mut TestAppContext) { + let (ep_store, mut requests, project, buffer) = init_sample_capture_test( + json!({ + "LICENSE": MIT_LICENSE, + "foo.md": (0..60).map(|ix| format!("line {ix}\n")).collect::(), + }), + cx, + ) + .await; + + buffer.update(cx, |buffer, cx| { + let offset = Point::new(1, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "prompted ")], None, cx); + }); + cx.run_until_parked(); + + let boundary = ep_store.update(cx, |ep_store, _cx| { + let project_state = ep_store.projects.get(&project.entity_id()).unwrap(); + PromptHistoryBoundary { + first_event_seq: project_state + .last_event + .as_ref() + .map_or(project_state.next_last_event_seq, |last_event| { + last_event.seq + }), + snapshot: project_state + .last_event + .as_ref() + .map(|last_event| last_event.new_snapshot.clone()), + } + }); + let editable_offset_range = enqueue_sample_capture( + &ep_store, + &project, + &buffer, + "prediction-sample", + Point::new(2, 0)..Point::new(3, 0), + CapturedPredictionContext { + repository_url: Some("https://example.com/repo.git".to_string()), + revision: Some("abc123".to_string()), + uncommitted_diff: Some("--- a/foo.md\n+++ b/foo.md\n".to_string()), + buffer_diagnostics: vec![zeta_prompt::ActiveBufferDiagnostic { + severity: Some(1), + message: "sample diagnostic".to_string(), + snippet: String::new(), + snippet_buffer_row_range: 0..0, + diagnostic_range_in_snippet: 0..0, + }], + }, + Some(boundary), + VecDeque::from([RecentFile { + path: Path::new("foo.md").into(), + cursor_position: Some(3), + }]), + cx, + ) + .await; + let expected_next_edit_cursor_offset = editable_offset_range.start; + + // A second capture whose user has data collection disabled; its sample + // and settled region must be redacted at send time. + let boundary = ep_store.update(cx, |ep_store, _cx| { + let project_state = ep_store.projects.get(&project.entity_id()).unwrap(); + PromptHistoryBoundary { + first_event_seq: project_state + .last_event + .as_ref() + .map_or(project_state.next_last_event_seq, |last_event| { + last_event.seq + }), + snapshot: project_state + .last_event + .as_ref() + .map(|last_event| last_event.new_snapshot.clone()), + } + }); + enqueue_sample_capture( + &ep_store, + &project, + &buffer, + "prediction-redacted", + Point::new(2, 0)..Point::new(3, 0), + CapturedPredictionContext { + repository_url: None, + revision: None, + uncommitted_diff: None, + buffer_diagnostics: Vec::new(), + }, + Some(boundary), + VecDeque::new(), + cx, + ) + .await; + ep_store.update(cx, |ep_store, _cx| { + ep_store + .projects + .get_mut(&project.entity_id()) + .unwrap() + .pending_prediction_captures + .last_mut() + .unwrap() + .can_collect_data = false; }); + cx.executor().advance_clock(LAST_CHANGE_GROUPING_TIME * 2); cx.run_until_parked(); + + for (ix, row) in [2, 20, 30, 40, 50].into_iter().enumerate() { + buffer.update(cx, |buffer, cx| { + let start = Point::new(row, 0).to_offset(buffer); + let end = Point::new(row, 4).to_offset(buffer); + buffer.edit(vec![(start..end, format!("future {ix}"))], None, cx); + }); + cx.run_until_parked(); + } + + cx.executor() + .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE); + cx.run_until_parked(); + + let mut settled_by_id = std::collections::HashMap::new(); + for _ in 0..2 { + let request = requests + .settled + .next() + .await + .expect("settled request should be sent"); + settled_by_id.insert(request.request_id.clone(), request); + } + + let redacted_request = settled_by_id.remove("prediction-redacted").unwrap(); + assert!(!redacted_request.can_collect_data); + assert_eq!(redacted_request.settled_editable_region, None); + assert_eq!(redacted_request.sample_data, None); + + let settled_request = settled_by_id.remove("prediction-sample").unwrap(); + let sample_data = settled_request + .sample_data + .expect("sample data should be sent after quiescence"); + assert_eq!( + sample_data.repository_url.as_deref(), + Some("https://example.com/repo.git") + ); + assert_eq!(sample_data.revision.as_deref(), Some("abc123")); + assert_eq!( + sample_data.uncommitted_diff.as_deref(), + Some("--- a/foo.md\n+++ b/foo.md\n") + ); + assert_eq!(sample_data.editable_path.as_ref(), Path::new("foo.md")); + assert_eq!(sample_data.editable_offset_range, editable_offset_range); + assert_eq!(sample_data.buffer_diagnostics.len(), 1); + assert_eq!(sample_data.future_edit_history_events.len(), 4); + assert_eq!(sample_data.navigation_history.len(), 1); + assert_eq!(sample_data.edit_events_before_quiescence, 5); + assert_eq!( + sample_data.next_edit_cursor_offset, + Some(expected_next_edit_cursor_offset) + ); + + let future_event_diffs = sample_data + .future_edit_history_events + .iter() + .map(|event| match event.as_ref() { + zeta_prompt::Event::BufferChange { diff, .. } => diff.as_str(), + }) + .collect::>(); + assert!(future_event_diffs.iter().all(|diff| { + diff.lines() + .all(|line| !line.starts_with("+prompted ") && line != "-line 1") + })); + assert!( + future_event_diffs + .iter() + .any(|diff| diff.contains("future 0")) + ); + assert!( + !future_event_diffs + .iter() + .any(|diff| diff.contains("future 4")) + ); +} + +#[gpui::test] +async fn test_edit_prediction_settled_sample_data_requires_observing_all_events_since_request( + cx: &mut TestAppContext, +) { + let (ep_store, mut requests, project, buffer) = init_sample_capture_test( + json!({ + "LICENSE": MIT_LICENSE, + "foo.md": (0..30).map(|ix| format!("line {ix}\n")).collect::(), + }), + cx, + ) + .await; + + // Two predictions are requested while no event is pending. + let (boundary_observed, boundary_missed) = ep_store.update(cx, |ep_store, _cx| { + let project_state = ep_store.projects.get(&project.entity_id()).unwrap(); + let first_event_seq = project_state + .last_event + .as_ref() + .map_or(project_state.next_last_event_seq, |last_event| { + last_event.seq + }); + let snapshot = project_state + .last_event + .as_ref() + .map(|last_event| last_event.new_snapshot.clone()); + ( + PromptHistoryBoundary { + first_event_seq, + snapshot: snapshot.clone(), + }, + PromptHistoryBoundary { + first_event_seq, + snapshot, + }, + ) + }); + assert!(boundary_observed.snapshot.is_none()); + + // The first capture is enqueued immediately, so it observes both + // subsequent events. + enqueue_sample_capture( + &ep_store, + &project, + &buffer, + "prediction-observed", + Point::new(1, 0)..Point::new(2, 0), + CapturedPredictionContext { + repository_url: None, + revision: None, + uncommitted_diff: None, + buffer_diagnostics: Vec::new(), + }, + Some(boundary_observed), + VecDeque::new(), + cx, + ) + .await; + + buffer.update(cx, |buffer, cx| { + let offset = Point::new(1, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "first ")], None, cx); + }); + cx.run_until_parked(); + buffer.update(cx, |buffer, cx| { + let offset = Point::new(20, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "second ")], None, cx); + }); + cx.run_until_parked(); + + // The second capture is enqueued only after the first event was + // finalized, so its future history has a gap and it must be dropped. + enqueue_sample_capture( + &ep_store, + &project, + &buffer, + "prediction-missed", + Point::new(1, 0)..Point::new(2, 0), + CapturedPredictionContext { + repository_url: None, + revision: None, + uncommitted_diff: None, + buffer_diagnostics: Vec::new(), + }, + Some(boundary_missed), + VecDeque::new(), + cx, + ) + .await; + + cx.executor() + .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE); + cx.run_until_parked(); + + let mut settled_by_id = std::collections::HashMap::new(); + for _ in 0..2 { + let request = requests + .settled + .next() + .await + .expect("settled request should be sent"); + settled_by_id.insert(request.request_id.clone(), request); + } + + let observed_request = settled_by_id.remove("prediction-observed").unwrap(); + let sample_data = observed_request + .sample_data + .expect("sample data should be sent"); + assert_eq!(sample_data.edit_events_before_quiescence, 2); + assert_eq!(sample_data.future_edit_history_events.len(), 2); + + let missed_request = settled_by_id.remove("prediction-missed").unwrap(); + assert_eq!(missed_request.sample_data, None); +} + +#[gpui::test] +async fn test_edit_prediction_settled_drops_future_events_when_their_oss_status_is_unknown( + cx: &mut TestAppContext, +) { + let (ep_store, mut requests, project, buffer) = init_sample_capture_test( + json!({ "foo.md": (0..30).map(|ix| format!("line {ix}\n")).collect::() }), + cx, + ) + .await; + + enqueue_sample_capture( + &ep_store, + &project, + &buffer, + "prediction-non-oss", + Point::new(1, 0)..Point::new(2, 0), + CapturedPredictionContext { + repository_url: None, + revision: None, + uncommitted_diff: None, + buffer_diagnostics: Vec::new(), + }, + None, + VecDeque::new(), + cx, + ) + .await; + + // There is no LICENSE file, so this event's open-source status is unknown + // and the sample must be dropped. + buffer.update(cx, |buffer, cx| { + let offset = Point::new(20, 0).to_offset(buffer); + buffer.edit(vec![(offset..offset, "outside ")], None, cx); + }); + cx.run_until_parked(); + cx.executor() .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE); cx.run_until_parked(); @@ -4025,9 +4410,7 @@ async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disable .next() .await .expect("settled request should be sent"); - assert!(!settled_request.can_collect_data); - assert_eq!(settled_request.settled_editable_region, None); - assert_eq!(settled_request.example, None); + assert_eq!(settled_request.sample_data, None); } #[gpui::test] diff --git a/crates/edit_prediction/src/jump_example.rs b/crates/edit_prediction/src/jump_example.rs deleted file mode 100644 index 33991df9c3..0000000000 --- a/crates/edit_prediction/src/jump_example.rs +++ /dev/null @@ -1,372 +0,0 @@ -use std::{ - ops::Range, - sync::Arc, - time::{Duration, Instant}, -}; - -use anyhow::{Context as _, Result}; -pub use cloud_api_types::JumpExampleTrigger; -use cloud_api_types::{ - JumpExampleRecentFile, SubmitEditPredictionJumpExampleBody, - SubmitEditPredictionJumpExampleResponse, -}; -use futures::future::Shared; -use gpui::{AppContext as _, AsyncApp, Context, Entity, Task, TaskExt as _, WeakEntity}; -use language::{BufferSnapshot, File, Point}; -use project::{Project, WorktreeId}; -use release_channel::AppVersion; - -use text::ToPoint as _; -use util::rel_path::RelPath; - -use crate::{ - EditPredictionStore, ProjectState, StoredEvent, - data_collection::{ - UncommittedDiffResult, compute_cursor_excerpt, compute_uncommitted_diff, - estimate_uncommitted_diff_byte_size, format_cursor_excerpt, - }, - example_spec::RecentFile, - zeta, -}; - -pub const JUMP_EXAMPLE_MAX_PENDING_CAPTURE_COUNT: usize = 10; -pub const JUMP_EXAMPLE_FUTURE_EVENT_COUNT: usize = 2; -pub const JUMP_EXAMPLE_TTL: Duration = Duration::from_secs(60 * 2); -pub const JUMP_EXAMPLE_NAVIGATION_COUNT: usize = 20; -pub const JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024; - -pub struct PendingJumpExampleCapture { - key: PendingJumpExampleCaptureKey, - trigger: JumpExampleTrigger, - file: Arc, - edit_history: Vec>, - recently_opened_files: Vec, - recently_viewed_files: Vec, - worktree_root_name: String, - cursor_position: String, - started_at: Instant, - uncommitted_diff: Option, - pub future_events: Vec>, - pub navigation_history: Vec, - diagnostics: Vec, - repository_url: Option, - revision: Option, - can_collect_data: bool, - is_in_open_source_repo: bool, -} - -#[derive(Eq, PartialEq, Hash, Clone)] -pub struct PendingJumpExampleCaptureKey { - worktree_id: WorktreeId, - file_path: Arc, - row_bucket: u32, -} - -pub fn try_start_jump_example_capture( - project_state: &ProjectState, - uncommitted_diffs: Shared>, - project: Entity, - snapshot: BufferSnapshot, - position: language::Anchor, - trigger: JumpExampleTrigger, - stored_events: Vec, - diagnostic_search_range: Range, - can_collect_data: bool, - is_in_open_source_repo: bool, - cx: &mut Context, -) { - let Some(file) = snapshot.file().cloned() else { - return; - }; - - let example_key = PendingJumpExampleCaptureKey { - worktree_id: file.worktree_id(cx), - file_path: file.path().clone(), - row_bucket: position.to_point(&snapshot).row / 10, - }; - let should_capture_example = project_state.pending_jump_example_captures.len() - < JUMP_EXAMPLE_MAX_PENDING_CAPTURE_COUNT - && !project_state - .starting_jump_example_captures - .contains(&example_key) - && !project_state - .pending_jump_example_captures - .iter() - .any(|capture| &capture.key == &example_key); - - if !should_capture_example { - return; - } - - let _project = project.clone(); - let _example_key = example_key.clone(); - let task = cx.spawn(async move |ep_store, cx| { - let project = _project; - let example_key = _example_key; - let Some(ep_store) = ep_store.upgrade() else { - return anyhow::Ok(()); - }; - ep_store.update(cx, |ep_store, cx| { - let project_state = ep_store.get_or_init_project(&project, cx); - project_state - .starting_jump_example_captures - .push(example_key.clone()); - }); - - let (repository, worktree) = project.read_with(cx, |project, cx| { - let repository = project.active_repository(cx); - let worktree_id = file.worktree_id(cx); - let worktree = project.worktree_for_id(worktree_id, cx); - (repository, worktree) - }); - let Some(worktree) = worktree else { - return Ok(()); - }; - - let diagnostics = zeta::active_buffer_diagnostics( - &snapshot, - diagnostic_search_range.clone(), - position.to_point(&snapshot).row, - 100, - ); - - let uncommitted_diff = 'uncommitted_diff: { - if repository.is_none() { - break 'uncommitted_diff None; - } - let uncommitted_diff_snapshot = uncommitted_diffs - .await - .map_err(|error| anyhow::anyhow!("{error:?}")) - .context("failed to capture uncommitted diff")?; - let estimated_byte_size = - estimate_uncommitted_diff_byte_size(&uncommitted_diff_snapshot); - if estimated_byte_size > JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE { - break 'uncommitted_diff None; - } - - let uncommitted_diff = cx - .background_executor() - .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) }) - .await; - if uncommitted_diff.len() > JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE { - break 'uncommitted_diff None; - } - Some(uncommitted_diff) - }; - - let edit_history = stored_events - .iter() - .map(|e| e.event.clone()) - .collect::>(); - let (repository_url, revision) = if let Some(repository) = &repository { - repository.read_with(cx, |repository, _| { - let snapshot = repository.snapshot(); - ( - snapshot - .remote_origin_url - .clone() - .or_else(|| snapshot.remote_upstream_url.clone()), - snapshot - .head_commit - .as_ref() - .map(|commit| commit.sha.to_string()), - ) - }) - } else { - (None, None) - }; - let line_comment_prefix = snapshot - .language() - .and_then(|language| language.config().line_comments.first()) - .map(|prefix| prefix.to_string()) - .unwrap_or_default(); - let (cursor_excerpt, cursor_offset_in_excerpt, _) = cx - .background_executor() - .spawn(async move { compute_cursor_excerpt(&snapshot, position) }) - .await; - let cursor_position = format_cursor_excerpt( - &cursor_excerpt, - cursor_offset_in_excerpt, - &line_comment_prefix, - ); - let now = cx.background_executor().now(); - ep_store.update(cx, |ep_store, cx| { - let recently_opened_files = ep_store.recently_opened_files_for_project(&project); - let recently_viewed_files = ep_store.recently_viewed_files_for_project(&project); - let project_state = ep_store.get_or_init_project(&project, cx); - project_state - .pending_jump_example_captures - .push(PendingJumpExampleCapture { - key: example_key, - trigger, - file, - uncommitted_diff, - edit_history, - recently_opened_files, - recently_viewed_files, - repository_url, - revision, - diagnostics, - worktree_root_name: worktree.read(cx).root_name_str().to_owned(), - cursor_position, - started_at: now, - future_events: Vec::new(), - navigation_history: Vec::new(), - is_in_open_source_repo, - can_collect_data, - }); - drain_completed_jump_example_captures(project_state, cx); - }); - Ok(()) - }); - cx.spawn(async move |ep_store, cx| { - let result = task.await; - ep_store - .update(cx, |ep_store, cx| { - ep_store - .get_or_init_project(&project, cx) - .starting_jump_example_captures - .retain(|key| key != &example_key); - }) - .ok(); - result - }) - .detach_and_log_err(cx); -} - -pub fn drain_completed_jump_example_captures( - project_state: &mut ProjectState, - cx: &mut Context, -) { - let now = cx.background_executor().now(); - - let mut capture_index = 0; - while capture_index < project_state.pending_jump_example_captures.len() { - let capture = &project_state.pending_jump_example_captures[capture_index]; - let finished = capture.future_events.len() >= JUMP_EXAMPLE_FUTURE_EVENT_COUNT - || now.saturating_duration_since(capture.started_at) >= JUMP_EXAMPLE_TTL; - if !finished { - capture_index += 1; - continue; - } - - let capture = project_state - .pending_jump_example_captures - .remove(capture_index); - cx.spawn(async move |this, cx| { - let result = submit_jump_example_capture_task(this, capture, cx).await; - if let Err(error) = result { - log::error!("failed to submit jump opportunity capture: {error:?}"); - } - }) - .detach(); - } -} - -fn submit_jump_example_capture_task( - this: WeakEntity, - capture: PendingJumpExampleCapture, - cx: &mut AsyncApp, -) -> Task> { - let Some((organization_id, client, llm_token, app_version)) = this - .update(cx, |this, cx| { - ( - this.user_store - .read(cx) - .current_organization() - .map(|organization| organization.id.clone()), - this.client.clone(), - this.llm_token.clone(), - AppVersion::global(cx), - ) - }) - .ok() - else { - return Task::ready(Ok(())); - }; - cx.background_spawn(async move { - let PendingJumpExampleCapture { - key: _, - trigger, - file, - edit_history, - recently_opened_files, - recently_viewed_files, - worktree_root_name, - cursor_position, - started_at: _, - uncommitted_diff, - future_events, - navigation_history, - diagnostics, - repository_url, - revision, - is_in_open_source_repo, - can_collect_data, - } = capture; - let future_edit_history = render_jump_example_events(&future_events, &worktree_root_name); - - let cursor_path = file.path().as_std_path().into(); - let example = SubmitEditPredictionJumpExampleBody { - request_id: uuid::Uuid::new_v4(), - trigger, - repository_url, - revision, - uncommitted_diff, - recently_opened_files: jump_example_recent_files(recently_opened_files), - recently_viewed_files: jump_example_recent_files(recently_viewed_files), - cursor_path, - cursor_position, - edit_history, - diagnostics, - future_edit_history, - navigation_history: jump_example_recent_files(navigation_history), - is_in_open_source_repo, - can_collect_data, - }; - let json_bytes = serde_json::to_vec(&example)?; - let compressed = zstd::encode_all(&json_bytes[..], 3)?; - let url = client - .http_client() - .build_zed_llm_url("/predict_edits/jump_example", &[])?; - EditPredictionStore::send_api_request::( - |builder| { - Ok(builder - .uri(url.as_ref()) - .header("Content-Encoding", "zstd") - .body(compressed.clone().into())?) - }, - client, - llm_token, - organization_id, - app_version, - ) - .await?; - Ok(()) - }) -} - -fn jump_example_recent_files(files: Vec) -> Vec { - files - .into_iter() - .map(|file| JumpExampleRecentFile { - path: file.path, - cursor_position: file.cursor_position, - }) - .collect() -} - -fn render_jump_example_events(events: &[Arc], root_name: &str) -> String { - let mut edit_history = String::new(); - for event in events { - crate::capture_example::write_event_with_relative_paths( - &mut edit_history, - event, - root_name, - ); - if !edit_history.ends_with('\n') { - edit_history.push('\n'); - } - } - edit_history -} diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index 2631d0fd22..8b388ef180 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -1,15 +1,15 @@ use crate::{ CloudRequestTimeoutError, CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent, - EditPredictionStore, ZedUpdateRequiredError, buffer_path_with_id_fallback, + EditPredictionStore, PromptHistoryBoundary, ZedUpdateRequiredError, + buffer_path_with_id_fallback, + capture_prediction_context::CapturedPredictionContext, cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, - data_collection::UncommittedDiffResult, prediction::EditPredictionResult, }; -use anyhow::{Context as _, Result}; +use anyhow::Result; use cloud_llm_client::{AcceptEditPredictionBody, 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 _, @@ -33,7 +33,7 @@ use crate::open_ai_compatible::{ load_open_ai_compatible_api_key_if_needed, send_custom_server_request, }; -pub fn request_prediction_with_zeta( +pub(crate) fn request_prediction_with_zeta( store: &mut EditPredictionStore, EditPredictionModelInput { buffer, @@ -41,7 +41,6 @@ pub fn request_prediction_with_zeta( position, related_files, events, - stored_events, debug_tx, mode, trigger, @@ -51,7 +50,8 @@ pub fn request_prediction_with_zeta( is_open_source, .. }: EditPredictionModelInput, - capture_data: Option>>, + context_task: Option>>, + prompt_history_boundary: Option, repo_url: Option, cx: &mut Context, ) -> Task>> { @@ -413,49 +413,7 @@ 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(|uncommitted_diffs| { - let (recently_opened_files, recently_viewed_files) = this - .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 - .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, - stored_events, - recently_opened_files, - recently_viewed_files, - uncommitted_diffs, - false, - cx, - ) - }) else { - return Err(anyhow::anyhow!("failed to capture example")); - }; - task.await - } - })) - }); cx.spawn(async move |cx| { - let example_spec = if let Some(task) = example_task { - task.await.ok() - } else { - None - }; - weak_this .update(cx, |this, cx| { this.enqueue_settled_prediction( @@ -465,7 +423,8 @@ pub fn request_prediction_with_zeta( &edited_buffer_snapshot, editable_range_in_buffer, &edit_preview, - example_spec, + context_task, + prompt_history_boundary, model_version, request_duration, cx, @@ -549,8 +508,8 @@ fn handle_api_response( const ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT: usize = 100; const MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT: usize = 20; -const MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT: usize = 512; -const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512; +pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT: usize = 512; +pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512; pub(crate) fn active_buffer_diagnostics( snapshot: &language::BufferSnapshot, diff --git a/crates/edit_prediction_cli/src/qa.rs b/crates/edit_prediction_cli/src/qa.rs index 81a2a77b27..d47c8b0526 100644 --- a/crates/edit_prediction_cli/src/qa.rs +++ b/crates/edit_prediction_cli/src/qa.rs @@ -87,8 +87,7 @@ pub fn build_prompt(example: &Example) -> Result { path, old_path, diff, - predicted: _, - in_open_source_repo: _, + .. } = event.as_ref(); edit_history.push_str(&format!("--- a{}\n", old_path.display())); edit_history.push_str(&format!("+++ b{}\n", path.display())); diff --git a/crates/edit_prediction_metrics/src/reversal.rs b/crates/edit_prediction_metrics/src/reversal.rs index fb84f57aa4..05b82f0c9c 100644 --- a/crates/edit_prediction_metrics/src/reversal.rs +++ b/crates/edit_prediction_metrics/src/reversal.rs @@ -1181,6 +1181,8 @@ mod tests { -old +new"} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: true, }), @@ -1192,6 +1194,8 @@ mod tests { -a +b"} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: true, }), @@ -1203,6 +1207,8 @@ mod tests { -x +y"} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: true, }), @@ -1931,6 +1937,8 @@ mod tests { line2 "} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: false, })], @@ -1969,6 +1977,8 @@ mod tests { line11 "} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: false, })], @@ -2029,6 +2039,8 @@ mod tests { line2 "} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: false, })], @@ -2066,6 +2078,8 @@ mod tests { more_wrong "} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: false, })], @@ -2134,6 +2148,8 @@ mod tests { line2 "} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: false, }), @@ -2147,6 +2163,8 @@ mod tests { line2 "} .into(), + old_range: 0..1, + new_range: 0..1, predicted: false, in_open_source_repo: false, }), diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 026722e4bf..b610e193c7 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -167,6 +167,10 @@ impl ZetaFormat { } } +fn empty_range() -> Range { + 0..0 +} + #[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)] #[serde(tag = "event")] pub enum Event { @@ -174,6 +178,10 @@ pub enum Event { path: Arc, old_path: Arc, diff: String, + #[serde(default = "empty_range")] + old_range: Range, + #[serde(default = "empty_range")] + new_range: Range, predicted: bool, in_open_source_repo: bool, }, @@ -203,7 +211,7 @@ pub fn write_event(prompt: &mut String, event: &Event) { old_path, diff, predicted, - in_open_source_repo: _, + .. } => { if *predicted { prompt.push_str("// User accepted prediction:\n"); @@ -5071,6 +5079,8 @@ mod tests { path: Path::new(path).into(), old_path: Path::new(path).into(), diff: diff.to_string(), + old_range: 0..diff.len(), + new_range: 0..diff.len(), predicted: false, in_open_source_repo: false, }