diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index 63240ddd53108f0b2450386150958e23f975d7ed..2347a731cb5b5f3590dafcf0a57dc0bab88c380c 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -75,6 +75,7 @@ pub mod zeta; #[cfg(test)] mod edit_prediction_tests; +use crate::cursor_excerpt::expand_context_syntactically_then_linewise; use crate::example_spec::ExampleSpec; use crate::license_detection::LicenseDetectionWatcher; use crate::mercury::Mercury; @@ -99,8 +100,9 @@ actions!( ); /// Maximum number of events to track. -const EVENT_COUNT_MAX: usize = 6; +const EVENT_COUNT_MAX: usize = 10; const CHANGE_GROUPING_LINE_SPAN: u32 = 8; +const COLLABORATOR_EDIT_LOCALITY_CONTEXT_TOKENS: usize = 512; const LAST_CHANGE_GROUPING_TIME: Duration = Duration::from_secs(1); const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice"; const REJECT_REQUEST_DEBOUNCE: Duration = Duration::from_secs(15); @@ -242,21 +244,31 @@ pub enum UserActionType { pub struct StoredEvent { pub event: Arc, pub old_snapshot: TextBufferSnapshot, - pub edit_range: Range, + pub new_snapshot_version: clock::Global, + pub total_edit_range: Range, } impl StoredEvent { fn can_merge( &self, - next_old_event: &&&StoredEvent, - new_snapshot: &TextBufferSnapshot, - last_edit_range: &Range, + next_old_event: &StoredEvent, + latest_snapshot: &TextBufferSnapshot, + latest_edit_range: &Range, ) -> bool { - // Events must be for the same buffer + // Events must be for the same buffer and be contiguous across included snapshots to be mergeable. if self.old_snapshot.remote_id() != next_old_event.old_snapshot.remote_id() { return false; } - if self.old_snapshot.remote_id() != new_snapshot.remote_id() { + if self.old_snapshot.remote_id() != latest_snapshot.remote_id() { + return false; + } + if self.new_snapshot_version != next_old_event.old_snapshot.version { + return false; + } + if !latest_snapshot + .version + .observed_all(&next_old_event.new_snapshot_version) + { return false; } @@ -281,9 +293,9 @@ impl StoredEvent { return false; } - let left_range = self.edit_range.to_point(new_snapshot); - let right_range = next_old_event.edit_range.to_point(new_snapshot); - let latest_range = last_edit_range.to_point(&new_snapshot); + let left_range = self.total_edit_range.to_point(latest_snapshot); + let right_range = next_old_event.total_edit_range.to_point(latest_snapshot); + let latest_range = latest_edit_range.to_point(latest_snapshot); // Events near to the latest edit are not merged if their sources differ. if lines_between_ranges(&left_range, &latest_range) @@ -516,7 +528,9 @@ struct LastEvent { new_snapshot: TextBufferSnapshot, old_file: Option>, new_file: Option>, - edit_range: Option>, + latest_edit_range: Range, + total_edit_range: Range, + total_edit_range_at_last_pause_boundary: Option>, predicted: bool, snapshot_after_last_editing_pause: Option, last_edit_time: Option, @@ -542,8 +556,11 @@ impl LastEvent { }) }); - let (diff, edit_range) = - compute_diff_between_snapshots(&self.old_snapshot, &self.new_snapshot)?; + let (diff, edit_range) = compute_diff_between_snapshots_in_range( + &self.old_snapshot, + &self.new_snapshot, + &self.total_edit_range, + )?; if path == old_path && diff.is_empty() { None @@ -556,9 +573,10 @@ impl LastEvent { in_open_source_repo, predicted: self.predicted, }), - edit_range: self.new_snapshot.anchor_before(edit_range.start) - ..self.new_snapshot.anchor_before(edit_range.end), 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), }) } } @@ -568,12 +586,28 @@ impl LastEvent { 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(), - edit_range: None, + latest_edit_range: latest_edit_range_before_pause, + 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, @@ -584,7 +618,9 @@ impl LastEvent { new_snapshot: self.new_snapshot.clone(), old_file: self.old_file.clone(), new_file: self.new_file.clone(), - edit_range: None, + 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, @@ -594,21 +630,78 @@ impl LastEvent { } } -pub(crate) fn compute_diff_between_snapshots( +fn compute_total_edit_range_between_snapshots( old_snapshot: &TextBufferSnapshot, new_snapshot: &TextBufferSnapshot, -) -> Option<(String, Range)> { +) -> Option> { let edits: Vec> = new_snapshot .edits_since::(&old_snapshot.version) .collect(); let (first_edit, last_edit) = edits.first().zip(edits.last())?; - - let old_start_point = old_snapshot.offset_to_point(first_edit.old.start); - let old_end_point = old_snapshot.offset_to_point(last_edit.old.end); let new_start_point = new_snapshot.offset_to_point(first_edit.new.start); let new_end_point = new_snapshot.offset_to_point(last_edit.new.end); + Some(new_snapshot.anchor_before(new_start_point)..new_snapshot.anchor_before(new_end_point)) +} + +fn compute_old_range_for_new_range( + old_snapshot: &TextBufferSnapshot, + new_snapshot: &TextBufferSnapshot, + total_edit_range: &Range, +) -> Option> { + 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 edits: Vec> = new_snapshot + .edits_since::(&old_snapshot.version) + .collect(); + let mut old_start_offset = None; + let mut old_end_offset = None; + let mut delta: isize = 0; + + for edit in &edits { + if old_start_offset.is_none() && new_start_offset <= edit.new.end { + old_start_offset = Some(if new_start_offset < edit.new.start { + new_start_offset.checked_add_signed(-delta)? + } else { + edit.old.start + }); + } + + if old_end_offset.is_none() && new_end_offset <= edit.new.end { + old_end_offset = Some(if new_end_offset < edit.new.start { + new_end_offset.checked_add_signed(-delta)? + } else { + edit.old.end + }); + } + + delta += edit.new.len() as isize - edit.old.len() as isize; + } + + let old_start_offset = + old_start_offset.unwrap_or_else(|| new_start_offset.saturating_add_signed(-delta)); + let old_end_offset = + old_end_offset.unwrap_or_else(|| new_end_offset.saturating_add_signed(-delta)); + + Some( + old_snapshot.offset_to_point(old_start_offset) + ..old_snapshot.offset_to_point(old_end_offset), + ) +} + +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); + 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; + const CONTEXT_LINES: u32 = 3; let old_context_start_row = old_start_point.row.saturating_sub(CONTEXT_LINES); @@ -1198,10 +1291,12 @@ impl EditPredictionStore { cx.subscribe(buffer, { let project = project.downgrade(); move |this, buffer, event, cx| { - if let language::BufferEvent::Edited { .. } = event + if let language::BufferEvent::Edited { is_local } = event && let Some(project) = project.upgrade() { - this.report_changes_for_buffer(&buffer, &project, false, cx); + this.report_changes_for_buffer( + &buffer, &project, false, *is_local, cx, + ); } } }), @@ -1223,6 +1318,7 @@ impl EditPredictionStore { buffer: &Entity, project: &Entity, is_predicted: bool, + is_local: bool, cx: &mut Context, ) { let project_state = self.get_or_init_project(project, cx); @@ -1234,7 +1330,6 @@ impl EditPredictionStore { if new_snapshot.version == registered_buffer.snapshot.version { return; } - let old_file = mem::replace(&mut registered_buffer.file, new_file.clone()); let old_snapshot = mem::replace(&mut registered_buffer.snapshot, new_snapshot.clone()); let mut num_edits = 0usize; @@ -1267,28 +1362,44 @@ impl EditPredictionStore { } } - let action_type = match (total_deleted, total_inserted, num_edits) { - (0, ins, n) if ins == n => UserActionType::InsertChar, - (0, _, _) => UserActionType::InsertSelection, - (del, 0, n) if del == n => UserActionType::DeleteChar, - (_, 0, _) => UserActionType::DeleteSelection, - (_, ins, n) if ins == n => UserActionType::InsertChar, - (_, _, _) => UserActionType::InsertSelection, - }; + let include_in_history = is_local + || collaborator_edit_overlaps_locality_region( + project_state, + project, + buffer, + &buf.snapshot(), + &edit_range, + cx, + ); - if let Some(offset) = last_offset { - let point = new_snapshot.offset_to_point(offset); - let timestamp_epoch_ms = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); - project_state.record_user_action(UserActionRecord { - action_type, - buffer_id: buffer.entity_id(), - line_number: point.row, - offset, - timestamp_epoch_ms, - }); + if is_local { + let action_type = match (total_deleted, total_inserted, num_edits) { + (0, ins, n) if ins == n => UserActionType::InsertChar, + (0, _, _) => UserActionType::InsertSelection, + (del, 0, n) if del == n => UserActionType::DeleteChar, + (_, 0, _) => UserActionType::DeleteSelection, + (_, ins, n) if ins == n => UserActionType::InsertChar, + (_, _, _) => UserActionType::InsertSelection, + }; + + if let Some(offset) = last_offset { + let point = new_snapshot.offset_to_point(offset); + let timestamp_epoch_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + project_state.record_user_action(UserActionRecord { + action_type, + buffer_id: buffer.entity_id(), + line_number: point.row, + offset, + timestamp_epoch_ms, + }); + } + } + + if !include_in_history { + return; } let events = &mut project_state.events; @@ -1302,15 +1413,10 @@ impl EditPredictionStore { let should_coalesce = is_next_snapshot_of_same_buffer && !prediction_source_changed - && last_event - .edit_range - .as_ref() - .is_some_and(|last_edit_range| { - lines_between_ranges( - &edit_range.to_point(&new_snapshot), - &last_edit_range.to_point(&new_snapshot), - ) <= CHANGE_GROUPING_LINE_SPAN - }); + && lines_between_ranges( + &edit_range.to_point(&new_snapshot), + &last_event.latest_edit_range.to_point(&new_snapshot), + ) <= CHANGE_GROUPING_LINE_SPAN; if should_coalesce { let pause_elapsed = last_event @@ -1320,9 +1426,13 @@ impl EditPredictionStore { if pause_elapsed { last_event.snapshot_after_last_editing_pause = Some(last_event.new_snapshot.clone()); + last_event.total_edit_range_at_last_pause_boundary = + Some(last_event.total_edit_range.clone()); } - last_event.edit_range = Some(edit_range); + last_event.latest_edit_range = edit_range.clone(); + last_event.total_edit_range = + merge_anchor_ranges(&last_event.total_edit_range, &edit_range, &new_snapshot); last_event.new_snapshot = new_snapshot; last_event.last_edit_time = Some(now); return; @@ -1345,7 +1455,9 @@ impl EditPredictionStore { new_file, old_snapshot, new_snapshot, - edit_range: Some(edit_range), + latest_edit_range: edit_range.clone(), + total_edit_range: edit_range, + total_edit_range_at_last_pause_boundary: None, predicted: is_predicted, snapshot_after_last_editing_pause: None, last_edit_time: Some(now), @@ -1401,7 +1513,13 @@ impl EditPredictionStore { return; }; - self.report_changes_for_buffer(¤t_prediction.prediction.buffer, project, true, cx); + self.report_changes_for_buffer( + ¤t_prediction.prediction.buffer, + project, + true, + true, + cx, + ); // can't hold &mut project_state ref across report_changes_for_buffer_call let Some(project_state) = self.projects.get_mut(&project.entity_id()) else { @@ -2670,6 +2788,32 @@ impl EditPredictionStore { } } +fn collaborator_edit_overlaps_locality_region( + project_state: &ProjectState, + project: &Entity, + buffer: &Entity, + snapshot: &BufferSnapshot, + edit_range: &Range, + cx: &App, +) -> bool { + let Some((active_buffer, Some(position))) = project_state.active_buffer(project, cx) else { + return false; + }; + + if active_buffer.entity_id() != buffer.entity_id() { + return false; + } + + let locality_point_range = expand_context_syntactically_then_linewise( + snapshot, + (position..position).to_point(snapshot), + COLLABORATOR_EDIT_LOCALITY_CONTEXT_TOKENS, + ); + let locality_anchor_range = snapshot.anchor_range_around(locality_point_range); + + edit_range.overlaps(&locality_anchor_range, snapshot) +} + fn merge_trailing_events_if_needed( events: &mut VecDeque, end_snapshot: &TextBufferSnapshot, @@ -2680,13 +2824,19 @@ fn merge_trailing_events_if_needed( if last_event.old_snapshot.remote_id() != latest_snapshot.remote_id() { return; } + if !latest_snapshot + .version + .observed_all(&last_event.new_snapshot_version) + { + return; + } } let mut next_old_event = None; let mut mergeable_count = 0; for old_event in events.iter().rev() { - if let Some(next_old_event) = &next_old_event - && !old_event.can_merge(&next_old_event, latest_snapshot, latest_edit_range) + if let Some(next_old_event) = next_old_event + && !old_event.can_merge(next_old_event, latest_snapshot, latest_edit_range) { break; } @@ -2701,10 +2851,19 @@ fn merge_trailing_events_if_needed( let mut events_to_merge = events.range(events.len() - mergeable_count..).peekable(); let oldest_event = events_to_merge.peek().unwrap(); let oldest_snapshot = oldest_event.old_snapshot.clone(); + let newest_snapshot = end_snapshot; + let mut merged_edit_range = oldest_event.total_edit_range.clone(); - if let Some((diff, edited_range)) = - compute_diff_between_snapshots(&oldest_snapshot, end_snapshot) - { + for event in events.range(events.len() - mergeable_count + 1..) { + merged_edit_range = + merge_anchor_ranges(&merged_edit_range, &event.total_edit_range, latest_snapshot); + } + + if let Some((diff, edit_range)) = compute_diff_between_snapshots_in_range( + &oldest_snapshot, + newest_snapshot, + &merged_edit_range, + ) { let merged_event = match oldest_event.event.as_ref() { zeta_prompt::Event::BufferChange { old_path, @@ -2728,8 +2887,9 @@ fn merge_trailing_events_if_needed( }), }), old_snapshot: oldest_snapshot.clone(), - edit_range: end_snapshot.anchor_before(edited_range.start) - ..end_snapshot.anchor_before(edited_range.end), + new_snapshot_version: newest_snapshot.version.clone(), + total_edit_range: newest_snapshot.anchor_before(edit_range.start) + ..newest_snapshot.anchor_before(edit_range.end), }, }; events.truncate(events.len() - mergeable_count); @@ -2737,6 +2897,24 @@ fn merge_trailing_events_if_needed( } } +fn merge_anchor_ranges( + left: &Range, + right: &Range, + snapshot: &TextBufferSnapshot, +) -> Range { + let start = if left.start.cmp(&right.start, snapshot).is_le() { + left.start + } else { + right.start + }; + let end = if left.end.cmp(&right.end, snapshot).is_ge() { + left.end + } else { + right.end + }; + start..end +} + #[derive(Error, Debug)] #[error( "You must update to Zed version {minimum_version} or higher to continue using edit predictions." diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 8f97df2c308980e1c2c89838609b30e1aedb1917..f377f3f705f8d3e04fd4718bbfd650ae4189ba37 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -1,7 +1,8 @@ use super::*; -use crate::{compute_diff_between_snapshots, udiff::apply_diff_to_string}; +use crate::udiff::apply_diff_to_string; use client::{UserStore, test::FakeServer}; use clock::FakeSystemClock; +use clock::ReplicaId; use cloud_api_types::{CreateLlmTokenResponse, LlmToken}; use cloud_llm_client::{ EditPredictionRejectReason, EditPredictionRejection, RejectEditPredictionsBody, @@ -18,8 +19,8 @@ use gpui::{ }; use indoc::indoc; use language::{ - Anchor, Buffer, CursorShape, Diagnostic, DiagnosticEntry, DiagnosticSet, DiagnosticSeverity, - Operation, Point, Selection, SelectionGoal, + Anchor, Buffer, Capability, CursorShape, Diagnostic, DiagnosticEntry, DiagnosticSet, + DiagnosticSeverity, Operation, Point, Selection, SelectionGoal, }; use language_model::RefreshLlmTokenListener; use lsp::LanguageServerId; @@ -28,7 +29,7 @@ use pretty_assertions::{assert_eq, assert_matches}; use project::{FakeFs, Project}; use serde_json::json; use settings::SettingsStore; -use std::{path::Path, sync::Arc, time::Duration}; +use std::{ops::Range, path::Path, sync::Arc, time::Duration}; use util::{ path, test::{TextRangeMarker, marked_text_ranges_by}, @@ -370,6 +371,12 @@ async fn test_edit_history_getter_pause_splits_last_event(cx: &mut TestAppContex ep_store.edit_history_for_project(&project, cx) }); assert_eq!(events.len(), 2); + + let first_total_edit_range = buffer.read_with(cx, |buffer, _| { + events[0].total_edit_range.to_point(&buffer.snapshot()) + }); + assert_eq!(first_total_edit_range, Point::new(1, 0)..Point::new(1, 3)); + let zeta_prompt::Event::BufferChange { diff, .. } = events[0].event.as_ref(); assert_eq!( diff.as_str(), @@ -382,6 +389,11 @@ async fn test_edit_history_getter_pause_splits_last_event(cx: &mut TestAppContex "} ); + let second_total_edit_range = buffer.read_with(cx, |buffer, _| { + events[1].total_edit_range.to_point(&buffer.snapshot()) + }); + assert_eq!(second_total_edit_range, Point::new(1, 3)..Point::new(1, 13)); + let zeta_prompt::Event::BufferChange { diff, .. } = events[1].event.as_ref(); assert_eq!( diff.as_str(), @@ -598,6 +610,240 @@ fn render_events_with_predicted(events: &[StoredEvent]) -> Vec { .collect() } +fn make_collaborator_replica( + buffer: &Entity, + cx: &mut TestAppContext, +) -> (Entity, clock::Global) { + let (state, version) = + buffer.read_with(cx, |buffer, _cx| (buffer.to_proto(_cx), buffer.version())); + let collaborator = cx.new(|_cx| { + Buffer::from_proto(ReplicaId::new(1), Capability::ReadWrite, state, None).unwrap() + }); + (collaborator, version) +} + +async fn apply_collaborator_edit( + collaborator: &Entity, + buffer: &Entity, + since_version: &mut clock::Global, + edit_range: Range, + new_text: &str, + cx: &mut TestAppContext, +) { + collaborator.update(cx, |collaborator, cx| { + collaborator.edit([(edit_range, new_text)], None, cx); + }); + + let serialize_task = collaborator.read_with(cx, |collaborator, cx| { + collaborator.serialize_ops(Some(since_version.clone()), cx) + }); + let ops = serialize_task.await; + *since_version = collaborator.read_with(cx, |collaborator, _cx| collaborator.version()); + + buffer.update(cx, |buffer, cx| { + buffer.apply_ops( + ops.into_iter() + .map(|op| language::proto::deserialize_operation(op).unwrap()), + cx, + ); + }); +} + +#[gpui::test] +async fn test_nearby_collaborator_edits_are_kept_in_history(cx: &mut TestAppContext) { + let (ep_store, _requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "foo.rs": "line 0\nline 1\nline 2\nline 3\nline 4\nline 5\nline 6\nline 7\nline 8\nline 9\nline 10\nline 11\nline 12\nline 13\nline 14\n" + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap(); + project.set_active_path(Some(path.clone()), cx); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + + let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0))); + + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&buffer, &project, cx); + let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx); + }); + + buffer.update(cx, |buffer, cx| { + buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx); + }); + + let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx); + + let (line_one_start, line_one_len) = collaborator.read_with(cx, |buffer, _cx| { + (Point::new(1, 0).to_offset(buffer), buffer.line_len(1)) + }); + + apply_collaborator_edit( + &collaborator, + &buffer, + &mut collaborator_version, + line_one_start..line_one_start + line_one_len as usize, + "REMOTE ONE", + cx, + ) + .await; + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + + assert_eq!( + render_events_with_predicted(&events), + vec![indoc! {" + manual + @@ -1,5 +1,5 @@ + -line 0 + -line 1 + +LOCAL ZERO + +REMOTE ONE + line 2 + line 3 + line 4 + "}] + ); +} + +#[gpui::test] +async fn test_distant_collaborator_edits_are_omitted_from_history(cx: &mut TestAppContext) { + let (ep_store, _requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "foo.rs": (0..1000) + .map(|i| format!("line {i}\n")) + .collect::() + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap(); + project.set_active_path(Some(path.clone()), cx); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + + let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0))); + + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&buffer, &project, cx); + let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx); + }); + + buffer.update(cx, |buffer, cx| { + buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx); + }); + + let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx); + + let far_line_start = buffer.read_with(cx, |buffer, _cx| Point::new(900, 0).to_offset(buffer)); + + apply_collaborator_edit( + &collaborator, + &buffer, + &mut collaborator_version, + far_line_start..far_line_start + 7, + "REMOTE FAR", + cx, + ) + .await; + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + + assert_eq!( + render_events_with_predicted(&events), + vec![indoc! {" + manual + @@ -1,4 +1,4 @@ + -line 0 + +LOCAL ZERO + line 1 + line 2 + line 3 + "}] + ); +} + +#[gpui::test] +async fn test_irrelevant_collaborator_edits_in_different_files_are_omitted_from_history( + cx: &mut TestAppContext, +) { + let (ep_store, _requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "foo.rs": "line 0\nline 1\nline 2\nline 3\n", + "bar.rs": "line 0\nline 1\nline 2\nline 3\n" + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let foo_buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap(); + project.set_active_path(Some(path.clone()), cx); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + let bar_buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/bar.rs"), cx).unwrap(); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + + let foo_cursor = foo_buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0))); + + ep_store.update(cx, |ep_store, cx| { + ep_store.register_buffer(&foo_buffer, &project, cx); + ep_store.register_buffer(&bar_buffer, &project, cx); + let _ = ep_store.prediction_at(&foo_buffer, Some(foo_cursor), &project, cx); + }); + + let (bar_collaborator, mut bar_version) = make_collaborator_replica(&bar_buffer, cx); + + apply_collaborator_edit( + &bar_collaborator, + &bar_buffer, + &mut bar_version, + 0..6, + "REMOTE BAR", + cx, + ) + .await; + + let events = ep_store.update(cx, |ep_store, cx| { + ep_store.edit_history_for_project(&project, cx) + }); + + assert!(events.is_empty()); +} + #[gpui::test] async fn test_predicted_flag_coalescing(cx: &mut TestAppContext) { let (ep_store, _requests) = init_test_with_fake_client(cx); @@ -680,7 +926,7 @@ async fn test_predicted_flag_coalescing(cx: &mut TestAppContext) { let end = Point::new(2, 6).to_offset(buffer); buffer.edit(vec![(offset..end, "LINE TWO")], None, cx); }); - ep_store.report_changes_for_buffer(&buffer, &project, true, cx); + ep_store.report_changes_for_buffer(&buffer, &project, true, true, cx); }); let events = ep_store.update(cx, |ep_store, cx| { @@ -722,7 +968,7 @@ async fn test_predicted_flag_coalescing(cx: &mut TestAppContext) { let end = Point::new(3, 6).to_offset(buffer); buffer.edit(vec![(offset..end, "LINE THREE")], None, cx); }); - ep_store.report_changes_for_buffer(&buffer, &project, true, cx); + ep_store.report_changes_for_buffer(&buffer, &project, true, true, cx); }); let events = ep_store.update(cx, |ep_store, cx| { @@ -2420,74 +2666,6 @@ async fn test_unauthenticated_without_custom_url_blocks_prediction_impl(cx: &mut ); } -#[gpui::test] -fn test_compute_diff_between_snapshots(cx: &mut TestAppContext) { - let buffer = cx.new(|cx| { - Buffer::local( - indoc! {" - zero - one - two - three - four - five - six - seven - eight - nine - ten - eleven - twelve - thirteen - fourteen - fifteen - sixteen - seventeen - eighteen - nineteen - twenty - twenty-one - twenty-two - twenty-three - twenty-four - "}, - cx, - ) - }); - - let old_snapshot = buffer.read_with(cx, |buffer, _| buffer.text_snapshot()); - - buffer.update(cx, |buffer, cx| { - let point = Point::new(12, 0); - buffer.edit([(point..point, "SECOND INSERTION\n")], None, cx); - let point = Point::new(8, 0); - buffer.edit([(point..point, "FIRST INSERTION\n")], None, cx); - }); - - let new_snapshot = buffer.read_with(cx, |buffer, _| buffer.text_snapshot()); - - let (diff, _) = compute_diff_between_snapshots(&old_snapshot, &new_snapshot).unwrap(); - - assert_eq!( - diff, - indoc! {" - @@ -6,10 +6,12 @@ - five - six - seven - +FIRST INSERTION - eight - nine - ten - eleven - +SECOND INSERTION - twelve - thirteen - fourteen - "} - ); -} - #[gpui::test] async fn test_diagnostic_jump_excludes_collaborator_regions(cx: &mut TestAppContext) { fn set_collaborator_cursor(buffer: &Entity, row: u32, cx: &mut TestAppContext) { diff --git a/crates/zeta_prompt/src/zeta_prompt.rs b/crates/zeta_prompt/src/zeta_prompt.rs index 41d02478c33ce807bf1771cf25799c9a427e63ed..8dd4d88e2a89cadc39e1335b4bcdc18a0a144571 100644 --- a/crates/zeta_prompt/src/zeta_prompt.rs +++ b/crates/zeta_prompt/src/zeta_prompt.rs @@ -479,6 +479,7 @@ pub fn format_prompt_with_budget_for_format( "<|file_sep|>", "edit history", budget_after_cursor, + max_edit_event_count_for_format(&format), ); let edit_history_tokens = estimate_tokens(edit_history_section.len()); let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens); @@ -516,6 +517,22 @@ pub fn filter_redundant_excerpts( related_files } +pub fn max_edit_event_count_for_format(format: &ZetaFormat) -> usize { + match format { + ZetaFormat::V0112MiddleAtEnd + | ZetaFormat::V0113Ordered + | ZetaFormat::V0114180EditableRegion + | ZetaFormat::V0120GitMergeMarkers + | ZetaFormat::V0131GitMergeMarkersPrefix + | ZetaFormat::V0211Prefill + | ZetaFormat::V0211SeedCoder + | ZetaFormat::v0226Hashline + | ZetaFormat::V0304SeedNoEdits + | ZetaFormat::V0304VariableEdit + | ZetaFormat::V0306SeedMultiRegions => 6, + } +} + pub fn get_prefill_for_format( format: ZetaFormat, context: &str, @@ -682,6 +699,7 @@ fn format_edit_history_within_budget( file_marker: &str, edit_history_name: &str, max_tokens: usize, + max_edit_event_count: usize, ) -> String { let header = format!("{}{}\n", file_marker, edit_history_name); let header_tokens = estimate_tokens(header.len()); @@ -692,7 +710,7 @@ fn format_edit_history_within_budget( let mut event_strings: Vec = Vec::new(); let mut total_tokens = header_tokens; - for event in events.iter().rev() { + for event in events.iter().rev().take(max_edit_event_count) { let mut event_str = String::new(); write_event(&mut event_str, event); let event_tokens = estimate_tokens(event_str.len()); @@ -2698,6 +2716,7 @@ pub mod seed_coder { FILE_MARKER, "edit_history", budget_after_cursor, + max_edit_event_count_for_format(&ZetaFormat::V0211SeedCoder), ); let edit_history_tokens = estimate_tokens(edit_history_section.len()); let budget_after_edit_history = budget_after_cursor.saturating_sub(edit_history_tokens); @@ -3824,7 +3843,13 @@ pub mod zeta1 { /// Formats events in zeta1 style (oldest first). fn format_zeta1_events(events: &[Arc]) -> String { let mut result = String::new(); - for event in events { + for event in + events + .iter() + .skip(events.len().saturating_sub(max_edit_event_count_for_format( + &ZetaFormat::V0114180EditableRegion, + ))) + { let event_string = format_zeta1_event(event); if event_string.is_empty() { continue; @@ -4781,6 +4806,87 @@ mod tests { ); } + #[test] + fn test_max_event_count() { + fn make_numbered_event(index: usize) -> Event { + return make_event( + &format!("event-{index}.rs"), + &format!("-old-{index}\n+new-{index}\n"), + ); + } + let input = make_input( + "x", + 0..1, + 0, + (0..3).map(make_numbered_event).collect(), + vec![], + ); + + let edit_history_section = format_edit_history_within_budget( + &input.events, + "<|file_sep|>", + "edit history", + usize::MAX, + 5, + ); + + assert_eq!( + &edit_history_section, + indoc!( + " + <|file_sep|>edit history + --- a/event-0.rs + +++ b/event-0.rs + -old-0 + +new-0 + --- a/event-1.rs + +++ b/event-1.rs + -old-1 + +new-1 + --- a/event-2.rs + +++ b/event-2.rs + -old-2 + +new-2 + " + ) + ); + + let edit_history_section = format_edit_history_within_budget( + &input.events, + "<|file_sep|>", + "edit history", + usize::MAX, + 2, + ); + + assert_eq!( + &edit_history_section, + indoc!( + " + <|file_sep|>edit history + --- a/event-1.rs + +++ b/event-1.rs + -old-1 + +new-1 + --- a/event-2.rs + +++ b/event-2.rs + -old-2 + +new-2 + " + ) + ); + + let edit_history_section = format_edit_history_within_budget( + &input.events, + "<|file_sep|>", + "edit history", + usize::MAX, + 0, + ); + + assert_eq!(&edit_history_section, ""); + } + #[test] fn test_clean_zeta1_model_output_basic() { let output = indoc! {"