environment.rs

  1use futures::{future::Shared, FutureExt};
  2use std::{path::Path, sync::Arc};
  3use util::ResultExt;
  4
  5use collections::HashMap;
  6use gpui::{AppContext, Context, Model, ModelContext, Task};
  7use settings::Settings as _;
  8use worktree::WorktreeId;
  9
 10use crate::{
 11    project_settings::{DirenvSettings, ProjectSettings},
 12    worktree_store::{WorktreeStore, WorktreeStoreEvent},
 13};
 14
 15pub struct ProjectEnvironment {
 16    cli_environment: Option<HashMap<String, String>>,
 17    environments: HashMap<WorktreeId, Shared<Task<Option<HashMap<String, String>>>>>,
 18    environment_error_messages: HashMap<WorktreeId, EnvironmentErrorMessage>,
 19}
 20
 21impl ProjectEnvironment {
 22    pub fn new(
 23        worktree_store: &Model<WorktreeStore>,
 24        cli_environment: Option<HashMap<String, String>>,
 25        cx: &mut AppContext,
 26    ) -> Model<Self> {
 27        cx.new_model(|cx| {
 28            cx.subscribe(worktree_store, |this: &mut Self, _, event, _| {
 29                if let WorktreeStoreEvent::WorktreeRemoved(_, id) = event {
 30                    this.remove_worktree_environment(*id);
 31                }
 32            })
 33            .detach();
 34
 35            Self {
 36                cli_environment,
 37                environments: Default::default(),
 38                environment_error_messages: Default::default(),
 39            }
 40        })
 41    }
 42
 43    pub(crate) fn remove_worktree_environment(&mut self, worktree_id: WorktreeId) {
 44        self.environment_error_messages.remove(&worktree_id);
 45        self.environments.remove(&worktree_id);
 46    }
 47
 48    /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
 49    pub(crate) fn get_cli_environment(&self) -> Option<HashMap<String, String>> {
 50        if let Some(mut env) = self.cli_environment.clone() {
 51            set_origin_marker(&mut env, EnvironmentOrigin::Cli);
 52            Some(env)
 53        } else {
 54            None
 55        }
 56    }
 57
 58    /// Returns an iterator over all pairs `(worktree_id, error_message)` of
 59    /// environment errors associated with this project environment.
 60    pub(crate) fn environment_errors(
 61        &self,
 62    ) -> impl Iterator<Item = (&WorktreeId, &EnvironmentErrorMessage)> {
 63        self.environment_error_messages.iter()
 64    }
 65
 66    pub(crate) fn remove_environment_error(&mut self, worktree_id: WorktreeId) {
 67        self.environment_error_messages.remove(&worktree_id);
 68    }
 69
 70    /// Returns the project environment, if possible.
 71    /// If the project was opened from the CLI, then the inherited CLI environment is returned.
 72    /// If it wasn't opened from the CLI, and a worktree is given, then a shell is spawned in
 73    /// the worktree's path, to get environment variables as if the user has `cd`'d into
 74    /// the worktrees path.
 75    pub(crate) fn get_environment(
 76        &mut self,
 77        worktree_id: Option<WorktreeId>,
 78        worktree_abs_path: Option<Arc<Path>>,
 79        cx: &ModelContext<Self>,
 80    ) -> Shared<Task<Option<HashMap<String, String>>>> {
 81        if cfg!(any(test, feature = "test-support")) {
 82            return Task::ready(Some(HashMap::default())).shared();
 83        }
 84
 85        if let Some(cli_environment) = self.get_cli_environment() {
 86            return cx
 87                .spawn(|_, _| async move {
 88                    let path = cli_environment
 89                        .get("PATH")
 90                        .map(|path| path.as_str())
 91                        .unwrap_or_default();
 92                    log::info!(
 93                        "using project environment variables from CLI. PATH={:?}",
 94                        path
 95                    );
 96                    Some(cli_environment)
 97                })
 98                .shared();
 99        }
100
101        let Some((worktree_id, worktree_abs_path)) = worktree_id.zip(worktree_abs_path) else {
102            return Task::ready(None).shared();
103        };
104
105        if let Some(task) = self.environments.get(&worktree_id) {
106            task.clone()
107        } else {
108            let task = self
109                .get_worktree_env(worktree_id, worktree_abs_path, cx)
110                .shared();
111            self.environments.insert(worktree_id, task.clone());
112            task
113        }
114    }
115
116    fn get_worktree_env(
117        &mut self,
118        worktree_id: WorktreeId,
119        worktree_abs_path: Arc<Path>,
120        cx: &ModelContext<Self>,
121    ) -> Task<Option<HashMap<String, String>>> {
122        let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
123
124        cx.spawn(|this, mut cx| async move {
125            let (mut shell_env, error_message) = cx
126                .background_executor()
127                .spawn({
128                    let worktree_abs_path = worktree_abs_path.clone();
129                    async move {
130                        load_worktree_shell_environment(&worktree_abs_path, &load_direnv).await
131                    }
132                })
133                .await;
134
135            if let Some(shell_env) = shell_env.as_mut() {
136                let path = shell_env
137                    .get("PATH")
138                    .map(|path| path.as_str())
139                    .unwrap_or_default();
140                log::info!(
141                    "using project environment variables shell launched in {:?}. PATH={:?}",
142                    worktree_abs_path,
143                    path
144                );
145
146                set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
147            }
148
149            if let Some(error) = error_message {
150                this.update(&mut cx, |this, _| {
151                    this.environment_error_messages.insert(worktree_id, error);
152                })
153                .log_err();
154            }
155
156            shell_env
157        })
158    }
159}
160
161fn set_origin_marker(env: &mut HashMap<String, String>, origin: EnvironmentOrigin) {
162    env.insert(ZED_ENVIRONMENT_ORIGIN_MARKER.to_string(), origin.into());
163}
164
165const ZED_ENVIRONMENT_ORIGIN_MARKER: &str = "ZED_ENVIRONMENT";
166
167enum EnvironmentOrigin {
168    Cli,
169    WorktreeShell,
170}
171
172impl From<EnvironmentOrigin> for String {
173    fn from(val: EnvironmentOrigin) -> Self {
174        match val {
175            EnvironmentOrigin::Cli => "cli".into(),
176            EnvironmentOrigin::WorktreeShell => "worktree-shell".into(),
177        }
178    }
179}
180
181pub struct EnvironmentErrorMessage(pub String);
182
183impl EnvironmentErrorMessage {
184    #[allow(dead_code)]
185    fn from_str(s: &str) -> Self {
186        Self(String::from(s))
187    }
188}
189
190async fn load_worktree_shell_environment(
191    worktree_abs_path: &Path,
192    load_direnv: &DirenvSettings,
193) -> (
194    Option<HashMap<String, String>>,
195    Option<EnvironmentErrorMessage>,
196) {
197    match smol::fs::metadata(worktree_abs_path).await {
198        Ok(meta) => {
199            let dir = if meta.is_dir() {
200                worktree_abs_path
201            } else if let Some(parent) = worktree_abs_path.parent() {
202                parent
203            } else {
204                return (
205                    None,
206                    Some(EnvironmentErrorMessage(format!(
207                        "Failed to load shell environment in {}: not a directory",
208                        worktree_abs_path.display()
209                    ))),
210                );
211            };
212
213            load_shell_environment(&dir, load_direnv).await
214        }
215        Err(err) => (
216            None,
217            Some(EnvironmentErrorMessage(format!(
218                "Failed to load shell environment in {}: {}",
219                worktree_abs_path.display(),
220                err
221            ))),
222        ),
223    }
224}
225
226#[cfg(any(test, feature = "test-support"))]
227async fn load_shell_environment(
228    _dir: &Path,
229    _load_direnv: &DirenvSettings,
230) -> (
231    Option<HashMap<String, String>>,
232    Option<EnvironmentErrorMessage>,
233) {
234    let fake_env = [("ZED_FAKE_TEST_ENV".into(), "true".into())]
235        .into_iter()
236        .collect();
237    (Some(fake_env), None)
238}
239
240#[cfg(all(target_os = "windows", not(any(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    // TODO the current code works with Unix $SHELL only, implement environment loading on windows
249    (None, None)
250}
251
252#[cfg(not(any(target_os = "windows", test, feature = "test-support")))]
253async fn load_shell_environment(
254    dir: &Path,
255    load_direnv: &DirenvSettings,
256) -> (
257    Option<HashMap<String, String>>,
258    Option<EnvironmentErrorMessage>,
259) {
260    use crate::direnv::{load_direnv_environment, DirenvError};
261    use std::path::PathBuf;
262    use util::parse_env_output;
263
264    fn message<T>(with: &str) -> (Option<T>, Option<EnvironmentErrorMessage>) {
265        let message = EnvironmentErrorMessage::from_str(with);
266        (None, Some(message))
267    }
268
269    let marker = "ZED_SHELL_START";
270    let Some(shell) = std::env::var("SHELL").log_err() else {
271        return message("Failed to get login environment. SHELL environment variable is not set");
272    };
273
274    // What we're doing here is to spawn a shell and then `cd` into
275    // the project directory to get the env in there as if the user
276    // `cd`'d into it. We do that because tools like direnv, asdf, ...
277    // hook into `cd` and only set up the env after that.
278    //
279    // If the user selects `Direct` for direnv, it would set an environment
280    // variable that later uses to know that it should not run the hook.
281    // We would include in `.envs` call so it is okay to run the hook
282    // even if direnv direct mode is enabled.
283    //
284    // In certain shells we need to execute additional_command in order to
285    // trigger the behavior of direnv, etc.
286    //
287    //
288    // The `exit 0` is the result of hours of debugging, trying to find out
289    // why running this command here, without `exit 0`, would mess
290    // up signal process for our process so that `ctrl-c` doesn't work
291    // anymore.
292    //
293    // We still don't know why `$SHELL -l -i -c '/usr/bin/env -0'`  would
294    // do that, but it does, and `exit 0` helps.
295    let additional_command = PathBuf::from(&shell)
296        .file_name()
297        .and_then(|f| f.to_str())
298        .and_then(|shell| match shell {
299            "fish" => Some("emit fish_prompt;"),
300            _ => None,
301        });
302
303    let command = format!(
304        "cd '{}';{} printf '%s' {marker}; /usr/bin/env; exit 0;",
305        dir.display(),
306        additional_command.unwrap_or("")
307    );
308
309    let Some(output) = smol::process::Command::new(&shell)
310        .args(["-l", "-i", "-c", &command])
311        .output()
312        .await
313        .log_err()
314    else {
315        return message("Failed to spawn login shell to source login environment variables. See logs for details");
316    };
317
318    if !output.status.success() {
319        log::error!("login shell exited with {}", output.status);
320        return message("Login shell exited with nonzero exit code. See logs for details");
321    }
322
323    let stdout = String::from_utf8_lossy(&output.stdout);
324    let Some(env_output_start) = stdout.find(marker) else {
325        let stderr = String::from_utf8_lossy(&output.stderr);
326        log::error!(
327            "failed to parse output of `env` command in login shell. stdout: {:?}, stderr: {:?}",
328            stdout,
329            stderr
330        );
331        return message("Failed to parse stdout of env command. See logs for the output");
332    };
333
334    let mut parsed_env = HashMap::default();
335    let env_output = &stdout[env_output_start + marker.len()..];
336
337    parse_env_output(env_output, |key, value| {
338        parsed_env.insert(key, value);
339    });
340
341    let (direnv_environment, direnv_error) = match load_direnv {
342        DirenvSettings::ShellHook => (None, None),
343        DirenvSettings::Direct => match load_direnv_environment(&parsed_env, dir).await {
344            Ok(env) => (Some(env), None),
345            Err(err) => (
346                None,
347                <Option<EnvironmentErrorMessage> as From<DirenvError>>::from(err),
348            ),
349        },
350    };
351
352    for (key, value) in direnv_environment.unwrap_or(HashMap::default()) {
353        parsed_env.insert(key, value);
354    }
355
356    (Some(parsed_env), direnv_error)
357}