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

Ben Kunkle created

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 ...

Change summary

crates/edit_prediction/src/edit_prediction.rs       | 308 +++++++++++--
crates/edit_prediction/src/edit_prediction_tests.rs | 326 +++++++++++---
crates/zeta_prompt/src/zeta_prompt.rs               | 110 ++++
3 files changed, 603 insertions(+), 141 deletions(-)

Detailed changes

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<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 🔗

@@ -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 🔗

@@ -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! {"