ep: Don't open unnecessary files during context collection (#57318) (cherry-pick to preview) (#57444)

Cherry-pick of #57318 to preview

----
Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX

checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [ ] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Closes #ISSUE

Release Notes:

- N/A or Added/Fixed/Improved ...

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

crates/edit_prediction_context/src/edit_prediction_context.rs

@@ -286,7 +286,11 @@ impl RelatedExcerptStore {
                             let buffer = buffer.upgrade()?;
                             let definitions = project
                                 .update(cx, |project, cx| {
-                                    project.definitions(&buffer, identifier.range.start, cx)
+                                    project.workspace_definitions(
+                                        &buffer,
+                                        identifier.range.start,
+                                        cx,
+                                    )
                                 })
                                 .ok()?;
                             let type_definitions = project
@@ -296,7 +300,11 @@ impl RelatedExcerptStore {
                                     if is_tombi_lsp_in_toml(project, &buffer, cx) {
                                         return Task::ready(Ok(None));
                                     }
-                                    project.type_definitions(&buffer, identifier.range.start, cx)
+                                    project.workspace_type_definitions(
+                                        &buffer,
+                                        identifier.range.start,
+                                        cx,
+                                    )
                                 })
                                 .ok()?;
                             Some((definitions, type_definitions))
@@ -304,7 +312,6 @@ impl RelatedExcerptStore {
                     };
 
                     let cx = async_cx.clone();
-                    let project = project.clone();
                     async move {
                         match task {
                             DefinitionTask::CacheHit(cache_entry) => {
@@ -323,39 +330,39 @@ impl RelatedExcerptStore {
                                     .flatten()
                                     .unwrap_or_default();
 
-                                Some(cx.update(|cx| {
-                                    let definitions: SmallVec<[CachedDefinition; 1]> =
-                                        definition_locations
-                                            .into_iter()
-                                            .filter_map(|location| {
-                                                process_definition(location, &project, cx)
-                                            })
-                                            .collect();
+                                let definitions: SmallVec<[CachedDefinition; 1]> =
+                                    definition_locations
+                                        .into_iter()
+                                        .filter_map(|location| {
+                                            let mut cx = cx.clone();
+                                            process_definition(location, &mut cx)
+                                        })
+                                        .collect();
 
-                                    let type_definitions: SmallVec<[CachedDefinition; 1]> =
-                                        type_definition_locations
-                                            .into_iter()
-                                            .filter_map(|location| {
-                                                process_definition(location, &project, cx)
+                                let type_definitions: SmallVec<[CachedDefinition; 1]> =
+                                    type_definition_locations
+                                        .into_iter()
+                                        .filter_map(|location| {
+                                            let mut cx = cx.clone();
+                                            process_definition(location, &mut cx)
+                                        })
+                                        .filter(|type_def| {
+                                            !definitions.iter().any(|def| {
+                                                def.buffer.entity_id()
+                                                    == type_def.buffer.entity_id()
+                                                    && def.anchor_range == type_def.anchor_range
                                             })
-                                            .filter(|type_def| {
-                                                !definitions.iter().any(|def| {
-                                                    def.buffer.entity_id()
-                                                        == type_def.buffer.entity_id()
-                                                        && def.anchor_range == type_def.anchor_range
-                                                })
-                                            })
-                                            .collect();
+                                        })
+                                        .collect();
 
-                                    (
-                                        identifier,
-                                        Arc::new(CacheEntry {
-                                            definitions,
-                                            type_definitions,
-                                        }),
-                                        Some(duration),
-                                    )
-                                }))
+                                Some((
+                                    identifier,
+                                    Arc::new(CacheEntry {
+                                        definitions,
+                                        type_definitions,
+                                    }),
+                                    Some(duration),
+                                ))
                             }
                         }
                     }
@@ -581,34 +588,29 @@ use language::ToPoint as _;
 
 const MAX_TARGET_LEN: usize = 128;
 
-fn process_definition(
-    location: LocationLink,
-    project: &Entity<Project>,
-    cx: &mut App,
-) -> Option<CachedDefinition> {
-    let buffer = location.target.buffer.read(cx);
-    let anchor_range = location.target.range;
-    let file = buffer.file()?;
-    let worktree = project.read(cx).worktree_for_id(file.worktree_id(cx), cx)?;
-    if worktree.read(cx).is_single_file() {
-        return None;
-    }
-
-    // If the target range is large, it likely means we requested the definition of an entire module.
-    // For individual definitions, the target range should be small as it only covers the symbol.
-    let buffer = location.target.buffer.read(cx);
-    let target_len = anchor_range.to_offset(&buffer).len();
-    if target_len > MAX_TARGET_LEN {
-        return None;
-    }
-
-    Some(CachedDefinition {
-        path: ProjectPath {
+fn process_definition(location: LocationLink, cx: &mut AsyncApp) -> Option<CachedDefinition> {
+    cx.update(|cx| {
+        let buffer = location.target.buffer;
+        let buffer_snapshot = buffer.read(cx);
+        let file = buffer_snapshot.file()?;
+        let path = ProjectPath {
             worktree_id: file.worktree_id(cx),
             path: file.path().clone(),
-        },
-        buffer: location.target.buffer,
-        anchor_range,
+        };
+        let anchor_range = location.target.range;
+
+        // If the target range is large, it likely means we requested the definition of an entire module.
+        // For individual definitions, the target range should be small as it only covers the symbol.
+        let target_len = anchor_range.to_offset(&buffer_snapshot).len();
+        if target_len > MAX_TARGET_LEN {
+            return None;
+        }
+
+        Some(CachedDefinition {
+            path,
+            buffer: buffer.clone(),
+            anchor_range,
+        })
     })
 }
 

crates/project/src/lsp_command.rs

@@ -189,6 +189,7 @@ pub(crate) struct PerformRename {
 #[derive(Debug, Clone, Copy)]
 pub struct GetDefinitions {
     pub position: PointUtf16,
+    pub workspace_only: bool,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -199,6 +200,7 @@ pub(crate) struct GetDeclarations {
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct GetTypeDefinitions {
     pub position: PointUtf16,
+    pub workspace_only: bool,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -689,7 +691,15 @@ impl LspCommand for GetDefinitions {
         server_id: LanguageServerId,
         cx: AsyncApp,
     ) -> Result<Vec<LocationLink>> {
-        location_links_from_lsp(message, lsp_store, buffer, server_id, cx).await
+        location_links_from_lsp(
+            message,
+            lsp_store,
+            buffer,
+            server_id,
+            self.workspace_only,
+            cx,
+        )
+        .await
     }
 
     fn to_proto(&self, project_id: u64, buffer: &Buffer) -> proto::GetDefinition {
@@ -700,6 +710,7 @@ impl LspCommand for GetDefinitions {
                 &buffer.anchor_before(self.position),
             )),
             version: serialize_version(&buffer.version()),
+            workspace_only: self.workspace_only,
         }
     }
 
@@ -720,6 +731,7 @@ impl LspCommand for GetDefinitions {
             .await?;
         Ok(Self {
             position: buffer.read_with(&cx, |buffer, _| position.to_point_utf16(buffer)),
+            workspace_only: message.workspace_only,
         })
     }
 
@@ -792,7 +804,7 @@ impl LspCommand for GetDeclarations {
         server_id: LanguageServerId,
         cx: AsyncApp,
     ) -> Result<Vec<LocationLink>> {
-        location_links_from_lsp(message, lsp_store, buffer, server_id, cx).await
+        location_links_from_lsp(message, lsp_store, buffer, server_id, false, cx).await
     }
 
     fn to_proto(&self, project_id: u64, buffer: &Buffer) -> proto::GetDeclaration {
@@ -894,7 +906,7 @@ impl LspCommand for GetImplementations {
         server_id: LanguageServerId,
         cx: AsyncApp,
     ) -> Result<Vec<LocationLink>> {
-        location_links_from_lsp(message, lsp_store, buffer, server_id, cx).await
+        location_links_from_lsp(message, lsp_store, buffer, server_id, false, cx).await
     }
 
     fn to_proto(&self, project_id: u64, buffer: &Buffer) -> proto::GetImplementation {
@@ -993,7 +1005,7 @@ impl LspCommand for GetTypeDefinitions {
         server_id: LanguageServerId,
         cx: AsyncApp,
     ) -> Result<Vec<LocationLink>> {
-        location_links_from_lsp(message, project, buffer, server_id, cx).await
+        location_links_from_lsp(message, project, buffer, server_id, self.workspace_only, cx).await
     }
 
     fn to_proto(&self, project_id: u64, buffer: &Buffer) -> proto::GetTypeDefinition {
@@ -1004,6 +1016,7 @@ impl LspCommand for GetTypeDefinitions {
                 &buffer.anchor_before(self.position),
             )),
             version: serialize_version(&buffer.version()),
+            workspace_only: self.workspace_only,
         }
     }
 
@@ -1024,6 +1037,7 @@ impl LspCommand for GetTypeDefinitions {
             .await?;
         Ok(Self {
             position: buffer.read_with(&cx, |buffer, _| position.to_point_utf16(buffer)),
+            workspace_only: message.workspace_only,
         })
     }
 
@@ -1148,6 +1162,7 @@ pub async fn location_links_from_lsp(
     lsp_store: Entity<LspStore>,
     buffer: Entity<Buffer>,
     server_id: LanguageServerId,
+    workspace_only: bool,
     mut cx: AsyncApp,
 ) -> Result<Vec<LocationLink>> {
     let message = match message {
@@ -1179,6 +1194,25 @@ pub async fn location_links_from_lsp(
     let (_, language_server) = language_server_for_buffer(&lsp_store, &buffer, server_id, &mut cx)?;
     let mut definitions = Vec::new();
     for (origin_range, target_uri, target_range) in unresolved_links {
+        if workspace_only
+            && !lsp_store.update(&mut cx, |this, cx| {
+                use util::paths::UrlExt as _;
+                let worktree_store = this.worktree_store().read(cx);
+                let path_style = worktree_store.path_style();
+                let Ok(abs_path) = target_uri.clone().to_file_path_ext(path_style) else {
+                    return false;
+                };
+                worktree_store
+                    .find_worktree(&abs_path, cx)
+                    .is_some_and(|(worktree, _)| {
+                        let worktree = worktree.read(cx);
+                        worktree.is_visible() && !worktree.is_single_file()
+                    })
+            })
+        {
+            continue;
+        }
+
         let target_buffer_handle = lsp_store
             .update(&mut cx, |this, cx| {
                 this.open_local_buffer_via_lsp(target_uri, language_server.server_id(), cx)

crates/project/src/lsp_store.rs

@@ -5938,9 +5938,31 @@ impl LspStore {
         buffer: &Entity<Buffer>,
         position: PointUtf16,
         cx: &mut Context<Self>,
+    ) -> Task<Result<Option<Vec<LocationLink>>>> {
+        self.definitions_with_filter(buffer, position, false, cx)
+    }
+
+    pub fn workspace_definitions(
+        &mut self,
+        buffer: &Entity<Buffer>,
+        position: PointUtf16,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<Option<Vec<LocationLink>>>> {
+        self.definitions_with_filter(buffer, position, true, cx)
+    }
+
+    fn definitions_with_filter(
+        &mut self,
+        buffer: &Entity<Buffer>,
+        position: PointUtf16,
+        workspace_only: bool,
+        cx: &mut Context<Self>,
     ) -> Task<Result<Option<Vec<LocationLink>>>> {
         if let Some((upstream_client, project_id)) = self.upstream_client() {
-            let request = GetDefinitions { position };
+            let request = GetDefinitions {
+                position,
+                workspace_only,
+            };
             if !self.is_capable_for_proto_request(buffer, &request, cx) {
                 return Task::ready(Ok(None));
             }
@@ -5965,7 +5987,11 @@ impl LspStore {
                     return Ok(None);
                 };
                 let actions = join_all(responses.payload.into_iter().map(|response| {
-                    GetDefinitions { position }.response_from_proto(
+                    GetDefinitions {
+                        position,
+                        workspace_only,
+                    }
+                    .response_from_proto(
                         response.response,
                         lsp_store.clone(),
                         buffer.clone(),
@@ -5988,7 +6014,10 @@ impl LspStore {
             let definitions_task = self.request_multiple_lsp_locally(
                 buffer,
                 Some(position),
-                GetDefinitions { position },
+                GetDefinitions {
+                    position,
+                    workspace_only,
+                },
                 cx,
             );
             cx.background_spawn(async move {
@@ -6078,9 +6107,31 @@ impl LspStore {
         buffer: &Entity<Buffer>,
         position: PointUtf16,
         cx: &mut Context<Self>,
+    ) -> Task<Result<Option<Vec<LocationLink>>>> {
+        self.type_definitions_with_filter(buffer, position, false, cx)
+    }
+
+    pub fn workspace_type_definitions(
+        &mut self,
+        buffer: &Entity<Buffer>,
+        position: PointUtf16,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<Option<Vec<LocationLink>>>> {
+        self.type_definitions_with_filter(buffer, position, true, cx)
+    }
+
+    fn type_definitions_with_filter(
+        &mut self,
+        buffer: &Entity<Buffer>,
+        position: PointUtf16,
+        workspace_only: bool,
+        cx: &mut Context<Self>,
     ) -> Task<Result<Option<Vec<LocationLink>>>> {
         if let Some((upstream_client, project_id)) = self.upstream_client() {
-            let request = GetTypeDefinitions { position };
+            let request = GetTypeDefinitions {
+                position,
+                workspace_only,
+            };
             if !self.is_capable_for_proto_request(buffer, &request, cx) {
                 return Task::ready(Ok(None));
             }
@@ -6103,7 +6154,11 @@ impl LspStore {
                     return Ok(None);
                 };
                 let actions = join_all(responses.payload.into_iter().map(|response| {
-                    GetTypeDefinitions { position }.response_from_proto(
+                    GetTypeDefinitions {
+                        position,
+                        workspace_only,
+                    }
+                    .response_from_proto(
                         response.response,
                         lsp_store.clone(),
                         buffer.clone(),
@@ -6126,7 +6181,10 @@ impl LspStore {
             let type_definitions_task = self.request_multiple_lsp_locally(
                 buffer,
                 Some(position),
-                GetTypeDefinitions { position },
+                GetTypeDefinitions {
+                    position,
+                    workspace_only,
+                },
                 cx,
             );
             cx.background_spawn(async move {

crates/project/src/lsp_store/lsp_ext_command.rs

@@ -443,6 +443,7 @@ impl LspCommand for GoToParentModule {
             lsp_store,
             buffer,
             server_id,
+            false,
             cx,
         )
         .await

crates/project/src/project.rs

@@ -4180,6 +4180,24 @@ impl Project {
         })
     }
 
+    pub fn workspace_definitions<T: ToPointUtf16>(
+        &mut self,
+        buffer: &Entity<Buffer>,
+        position: T,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<Option<Vec<LocationLink>>>> {
+        let position = position.to_point_utf16(buffer.read(cx));
+        let guard = self.retain_remotely_created_models(cx);
+        let task = self.lsp_store.update(cx, |lsp_store, cx| {
+            lsp_store.workspace_definitions(buffer, position, cx)
+        });
+        cx.background_spawn(async move {
+            let result = task.await;
+            drop(guard);
+            result
+        })
+    }
+
     pub fn declarations<T: ToPointUtf16>(
         &mut self,
         buffer: &Entity<Buffer>,
@@ -4216,6 +4234,24 @@ impl Project {
         })
     }
 
+    pub fn workspace_type_definitions<T: ToPointUtf16>(
+        &mut self,
+        buffer: &Entity<Buffer>,
+        position: T,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<Option<Vec<LocationLink>>>> {
+        let position = position.to_point_utf16(buffer.read(cx));
+        let guard = self.retain_remotely_created_models(cx);
+        let task = self.lsp_store.update(cx, |lsp_store, cx| {
+            lsp_store.workspace_type_definitions(buffer, position, cx)
+        });
+        cx.background_spawn(async move {
+            let result = task.await;
+            drop(guard);
+            result
+        })
+    }
+
     pub fn implementations<T: ToPointUtf16>(
         &mut self,
         buffer: &Entity<Buffer>,

crates/proto/proto/lsp.proto

@@ -8,6 +8,7 @@ message GetDefinition {
   uint64 buffer_id = 2;
   Anchor position = 3;
   repeated VectorClockEntry version = 4;
+  bool workspace_only = 5;
 }
 
 message GetDefinitionResponse {
@@ -30,6 +31,7 @@ message GetTypeDefinition {
   uint64 buffer_id = 2;
   Anchor position = 3;
   repeated VectorClockEntry version = 4;
+  bool workspace_only = 5;
 }
 
 message GetTypeDefinitionResponse {