Clean up environment loading a bit (#28356)

Closes #ISSUE

Release Notes:

- N/A

crates/project/src/debugger/dap_store.rs

@@ -339,8 +339,7 @@ impl DapStore {
             local_store.language_registry.clone(),
             local_store.toolchain_store.clone(),
             local_store.environment.update(cx, |env, cx| {
-                let worktree = worktree.read(cx);
+                env.get_worktree_environment(worktree.clone(), cx)
-                env.get_environment(worktree.abs_path().into(), cx)
             }),
         );
         let session_id = local_store.next_session_id();
@@ -414,8 +413,7 @@ impl DapStore {
             local_store.language_registry.clone(),
             local_store.toolchain_store.clone(),
             local_store.environment.update(cx, |env, cx| {
-                let worktree = worktree.read(cx);
+                env.get_worktree_environment(worktree.clone(), cx)
-                env.get_environment(Some(worktree.abs_path()), cx)
             }),
         );
         let session_id = local_store.next_session_id();

crates/project/src/environment.rs

@@ -1,22 +1,21 @@
-use futures::{
+use futures::{FutureExt, future::Shared};
-    FutureExt,
+use language::Buffer;
-    future::{Shared, WeakShared},
-};
 use std::{path::Path, sync::Arc};
 use util::ResultExt;
+use worktree::Worktree;
 
 use collections::HashMap;
-use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Task};
+use gpui::{AppContext as _, Context, Entity, EventEmitter, Task};
 use settings::Settings as _;
 
 use crate::{
     project_settings::{DirenvSettings, ProjectSettings},
-    worktree_store::{WorktreeStore, WorktreeStoreEvent},
+    worktree_store::WorktreeStore,
 };
 
 pub struct ProjectEnvironment {
     cli_environment: Option<HashMap<String, String>>,
-    environments: HashMap<Arc<Path>, WeakShared<Task<Option<HashMap<String, String>>>>>,
+    environments: HashMap<Arc<Path>, Shared<Task<Option<HashMap<String, String>>>>>,
     environment_error_messages: HashMap<Arc<Path>, EnvironmentErrorMessage>,
 }
 
@@ -27,27 +26,12 @@ pub enum ProjectEnvironmentEvent {
 impl EventEmitter<ProjectEnvironmentEvent> for ProjectEnvironment {}
 
 impl ProjectEnvironment {
-    pub fn new(
+    pub fn new(cli_environment: Option<HashMap<String, String>>) -> Self {
-        worktree_store: &Entity<WorktreeStore>,
-        cli_environment: Option<HashMap<String, String>>,
-        cx: &mut App,
-    ) -> Entity<Self> {
-        cx.new(|cx| {
-            cx.subscribe(worktree_store, |this: &mut Self, _, event, _| {
-                if let WorktreeStoreEvent::WorktreeRemoved(_, _) = event {
-                    this.environments.retain(|_, weak| weak.upgrade().is_some());
-                    this.environment_error_messages
-                        .retain(|abs_path, _| this.environments.contains_key(abs_path));
-                }
-            })
-            .detach();
-
         Self {
             cli_environment,
             environments: Default::default(),
             environment_error_messages: Default::default(),
         }
-        })
     }
 
     /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
@@ -73,54 +57,85 @@ impl ProjectEnvironment {
         cx.emit(ProjectEnvironmentEvent::ErrorsUpdated);
     }
 
-    /// Returns the project environment, if possible.
+    pub(crate) fn get_buffer_environment(
-    /// If the project was opened from the CLI, then the inherited CLI environment is returned.
-    /// If it wasn't opened from the CLI, and an absolute path is given, then a shell is spawned in
-    /// that directory, to get environment variables as if the user has `cd`'d there.
-    pub(crate) fn get_environment(
         &mut self,
-        abs_path: Option<Arc<Path>>,
+        buffer: Entity<Buffer>,
-        cx: &Context<Self>,
+        worktree_store: Entity<WorktreeStore>,
+        cx: &mut Context<Self>,
     ) -> Shared<Task<Option<HashMap<String, String>>>> {
         if cfg!(any(test, feature = "test-support")) {
             return Task::ready(Some(HashMap::default())).shared();
         }
 
         if let Some(cli_environment) = self.get_cli_environment() {
-            return cx
+            log::info!("using project environment variables from CLI");
-                .spawn(async move |_, _| {
+            return Task::ready(Some(cli_environment)).shared();
-                    let path = cli_environment
-                        .get("PATH")
-                        .map(|path| path.as_str())
-                        .unwrap_or_default();
-                    log::info!(
-                        "using project environment variables from CLI. PATH={:?}",
-                        path
-                    );
-                    Some(cli_environment)
-                })
-                .shared();
         }
 
-        let Some(abs_path) = abs_path else {
+        let Some(worktree) = buffer
+            .read(cx)
+            .file()
+            .map(|f| f.worktree_id(cx))
+            .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
+        else {
             return Task::ready(None).shared();
         };
 
-        if let Some(existing) = self
+        self.get_worktree_environment(worktree, cx)
-            .environments
-            .get(&abs_path)
-            .and_then(|weak| weak.upgrade())
-        {
-            existing
-        } else {
-            let env = get_directory_env(abs_path.clone(), cx).shared();
-            self.environments.insert(
-                abs_path.clone(),
-                env.downgrade()
-                    .expect("environment task has not been polled yet"),
-            );
-            env
     }
+
+    pub(crate) fn get_worktree_environment(
+        &mut self,
+        worktree: Entity<Worktree>,
+        cx: &mut Context<Self>,
+    ) -> Shared<Task<Option<HashMap<String, String>>>> {
+        if cfg!(any(test, feature = "test-support")) {
+            return Task::ready(Some(HashMap::default())).shared();
+        }
+
+        if let Some(cli_environment) = self.get_cli_environment() {
+            log::info!("using project environment variables from CLI");
+            return Task::ready(Some(cli_environment)).shared();
+        }
+
+        let mut abs_path = worktree.read(cx).abs_path();
+        if !worktree.read(cx).is_local() {
+            log::error!(
+                "attempted to get project environment for a non-local worktree at {abs_path:?}"
+            );
+            return Task::ready(None).shared();
+        } else if worktree.read(cx).is_single_file() {
+            let Some(parent) = abs_path.parent() else {
+                return Task::ready(None).shared();
+            };
+            abs_path = parent.into();
+        }
+
+        self.get_directory_environment(abs_path, cx)
+    }
+
+    /// Returns the project environment, if possible.
+    /// If the project was opened from the CLI, then the inherited CLI environment is returned.
+    /// If it wasn't opened from the CLI, and an absolute path is given, then a shell is spawned in
+    /// that directory, to get environment variables as if the user has `cd`'d there.
+    pub(crate) fn get_directory_environment(
+        &mut self,
+        abs_path: Arc<Path>,
+        cx: &mut Context<Self>,
+    ) -> Shared<Task<Option<HashMap<String, String>>>> {
+        if cfg!(any(test, feature = "test-support")) {
+            return Task::ready(Some(HashMap::default())).shared();
+        }
+
+        if let Some(cli_environment) = self.get_cli_environment() {
+            log::info!("using project environment variables from CLI");
+            return Task::ready(Some(cli_environment)).shared();
+        }
+
+        self.environments
+            .entry(abs_path.clone())
+            .or_insert_with(|| get_directory_env_impl(abs_path.clone(), cx).shared())
+            .clone()
     }
 }
 
@@ -338,7 +353,7 @@ async fn load_shell_environment(
     (Some(parsed_env), direnv_error)
 }
 
-fn get_directory_env(
+fn get_directory_env_impl(
     abs_path: Arc<Path>,
     cx: &Context<ProjectEnvironment>,
 ) -> Task<Option<HashMap<String, String>>> {
@@ -368,10 +383,7 @@ fn get_directory_env(
 
         if let Some(error) = error_message {
             this.update(cx, |this, cx| {
-                log::error!(
+                log::error!("{error}",);
-                    "error fetching environment for path {abs_path:?}: {}",
-                    error
-                );
                 this.environment_error_messages.insert(abs_path, error);
                 cx.emit(ProjectEnvironmentEvent::ErrorsUpdated)
             })

crates/project/src/git_store.rs

@@ -3776,7 +3776,7 @@ impl Repository {
                 .upgrade()
                 .ok_or_else(|| anyhow!("missing project environment"))?
                 .update(cx, |project_environment, cx| {
-                    project_environment.get_environment(Some(work_directory_abs_path.clone()), cx)
+                    project_environment.get_directory_environment(work_directory_abs_path.clone(), cx)
                 })?
                 .await
                 .unwrap_or_else(|| {

crates/project/src/lsp_store.rs

@@ -8291,16 +8291,10 @@ impl LspStore {
         buffer: &Entity<Buffer>,
         cx: &mut Context<Self>,
     ) -> Shared<Task<Option<HashMap<String, String>>>> {
-        let worktree_id = buffer.read(cx).file().map(|file| file.worktree_id(cx));
-        let worktree_abs_path = worktree_id.and_then(|worktree_id| {
-            self.worktree_store
-                .read(cx)
-                .worktree_for_id(worktree_id, cx)
-                .map(|entry| entry.read(cx).abs_path().clone())
-        });
-
         if let Some(environment) = &self.as_local().map(|local| local.environment.clone()) {
-            environment.update(cx, |env, cx| env.get_environment(worktree_abs_path, cx))
+            environment.update(cx, |env, cx| {
+                env.get_buffer_environment(buffer.clone(), self.worktree_store.clone(), cx)
+            })
         } else {
             Task::ready(None).shared()
         }
@@ -10134,10 +10128,8 @@ impl LocalLspAdapterDelegate {
         fs: Arc<dyn Fs>,
         cx: &mut App,
     ) -> Arc<Self> {
-        let worktree_abs_path = worktree.read(cx).abs_path();
-
         let load_shell_env_task = environment.update(cx, |env, cx| {
-            env.get_environment(Some(worktree_abs_path), cx)
+            env.get_worktree_environment(worktree.clone(), cx)
         });
 
         Arc::new(Self {

crates/project/src/project.rs

@@ -841,7 +841,7 @@ impl Project {
             cx.subscribe(&worktree_store, Self::on_worktree_store_event)
                 .detach();
 
-            let environment = ProjectEnvironment::new(&worktree_store, env, cx);
+            let environment = cx.new(|_| ProjectEnvironment::new(env));
             let toolchain_store = cx.new(|cx| {
                 ToolchainStore::local(
                     languages.clone(),
@@ -1043,7 +1043,7 @@ impl Project {
             cx.subscribe(&settings_observer, Self::on_settings_observer_event)
                 .detach();
 
-            let environment = ProjectEnvironment::new(&worktree_store, None, cx);
+            let environment = cx.new(|_| ProjectEnvironment::new(None));
 
             let lsp_store = cx.new(|cx| {
                 LspStore::new_remote(
@@ -1242,7 +1242,7 @@ impl Project {
             ImageStore::remote(worktree_store.clone(), client.clone().into(), remote_id, cx)
         })?;
 
-        let environment = cx.update(|cx| ProjectEnvironment::new(&worktree_store, None, cx))?;
+        let environment = cx.new(|_| ProjectEnvironment::new(None))?;
 
         let breakpoint_store =
             cx.new(|_| BreakpointStore::remote(remote_id, client.clone().into()))?;

crates/project/src/task_store.rs

@@ -295,10 +295,13 @@ fn local_task_context_for_location(
         .and_then(|worktree| worktree.read(cx).root_dir());
 
     cx.spawn(async move |cx| {
-        let worktree_abs_path = worktree_abs_path.clone();
         let project_env = environment
             .update(cx, |environment, cx| {
-                environment.get_environment(worktree_abs_path.clone(), cx)
+                environment.get_buffer_environment(
+                    location.buffer.clone(),
+                    worktree_store.clone(),
+                    cx,
+                )
             })
             .ok()?
             .await;

crates/project/src/toolchain_store.rs

@@ -317,21 +317,14 @@ impl LocalToolchainStore {
         cx: &App,
     ) -> Task<Option<ToolchainList>> {
         let registry = self.languages.clone();
-        let Some(root) = self
+        let Some(abs_path) = self.worktree_store.read(cx).absolutize(&path, cx) else {
-            .worktree_store
-            .read(cx)
-            .worktree_for_id(path.worktree_id, cx)
-            .map(|worktree| worktree.read(cx).abs_path())
-        else {
             return Task::ready(None);
         };
-
-        let abs_path = root.join(path.path);
         let environment = self.project_environment.clone();
         cx.spawn(async move |cx| {
             let project_env = environment
                 .update(cx, |environment, cx| {
-                    environment.get_environment(Some(root.clone()), cx)
+                    environment.get_directory_environment(abs_path.as_path().into(), cx)
                 })
                 .ok()?
                 .await;

crates/remote_server/src/headless_project.rs

@@ -9,7 +9,8 @@ use http_client::HttpClient;
 use language::{Buffer, BufferEvent, LanguageRegistry, proto::serialize_operation};
 use node_runtime::NodeRuntime;
 use project::{
-    LspStore, LspStoreEvent, PrettierStore, ProjectPath, ToolchainStore, WorktreeId,
+    LspStore, LspStoreEvent, PrettierStore, ProjectEnvironment, ProjectPath, ToolchainStore,
+    WorktreeId,
     buffer_store::{BufferStore, BufferStoreEvent},
     debugger::{breakpoint_store::BreakpointStore, dap_store::DapStore},
     git_store::GitStore,
@@ -85,7 +86,7 @@ impl HeadlessProject {
             store
         });
 
-        let environment = project::ProjectEnvironment::new(&worktree_store, None, cx);
+        let environment = cx.new(|_| ProjectEnvironment::new(None));
 
         let toolchain_store = cx.new(|cx| {
             ToolchainStore::local(