zeta: Extend the context window to include at least visible range

crates/copilot/src/copilot_completion_provider.rs

@@ -7,7 +7,7 @@ use language::{
     Buffer, OffsetRangeExt, ToOffset,
 };
 use settings::Settings;
-use std::{path::Path, time::Duration};
+use std::{ops::Range, path::Path, time::Duration};
 
 pub const COPILOT_DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(75);
 
@@ -76,6 +76,7 @@ impl InlineCompletionProvider for CopilotCompletionProvider {
         &mut self,
         buffer: Model<Buffer>,
         cursor_position: language::Anchor,
+        _visible_range: Option<Range<usize>>,
         debounce: bool,
         cx: &mut ModelContext<Self>,
     ) {

crates/editor/src/editor.rs

@@ -570,6 +570,7 @@ pub struct Editor {
     collaboration_hub: Option<Box<dyn CollaborationHub>>,
     blink_manager: Model<BlinkManager>,
     show_cursor_names: bool,
+    visible_range: Option<Range<Anchor>>,
     hovered_cursors: HashMap<HoveredCursor, Task<()>>,
     pub show_local_selections: bool,
     mode: EditorMode,
@@ -1310,6 +1311,7 @@ impl Editor {
             registered_buffers: HashMap::default(),
             _scroll_cursor_center_top_bottom_task: Task::ready(()),
             text_style_refinement: None,
+            visible_range: None,
         };
         this.tasks_update_task = Some(this.refresh_runnables(cx));
         this._subscriptions.extend(project_subscriptions);
@@ -4425,8 +4427,18 @@ impl Editor {
             return None;
         }
 
+        let excerpt_id = cursor.excerpt_id;
+        let visible_range = self.visible_range.as_ref().and_then(|visible_range| {
+            self.buffer
+                .read(cx)
+                .range_to_buffer_ranges(visible_range.clone(), cx)
+                .into_iter()
+                .find(|(_, _, buffer_excerpt_id)| *buffer_excerpt_id == excerpt_id)
+                .map(|(_, range, _)| range)
+        });
+
         self.update_visible_inline_completion(cx);
-        provider.refresh(buffer, cursor_buffer_position, debounce, cx);
+        provider.refresh(buffer, cursor_buffer_position, visible_range, debounce, cx);
         Some(())
     }
 
@@ -10715,6 +10727,10 @@ impl Editor {
         }
     }
 
+    pub fn set_visible_range(&mut self, range: Range<Anchor>) {
+        self.visible_range = Some(range);
+    }
+
     pub fn insert_blocks(
         &mut self,
         blocks: impl IntoIterator<Item = BlockProperties<Anchor>>,

crates/editor/src/element.rs

@@ -5395,6 +5395,10 @@ impl Element for EditorElement {
                         )
                     };
 
+                    self.editor.update(cx, |editor, _| {
+                        editor.set_visible_range(start_anchor..end_anchor);
+                    });
+
                     let highlighted_rows = self
                         .editor
                         .update(cx, |editor, cx| editor.highlighted_display_rows(cx));

crates/editor/src/inline_completion_tests.rs

@@ -330,6 +330,7 @@ impl InlineCompletionProvider for FakeInlineCompletionProvider {
         &mut self,
         _buffer: gpui::Model<language::Buffer>,
         _cursor_position: language::Anchor,
+        _visible_range: Option<Range<usize>>,
         _debounce: bool,
         _cx: &mut gpui::ModelContext<Self>,
     ) {

crates/inline_completion/src/inline_completion.rs

@@ -29,6 +29,7 @@ pub trait InlineCompletionProvider: 'static + Sized {
         &mut self,
         buffer: Model<Buffer>,
         cursor_position: language::Anchor,
+        visible_range: Option<Range<usize>>,
         debounce: bool,
         cx: &mut ModelContext<Self>,
     );
@@ -61,6 +62,7 @@ pub trait InlineCompletionProviderHandle {
         &self,
         buffer: Model<Buffer>,
         cursor_position: language::Anchor,
+        visible_range: Option<Range<usize>>,
         debounce: bool,
         cx: &mut AppContext,
     );
@@ -102,11 +104,12 @@ where
         &self,
         buffer: Model<Buffer>,
         cursor_position: language::Anchor,
+        visible_range: Option<Range<usize>>,
         debounce: bool,
         cx: &mut AppContext,
     ) {
         self.update(cx, |this, cx| {
-            this.refresh(buffer, cursor_position, debounce, cx)
+            this.refresh(buffer, cursor_position, visible_range, debounce, cx)
         })
     }
 

crates/supermaven/src/supermaven_completion_provider.rs

@@ -114,6 +114,7 @@ impl InlineCompletionProvider for SupermavenCompletionProvider {
         &mut self,
         buffer_handle: Model<Buffer>,
         cursor_position: Anchor,
+        _visible_range: Option<Range<usize>>,
         debounce: bool,
         cx: &mut ModelContext<Self>,
     ) {

crates/zeta/src/zeta.rs

@@ -273,6 +273,7 @@ impl Zeta {
         &mut self,
         buffer: &Model<Buffer>,
         position: language::Anchor,
+        visible_range: Option<Range<usize>>,
         cx: &mut ModelContext<Self>,
         perform_predict_edits: F,
     ) -> Task<Result<InlineCompletion>>
@@ -283,7 +284,9 @@ impl Zeta {
         let snapshot = self.report_changes_for_buffer(buffer, cx);
         let point = position.to_point(&snapshot);
         let offset = point.to_offset(&snapshot);
-        let excerpt_range = excerpt_range_for_position(point, &snapshot);
+        let visible_range = visible_range
+            .map(|range| range.start.to_point(&snapshot)..range.end.to_point(&snapshot));
+        let excerpt_range = excerpt_range_for_position(point, visible_range, &snapshot);
         let events = self.events.clone();
         let path = snapshot
             .file()
@@ -491,16 +494,23 @@ and then another
     ) -> Task<Result<InlineCompletion>> {
         use std::future::ready;
 
-        self.request_completion_impl(buffer, position, cx, |_, _, _| ready(Ok(response)))
+        self.request_completion_impl(buffer, position, None, cx, |_, _, _| ready(Ok(response)))
     }
 
     pub fn request_completion(
         &mut self,
         buffer: &Model<Buffer>,
         position: language::Anchor,
+        visible_range: Option<Range<usize>>,
         cx: &mut ModelContext<Self>,
     ) -> Task<Result<InlineCompletion>> {
-        self.request_completion_impl(buffer, position, cx, Self::perform_predict_edits)
+        self.request_completion_impl(
+            buffer,
+            position,
+            visible_range,
+            cx,
+            Self::perform_predict_edits,
+        )
     }
 
     fn perform_predict_edits(
@@ -797,21 +807,33 @@ fn prompt_for_excerpt(
     prompt_excerpt
 }
 
-fn excerpt_range_for_position(point: Point, snapshot: &BufferSnapshot) -> Range<usize> {
+fn excerpt_range_for_position(
+    point: Point,
+    visible_range: Option<Range<Point>>,
+    snapshot: &BufferSnapshot,
+) -> Range<usize> {
     const CONTEXT_LINES: u32 = 16;
 
-    let mut context_lines_before = CONTEXT_LINES;
-    let mut context_lines_after = CONTEXT_LINES;
-    if point.row < CONTEXT_LINES {
-        context_lines_after += CONTEXT_LINES - point.row;
-    } else if point.row + CONTEXT_LINES > snapshot.max_point().row {
-        context_lines_before += (point.row + CONTEXT_LINES) - snapshot.max_point().row;
+    let (mut excerpt_start_row, mut excerpt_end_row) = if let Some(visible) = visible_range {
+        (visible.start.row, visible.end.row)
+    } else {
+        (point.row, point.row)
+    };
+
+    let point_min_row = point.row.saturating_sub(CONTEXT_LINES);
+    let point_max_row = cmp::min(point.row + CONTEXT_LINES, snapshot.max_point().row);
+
+    if point_min_row < excerpt_start_row {
+        excerpt_start_row = point_min_row;
+    }
+
+    if point_max_row > excerpt_end_row {
+        excerpt_end_row = point_max_row;
     }
 
-    let excerpt_start_row = point.row.saturating_sub(context_lines_before);
     let excerpt_start = Point::new(excerpt_start_row, 0);
-    let excerpt_end_row = cmp::min(point.row + context_lines_after, snapshot.max_point().row);
     let excerpt_end = Point::new(excerpt_end_row, snapshot.line_len(excerpt_end_row));
+
     excerpt_start.to_offset(snapshot)..excerpt_end.to_offset(snapshot)
 }
 
@@ -941,6 +963,7 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
         &mut self,
         buffer: Model<Buffer>,
         position: language::Anchor,
+        visible_range: Option<Range<usize>>,
         debounce: bool,
         cx: &mut ModelContext<Self>,
     ) {
@@ -952,7 +975,7 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
 
             let completion_request = this.update(&mut cx, |this, cx| {
                 this.zeta.update(cx, |zeta, cx| {
-                    zeta.request_completion(&buffer, position, cx)
+                    zeta.request_completion(&buffer, position, visible_range, cx)
                 })
             });
 
@@ -1238,8 +1261,9 @@ mod tests {
         let zeta = cx.new_model(|cx| Zeta::new(client, cx));
         let buffer = cx.new_model(|cx| Buffer::local(buffer_content, cx));
         let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
-        let completion_task =
-            zeta.update(cx, |zeta, cx| zeta.request_completion(&buffer, cursor, cx));
+        let completion_task = zeta.update(cx, |zeta, cx| {
+            zeta.request_completion(&buffer, cursor, None, cx)
+        });
 
         let token_request = server.receive::<proto::GetLlmToken>().await.unwrap();
         server.respond(