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

ep: Avoid including collaborator edits in edit history sent to model (#51343)

Browse Source 
Closes #ISSUE

Before you mark this PR as ready for review, make sure that you have:
- [ ] Added a solid test coverage and/or screenshots from doing manual
testing
- [ ] Done a self-review taking into account security and performance
aspects
- [ ] Aligned any UI changes with the [UI
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)

Release Notes:

- N/A *or* Added/Fixed/Improved ...
This commit is contained in:
Ben Kunkle authored and GitHub committed 2026-03-12 16:56:37 -04:00
1 parent 5586fbf288
commit df8bafdccf
3 files changed
+603 -141
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/edit_prediction/src/edit_prediction.rs
+243 -65
View File
@@ -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<zeta_prompt::Event>,
pub old_snapshot: TextBufferSnapshot,
pub edit_range: Range<Anchor>,
pub new_snapshot_version: clock::Global,
pub total_edit_range: Range<Anchor>,
}
impl StoredEvent {
fn can_merge(
&self,
next_old_event: &&&StoredEvent,
new_snapshot: &TextBufferSnapshot,
last_edit_range: &Range<Anchor>,
next_old_event: &StoredEvent,
latest_snapshot: &TextBufferSnapshot,
latest_edit_range: &Range<Anchor>,
) -> 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<Arc<dyn File>>,
new_file: Option<Arc<dyn File>>,
edit_range: Option<Range<Anchor>>,
latest_edit_range: Range<Anchor>,
total_edit_range: Range<Anchor>,
total_edit_range_at_last_pause_boundary: Option<Range<Anchor>>,
predicted: bool,
snapshot_after_last_editing_pause: Option<TextBufferSnapshot>,
last_edit_time: Option<Instant>,
@@ -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<Point>)> {
) -> Option<Range<Anchor>> {
let edits: Vec<Edit<usize>> = new_snapshot
.edits_since::<usize>(&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<Anchor>,
) -> Option<Range<Point>> {
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<Edit<usize>> = new_snapshot
.edits_since::<usize>(&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<Anchor>,
) -> Option<(String, Range<Point>)> {
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<Buffer>,
project: &Entity<Project>,
is_predicted: bool,
is_local: bool,
cx: &mut Context<Self>,
) {
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(&current_prediction.prediction.buffer, project, true, cx);
self.report_changes_for_buffer(
&current_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<Project>,
buffer: &Entity<Buffer>,
snapshot: &BufferSnapshot,
edit_range: &Range<Anchor>,
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<StoredEvent>,
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<Anchor>,
right: &Range<Anchor>,
snapshot: &TextBufferSnapshot,
) -> Range<Anchor> {
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."
crates/edit_prediction/src/edit_prediction_tests.rs
+252 -74
View File
@@ -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<String> {
.collect()
}
fn make_collaborator_replica(
buffer: &Entity<Buffer>,
cx: &mut TestAppContext,
) -> (Entity<Buffer>, 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>,
buffer: &Entity<Buffer>,
since_version: &mut clock::Global,
edit_range: Range<usize>,
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::<String>()
}),
)
.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<Buffer>, row: u32, cx: &mut TestAppContext) {
crates/zeta_prompt/src/zeta_prompt.rs
+108 -2
View File
@@ -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<String> = 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<Event>]) -> 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! {"