environment.rs

  1use futures::{FutureExt, future::Shared};
  2use language::Buffer;
  3use std::{path::Path, sync::Arc};
  4use util::ResultExt;
  5use worktree::Worktree;
  6
  7use collections::HashMap;
  8use gpui::{AppContext as _, Context, Entity, EventEmitter, Task};
  9use settings::Settings as _;
 10
 11use crate::{
 12    project_settings::{DirenvSettings, ProjectSettings},
 13    worktree_store::WorktreeStore,
 14};
 15
 16pub struct ProjectEnvironment {
 17    cli_environment: Option<HashMap<String, String>>,
 18    environments: HashMap<Arc<Path>, Shared<Task<Option<HashMap<String, String>>>>>,
 19    environment_error_messages: HashMap<Arc<Path>, EnvironmentErrorMessage>,
 20}
 21
 22pub enum ProjectEnvironmentEvent {
 23    ErrorsUpdated,
 24}
 25
 26impl EventEmitter<ProjectEnvironmentEvent> for ProjectEnvironment {}
 27
 28impl ProjectEnvironment {
 29    pub fn new(cli_environment: Option<HashMap<String, String>>) -> Self {
 30        Self {
 31            cli_environment,
 32            environments: Default::default(),
 33            environment_error_messages: Default::default(),
 34        }
 35    }
 36
 37    /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
 38    pub(crate) fn get_cli_environment(&self) -> Option<HashMap<String, String>> {
 39        if let Some(mut env) = self.cli_environment.clone() {
 40            set_origin_marker(&mut env, EnvironmentOrigin::Cli);
 41            Some(env)
 42        } else {
 43            None
 44        }
 45    }
 46
 47    /// Returns an iterator over all pairs `(abs_path, error_message)` of
 48    /// environment errors associated with this project environment.
 49    pub(crate) fn environment_errors(
 50        &self,
 51    ) -> impl Iterator<Item = (&Arc<Path>, &EnvironmentErrorMessage)> {
 52        self.environment_error_messages.iter()
 53    }
 54
 55    pub(crate) fn remove_environment_error(&mut self, abs_path: &Path, cx: &mut Context<Self>) {
 56        self.environment_error_messages.remove(abs_path);
 57        cx.emit(ProjectEnvironmentEvent::ErrorsUpdated);
 58    }
 59
 60    pub(crate) fn get_buffer_environment(
 61        &mut self,
 62        buffer: &Entity<Buffer>,
 63        worktree_store: &Entity<WorktreeStore>,
 64        cx: &mut Context<Self>,
 65    ) -> Shared<Task<Option<HashMap<String, String>>>> {
 66        if cfg!(any(test, feature = "test-support")) {
 67            return Task::ready(Some(HashMap::default())).shared();
 68        }
 69
 70        if let Some(cli_environment) = self.get_cli_environment() {
 71            log::debug!("using project environment variables from CLI");
 72            return Task::ready(Some(cli_environment)).shared();
 73        }
 74
 75        let Some(worktree) = buffer
 76            .read(cx)
 77            .file()
 78            .map(|f| f.worktree_id(cx))
 79            .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
 80        else {
 81            return Task::ready(None).shared();
 82        };
 83
 84        self.get_worktree_environment(worktree, cx)
 85    }
 86
 87    pub fn get_worktree_environment(
 88        &mut self,
 89        worktree: Entity<Worktree>,
 90        cx: &mut Context<Self>,
 91    ) -> Shared<Task<Option<HashMap<String, String>>>> {
 92        if cfg!(any(test, feature = "test-support")) {
 93            return Task::ready(Some(HashMap::default())).shared();
 94        }
 95
 96        if let Some(cli_environment) = self.get_cli_environment() {
 97            log::debug!("using project environment variables from CLI");
 98            return Task::ready(Some(cli_environment)).shared();
 99        }
100
101        let mut abs_path = worktree.read(cx).abs_path();
102        if !worktree.read(cx).is_local() {
103            log::error!(
104                "attempted to get project environment for a non-local worktree at {abs_path:?}"
105            );
106            return Task::ready(None).shared();
107        } else if worktree.read(cx).is_single_file() {
108            let Some(parent) = abs_path.parent() else {
109                return Task::ready(None).shared();
110            };
111            abs_path = parent.into();
112        }
113
114        self.get_directory_environment(abs_path, cx)
115    }
116
117    /// Returns the project environment, if possible.
118    /// If the project was opened from the CLI, then the inherited CLI environment is returned.
119    /// If it wasn't opened from the CLI, and an absolute path is given, then a shell is spawned in
120    /// that directory, to get environment variables as if the user has `cd`'d there.
121    pub fn get_directory_environment(
122        &mut self,
123        abs_path: Arc<Path>,
124        cx: &mut Context<Self>,
125    ) -> Shared<Task<Option<HashMap<String, String>>>> {
126        if cfg!(any(test, feature = "test-support")) {
127            return Task::ready(Some(HashMap::default())).shared();
128        }
129
130        if let Some(cli_environment) = self.get_cli_environment() {
131            log::debug!("using project environment variables from CLI");
132            return Task::ready(Some(cli_environment)).shared();
133        }
134
135        self.environments
136            .entry(abs_path.clone())
137            .or_insert_with(|| get_directory_env_impl(abs_path.clone(), cx).shared())
138            .clone()
139    }
140}
141
142fn set_origin_marker(env: &mut HashMap<String, String>, origin: EnvironmentOrigin) {
143    env.insert(ZED_ENVIRONMENT_ORIGIN_MARKER.to_string(), origin.into());
144}
145
146const ZED_ENVIRONMENT_ORIGIN_MARKER: &str = "ZED_ENVIRONMENT";
147
148enum EnvironmentOrigin {
149    Cli,
150    WorktreeShell,
151}
152
153impl From<EnvironmentOrigin> for String {
154    fn from(val: EnvironmentOrigin) -> Self {
155        match val {
156            EnvironmentOrigin::Cli => "cli".into(),
157            EnvironmentOrigin::WorktreeShell => "worktree-shell".into(),
158        }
159    }
160}
161
162#[derive(Debug)]
163pub struct EnvironmentErrorMessage(pub String);
164
165impl std::fmt::Display for EnvironmentErrorMessage {
166    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167        write!(f, "{}", self.0)
168    }
169}
170
171impl EnvironmentErrorMessage {
172    #[allow(dead_code)]
173    fn from_str(s: &str) -> Self {
174        Self(String::from(s))
175    }
176}
177
178async fn load_directory_shell_environment(
179    abs_path: &Path,
180    load_direnv: &DirenvSettings,
181) -> (
182    Option<HashMap<String, String>>,
183    Option<EnvironmentErrorMessage>,
184) {
185    match smol::fs::metadata(abs_path).await {
186        Ok(meta) => {
187            let dir = if meta.is_dir() {
188                abs_path
189            } else if let Some(parent) = abs_path.parent() {
190                parent
191            } else {
192                return (
193                    None,
194                    Some(EnvironmentErrorMessage(format!(
195                        "Failed to load shell environment in {}: not a directory",
196                        abs_path.display()
197                    ))),
198                );
199            };
200
201            load_shell_environment(dir, load_direnv).await
202        }
203        Err(err) => (
204            None,
205            Some(EnvironmentErrorMessage(format!(
206                "Failed to load shell environment in {}: {}",
207                abs_path.display(),
208                err
209            ))),
210        ),
211    }
212}
213
214#[cfg(any(test, feature = "test-support"))]
215async fn load_shell_environment(
216    _dir: &Path,
217    _load_direnv: &DirenvSettings,
218) -> (
219    Option<HashMap<String, String>>,
220    Option<EnvironmentErrorMessage>,
221) {
222    let fake_env = [("ZED_FAKE_TEST_ENV".into(), "true".into())]
223        .into_iter()
224        .collect();
225    (Some(fake_env), None)
226}
227
228#[cfg(all(target_os = "windows", not(any(test, feature = "test-support"))))]
229async fn load_shell_environment(
230    _dir: &Path,
231    _load_direnv: &DirenvSettings,
232) -> (
233    Option<HashMap<String, String>>,
234    Option<EnvironmentErrorMessage>,
235) {
236    // TODO the current code works with Unix $SHELL only, implement environment loading on windows
237    (None, None)
238}
239
240#[cfg(not(any(target_os = "windows", test, feature = "test-support")))]
241async fn load_shell_environment(
242    dir: &Path,
243    load_direnv: &DirenvSettings,
244) -> (
245    Option<HashMap<String, String>>,
246    Option<EnvironmentErrorMessage>,
247) {
248    use crate::direnv::load_direnv_environment;
249    use util::shell_env;
250
251    let dir_ = dir.to_owned();
252    let mut envs = match smol::unblock(move || shell_env::capture(&dir_)).await {
253        Ok(envs) => envs,
254        Err(err) => {
255            util::log_err(&err);
256            return (
257                None,
258                Some(EnvironmentErrorMessage::from_str(
259                    "Failed to load environment variables. See log for details",
260                )),
261            );
262        }
263    };
264
265    // If the user selects `Direct` for direnv, it would set an environment
266    // variable that later uses to know that it should not run the hook.
267    // We would include in `.envs` call so it is okay to run the hook
268    // even if direnv direct mode is enabled.
269    let (direnv_environment, direnv_error) = match load_direnv {
270        DirenvSettings::ShellHook => (None, None),
271        DirenvSettings::Direct => match load_direnv_environment(&envs, dir).await {
272            Ok(env) => (Some(env), None),
273            Err(err) => (None, err.into()),
274        },
275    };
276    if let Some(direnv_environment) = direnv_environment {
277        for (key, value) in direnv_environment {
278            if let Some(value) = value {
279                envs.insert(key, value);
280            } else {
281                envs.remove(&key);
282            }
283        }
284    }
285
286    (Some(envs), direnv_error)
287}
288
289fn get_directory_env_impl(
290    abs_path: Arc<Path>,
291    cx: &Context<ProjectEnvironment>,
292) -> Task<Option<HashMap<String, String>>> {
293    let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
294
295    cx.spawn(async move |this, cx| {
296        let (mut shell_env, error_message) = cx
297            .background_spawn({
298                let abs_path = abs_path.clone();
299                async move { load_directory_shell_environment(&abs_path, &load_direnv).await }
300            })
301            .await;
302
303        if let Some(shell_env) = shell_env.as_mut() {
304            let path = shell_env
305                .get("PATH")
306                .map(|path| path.as_str())
307                .unwrap_or_default();
308            log::info!(
309                "using project environment variables shell launched in {:?}. PATH={:?}",
310                abs_path,
311                path
312            );
313
314            set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
315        }
316
317        if let Some(error) = error_message {
318            this.update(cx, |this, cx| {
319                log::error!("{error}",);
320                this.environment_error_messages.insert(abs_path, error);
321                cx.emit(ProjectEnvironmentEvent::ErrorsUpdated)
322            })
323            .log_err();
324        }
325
326        shell_env
327    })
328}