environment.rs

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