Add experimental LSP-based context retrieval system for edit prediction (#44036)

To do

* [x] Default to no context retrieval. Allow opting in to LSP-based
retrieval via a setting (for users in `zeta2` feature flag)
* [x] Feed this context to models when enabled
* [x] Make the zeta2 context view work well with LSP retrieval
* [x] Add a UI for the setting (for feature-flagged users)
* [x] Ensure Zeta CLI `context` command is usable

---

* [ ] Filter out LSP definitions that are too large / entire files (e.g.
modules)
* [ ] Introduce timeouts
* [ ] Test with other LSPs
* [ ] Figure out hangs

Release Notes:

- N/A

---------

Co-authored-by: Ben Kunkle <ben@zed.dev>
Co-authored-by: Agus Zubiaga <agus@zed.dev>

crates/language/src/buffer.rs

@@ -4022,6 +4022,20 @@ impl BufferSnapshot {
         })
     }
 
+    pub fn outline_items_as_offsets_containing<T: ToOffset>(
+        &self,
+        range: Range<T>,
+        include_extra_context: bool,
+        theme: Option<&SyntaxTheme>,
+    ) -> Vec<OutlineItem<usize>> {
+        self.outline_items_containing_internal(
+            range,
+            include_extra_context,
+            theme,
+            |buffer, range| range.to_offset(buffer),
+        )
+    }
+
     fn outline_items_containing_internal<T: ToOffset, U>(
         &self,
         range: Range<T>,

crates/language/src/buffer_tests.rs

@@ -784,28 +784,48 @@ async fn test_outline(cx: &mut gpui::TestAppContext) {
     .unindent();
 
     let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
-    let outline = buffer.update(cx, |buffer, _| buffer.snapshot().outline(None));
+    let snapshot = buffer.update(cx, |buffer, _| buffer.snapshot());
+    let outline = snapshot.outline(None);
 
-    assert_eq!(
+    pretty_assertions::assert_eq!(
         outline
             .items
             .iter()
-            .map(|item| (item.text.as_str(), item.depth))
+            .map(|item| (
+                item.text.as_str(),
+                item.depth,
+                item.to_point(&snapshot).body_range(&snapshot)
+                    .map(|range| minimize_space(&snapshot.text_for_range(range).collect::<String>()))
+            ))
             .collect::<Vec<_>>(),
         &[
-            ("struct Person", 0),
-            ("name", 1),
-            ("age", 1),
-            ("mod module", 0),
-            ("enum LoginState", 1),
-            ("LoggedOut", 2),
-            ("LoggingOn", 2),
-            ("LoggedIn", 2),
-            ("person", 3),
-            ("time", 3),
-            ("impl Eq for Person", 0),
-            ("impl Drop for Person", 0),
-            ("fn drop", 1),
+            ("struct Person", 0, Some("name: String, age: usize,".to_string())),
+            ("name", 1, None),
+            ("age", 1, None),
+            (
+                "mod module",
+                0,
+                Some(
+                    "enum LoginState { LoggedOut, LoggingOn, LoggedIn { person: Person, time: Instant, } }".to_string()
+                )
+            ),
+            (
+                "enum LoginState",
+                1,
+                Some("LoggedOut, LoggingOn, LoggedIn { person: Person, time: Instant, }".to_string())
+            ),
+            ("LoggedOut", 2, None),
+            ("LoggingOn", 2, None),
+            ("LoggedIn", 2, Some("person: Person, time: Instant,".to_string())),
+            ("person", 3, None),
+            ("time", 3, None),
+            ("impl Eq for Person", 0, None),
+            (
+                "impl Drop for Person",
+                0,
+                Some("fn drop(&mut self) { println!(\"bye\"); }".to_string())
+            ),
+            ("fn drop", 1, Some("println!(\"bye\");".to_string())),
         ]
     );
 
@@ -840,6 +860,11 @@ async fn test_outline(cx: &mut gpui::TestAppContext) {
         ]
     );
 
+    fn minimize_space(text: &str) -> String {
+        static WHITESPACE: LazyLock<Regex> = LazyLock::new(|| Regex::new("[\\n\\s]+").unwrap());
+        WHITESPACE.replace_all(text, " ").trim().to_string()
+    }
+
     async fn search<'a>(
         outline: &'a Outline<Anchor>,
         query: &'a str,

crates/language/src/language_registry.rs

@@ -437,26 +437,14 @@ impl LanguageRegistry {
         language_name: impl Into<LanguageName>,
         mut adapter: crate::FakeLspAdapter,
     ) -> futures::channel::mpsc::UnboundedReceiver<lsp::FakeLanguageServer> {
-        let language_name = language_name.into();
         let adapter_name = LanguageServerName(adapter.name.into());
         let capabilities = adapter.capabilities.clone();
         let initializer = adapter.initializer.take();
-        let adapter = CachedLspAdapter::new(Arc::new(adapter));
-        {
-            let mut state = self.state.write();
-            state
-                .lsp_adapters
-                .entry(language_name)
-                .or_default()
-                .push(adapter.clone());
-            state.all_lsp_adapters.insert(adapter.name(), adapter);
-        }
-
-        self.register_fake_language_server(adapter_name, capabilities, initializer)
+        self.register_fake_lsp_adapter(language_name, adapter);
+        self.register_fake_lsp_server(adapter_name, capabilities, initializer)
     }
 
     /// Register a fake lsp adapter (without the language server)
-    /// The returned channel receives a new instance of the language server every time it is started
     #[cfg(any(feature = "test-support", test))]
     pub fn register_fake_lsp_adapter(
         &self,
@@ -479,7 +467,7 @@ impl LanguageRegistry {
     /// Register a fake language server (without the adapter)
     /// The returned channel receives a new instance of the language server every time it is started
     #[cfg(any(feature = "test-support", test))]
-    pub fn register_fake_language_server(
+    pub fn register_fake_lsp_server(
         &self,
         lsp_name: LanguageServerName,
         capabilities: lsp::ServerCapabilities,

crates/language/src/language_settings.rs

@@ -373,6 +373,8 @@ impl InlayHintSettings {
 pub struct EditPredictionSettings {
     /// The provider that supplies edit predictions.
     pub provider: settings::EditPredictionProvider,
+    /// Whether to use the experimental edit prediction context retrieval system.
+    pub use_context: bool,
     /// A list of globs representing files that edit predictions should be disabled for.
     /// This list adds to a pre-existing, sensible default set of globs.
     /// Any additional ones you add are combined with them.
@@ -622,6 +624,11 @@ impl settings::Settings for AllLanguageSettings {
             .features
             .as_ref()
             .and_then(|f| f.edit_prediction_provider);
+        let use_edit_prediction_context = all_languages
+            .features
+            .as_ref()
+            .and_then(|f| f.experimental_edit_prediction_context_retrieval)
+            .unwrap_or_default();
 
         let edit_predictions = all_languages.edit_predictions.clone().unwrap();
         let edit_predictions_mode = edit_predictions.mode.unwrap();
@@ -668,6 +675,7 @@ impl settings::Settings for AllLanguageSettings {
                 } else {
                     EditPredictionProvider::None
                 },
+                use_context: use_edit_prediction_context,
                 disabled_globs: disabled_globs
                     .iter()
                     .filter_map(|g| {

crates/language/src/outline.rs

@@ -1,4 +1,4 @@
-use crate::{BufferSnapshot, Point, ToPoint};
+use crate::{BufferSnapshot, Point, ToPoint, ToTreeSitterPoint};
 use fuzzy::{StringMatch, StringMatchCandidate};
 use gpui::{BackgroundExecutor, HighlightStyle};
 use std::ops::Range;
@@ -48,6 +48,54 @@ impl<T: ToPoint> OutlineItem<T> {
                 .map(|r| r.start.to_point(buffer)..r.end.to_point(buffer)),
         }
     }
+
+    pub fn body_range(&self, buffer: &BufferSnapshot) -> Option<Range<Point>> {
+        if let Some(range) = self.body_range.as_ref() {
+            return Some(range.start.to_point(buffer)..range.end.to_point(buffer));
+        }
+
+        let range = self.range.start.to_point(buffer)..self.range.end.to_point(buffer);
+        let start_indent = buffer.indent_size_for_line(range.start.row);
+        let node = buffer.syntax_ancestor(range.clone())?;
+
+        let mut cursor = node.walk();
+        loop {
+            let node = cursor.node();
+            if node.start_position() >= range.start.to_ts_point()
+                && node.end_position() <= range.end.to_ts_point()
+            {
+                break;
+            }
+            cursor.goto_first_child_for_point(range.start.to_ts_point());
+        }
+
+        if !cursor.goto_last_child() {
+            return None;
+        }
+        let body_node = loop {
+            let node = cursor.node();
+            if node.child_count() > 0 {
+                break node;
+            }
+            if !cursor.goto_previous_sibling() {
+                return None;
+            }
+        };
+
+        let mut start_row = body_node.start_position().row as u32;
+        let mut end_row = body_node.end_position().row as u32;
+
+        while start_row < end_row && buffer.indent_size_for_line(start_row) == start_indent {
+            start_row += 1;
+        }
+        while start_row < end_row && buffer.indent_size_for_line(end_row - 1) == start_indent {
+            end_row -= 1;
+        }
+        if start_row < end_row {
+            return Some(Point::new(start_row, 0)..Point::new(end_row, 0));
+        }
+        None
+    }
 }
 
 impl<T> Outline<T> {

crates/language/src/syntax_map.rs

@@ -1215,6 +1215,19 @@ impl<'a> SyntaxMapMatches<'a> {
 
         true
     }
+
+    // pub fn set_byte_range(&mut self, range: Range<usize>) {
+    //     for layer in &mut self.layers {
+    //         layer.matches.set_byte_range(range.clone());
+    //         layer.advance();
+    //     }
+    //     self.layers.sort_unstable_by_key(|layer| layer.sort_key());
+    //     self.active_layer_count = self
+    //         .layers
+    //         .iter()
+    //         .position(|layer| !layer.has_next)
+    //         .unwrap_or(self.layers.len());
+    // }
 }
 
 impl SyntaxMapCapturesLayer<'_> {