environment.rs

  1use anyhow::{Context as _, bail};
  2use futures::{FutureExt, StreamExt as _, channel::mpsc, future::Shared};
  3use language::Buffer;
  4use remote::RemoteClient;
  5use rpc::proto::{self, REMOTE_SERVER_PROJECT_ID};
  6use std::{collections::VecDeque, path::Path, sync::Arc};
  7use task::{Shell, shell_to_proto};
  8use util::{ResultExt, command::new_command};
  9use worktree::Worktree;
 10
 11use collections::HashMap;
 12use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Task, WeakEntity};
 13use settings::Settings as _;
 14
 15use crate::{
 16    project_settings::{DirenvSettings, ProjectSettings},
 17    worktree_store::WorktreeStore,
 18};
 19
 20pub struct ProjectEnvironment {
 21    cli_environment: Option<HashMap<String, String>>,
 22    local_environments: HashMap<(Shell, Arc<Path>), Shared<Task<Option<HashMap<String, String>>>>>,
 23    remote_environments: HashMap<(Shell, Arc<Path>), Shared<Task<Option<HashMap<String, String>>>>>,
 24    environment_error_messages: VecDeque<String>,
 25    environment_error_messages_tx: mpsc::UnboundedSender<String>,
 26    worktree_store: WeakEntity<WorktreeStore>,
 27    remote_client: Option<WeakEntity<RemoteClient>>,
 28    is_remote_project: bool,
 29    _tasks: Vec<Task<()>>,
 30}
 31
 32pub enum ProjectEnvironmentEvent {
 33    ErrorsUpdated,
 34}
 35
 36impl EventEmitter<ProjectEnvironmentEvent> for ProjectEnvironment {}
 37
 38impl ProjectEnvironment {
 39    pub fn new(
 40        cli_environment: Option<HashMap<String, String>>,
 41        worktree_store: WeakEntity<WorktreeStore>,
 42        remote_client: Option<WeakEntity<RemoteClient>>,
 43        is_remote_project: bool,
 44        cx: &mut Context<Self>,
 45    ) -> Self {
 46        let (tx, mut rx) = mpsc::unbounded();
 47        let task = cx.spawn(async move |this, cx| {
 48            while let Some(message) = rx.next().await {
 49                this.update(cx, |this, cx| {
 50                    this.environment_error_messages.push_back(message);
 51                    cx.emit(ProjectEnvironmentEvent::ErrorsUpdated);
 52                })
 53                .ok();
 54            }
 55        });
 56        Self {
 57            cli_environment,
 58            local_environments: Default::default(),
 59            remote_environments: Default::default(),
 60            environment_error_messages: Default::default(),
 61            environment_error_messages_tx: tx,
 62            worktree_store,
 63            remote_client,
 64            is_remote_project,
 65            _tasks: vec![task],
 66        }
 67    }
 68
 69    /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
 70    pub(crate) fn get_cli_environment(&self) -> Option<HashMap<String, String>> {
 71        if cfg!(any(test, feature = "test-support")) {
 72            return Some(HashMap::default());
 73        }
 74        if let Some(mut env) = self.cli_environment.clone() {
 75            set_origin_marker(&mut env, EnvironmentOrigin::Cli);
 76            Some(env)
 77        } else {
 78            None
 79        }
 80    }
 81
 82    pub fn buffer_environment(
 83        &mut self,
 84        buffer: &Entity<Buffer>,
 85        worktree_store: &Entity<WorktreeStore>,
 86        cx: &mut Context<Self>,
 87    ) -> Shared<Task<Option<HashMap<String, String>>>> {
 88        if let Some(cli_environment) = self.get_cli_environment() {
 89            log::debug!("using project environment variables from CLI");
 90            return Task::ready(Some(cli_environment)).shared();
 91        }
 92
 93        let Some(worktree) = buffer
 94            .read(cx)
 95            .file()
 96            .map(|f| f.worktree_id(cx))
 97            .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
 98        else {
 99            return Task::ready(None).shared();
100        };
101        self.worktree_environment(worktree, cx)
102    }
103
104    pub fn worktree_environment(
105        &mut self,
106        worktree: Entity<Worktree>,
107        cx: &mut App,
108    ) -> Shared<Task<Option<HashMap<String, String>>>> {
109        if let Some(cli_environment) = self.get_cli_environment() {
110            log::debug!("using project environment variables from CLI");
111            return Task::ready(Some(cli_environment)).shared();
112        }
113
114        let worktree = worktree.read(cx);
115        let mut abs_path = worktree.abs_path();
116        if worktree.is_single_file() {
117            let Some(parent) = abs_path.parent() else {
118                return Task::ready(None).shared();
119            };
120            abs_path = parent.into();
121        }
122
123        let remote_client = self.remote_client.as_ref().and_then(|it| it.upgrade());
124        match remote_client {
125            Some(remote_client) => remote_client.clone().read(cx).shell().map(|shell| {
126                self.remote_directory_environment(
127                    &Shell::Program(shell),
128                    abs_path,
129                    remote_client,
130                    cx,
131                )
132            }),
133            None if self.is_remote_project => {
134                Some(self.local_directory_environment(&Shell::System, abs_path, cx))
135            }
136            None => Some(self.local_directory_environment(&Shell::System, abs_path, cx)),
137        }
138        .unwrap_or_else(|| Task::ready(None).shared())
139    }
140
141    pub fn directory_environment(
142        &mut self,
143        abs_path: Arc<Path>,
144        cx: &mut App,
145    ) -> Shared<Task<Option<HashMap<String, String>>>> {
146        let remote_client = self.remote_client.as_ref().and_then(|it| it.upgrade());
147        match remote_client {
148            Some(remote_client) => remote_client.clone().read(cx).shell().map(|shell| {
149                self.remote_directory_environment(
150                    &Shell::Program(shell),
151                    abs_path,
152                    remote_client,
153                    cx,
154                )
155            }),
156            None if self.is_remote_project => {
157                Some(self.local_directory_environment(&Shell::System, abs_path, cx))
158            }
159            None => self
160                .worktree_store
161                .read_with(cx, |worktree_store, cx| {
162                    worktree_store.find_worktree(&abs_path, cx)
163                })
164                .ok()
165                .map(|_| self.local_directory_environment(&Shell::System, abs_path, cx)),
166        }
167        .unwrap_or_else(|| Task::ready(None).shared())
168    }
169
170    /// Returns the project environment using the default worktree path.
171    /// This ensures that project-specific environment variables (e.g. from `.envrc`)
172    /// are loaded from the project directory rather than the home directory.
173    pub fn default_environment(
174        &mut self,
175        cx: &mut App,
176    ) -> Shared<Task<Option<HashMap<String, String>>>> {
177        let abs_path = self
178            .worktree_store
179            .read_with(cx, |worktree_store, cx| {
180                crate::Project::default_visible_worktree_paths(worktree_store, cx)
181                    .into_iter()
182                    .next()
183            })
184            .ok()
185            .flatten()
186            .map(|path| Arc::<Path>::from(path))
187            .unwrap_or_else(|| paths::home_dir().as_path().into());
188        self.local_directory_environment(&Shell::System, abs_path, cx)
189    }
190
191    /// Returns the project environment, if possible.
192    /// If the project was opened from the CLI, then the inherited CLI environment is returned.
193    /// If it wasn't opened from the CLI, and an absolute path is given, then a shell is spawned in
194    /// that directory, to get environment variables as if the user has `cd`'d there.
195    pub fn local_directory_environment(
196        &mut self,
197        shell: &Shell,
198        abs_path: Arc<Path>,
199        cx: &mut App,
200    ) -> Shared<Task<Option<HashMap<String, String>>>> {
201        if let Some(cli_environment) = self.get_cli_environment() {
202            log::debug!("using project environment variables from CLI");
203            return Task::ready(Some(cli_environment)).shared();
204        }
205
206        self.local_environments
207            .entry((shell.clone(), abs_path.clone()))
208            .or_insert_with(|| {
209                let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
210                let shell = shell.clone();
211                let tx = self.environment_error_messages_tx.clone();
212                cx.spawn(async move |cx| {
213                    let mut shell_env = match cx
214                        .background_spawn(load_directory_shell_environment(
215                            shell,
216                            abs_path.clone(),
217                            load_direnv,
218                            tx,
219                        ))
220                        .await
221                    {
222                        Ok(shell_env) => Some(shell_env),
223                        Err(e) => {
224                            log::error!(
225                                "Failed to load shell environment for directory {abs_path:?}: {e:#}"
226                            );
227                            None
228                        }
229                    };
230
231                    if let Some(shell_env) = shell_env.as_mut() {
232                        let path = shell_env
233                            .get("PATH")
234                            .map(|path| path.as_str())
235                            .unwrap_or_default();
236                        log::debug!(
237                            "using project environment variables shell launched in {:?}. PATH={:?}",
238                            abs_path,
239                            path
240                        );
241
242                        set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
243                    }
244
245                    shell_env
246                })
247                .shared()
248            })
249            .clone()
250    }
251
252    pub fn remote_directory_environment(
253        &mut self,
254        shell: &Shell,
255        abs_path: Arc<Path>,
256        remote_client: Entity<RemoteClient>,
257        cx: &mut App,
258    ) -> Shared<Task<Option<HashMap<String, String>>>> {
259        if cfg!(any(test, feature = "test-support")) {
260            return Task::ready(Some(HashMap::default())).shared();
261        }
262
263        self.remote_environments
264            .entry((shell.clone(), abs_path.clone()))
265            .or_insert_with(|| {
266                let response =
267                    remote_client
268                        .read(cx)
269                        .proto_client()
270                        .request(proto::GetDirectoryEnvironment {
271                            project_id: REMOTE_SERVER_PROJECT_ID,
272                            shell: Some(shell_to_proto(shell.clone())),
273                            directory: abs_path.to_string_lossy().to_string(),
274                        });
275                cx.background_spawn(async move {
276                    let environment = response.await.log_err()?;
277                    Some(environment.environment.into_iter().collect())
278                })
279                .shared()
280            })
281            .clone()
282    }
283
284    pub fn peek_environment_error(&self) -> Option<&String> {
285        self.environment_error_messages.front()
286    }
287
288    pub fn pop_environment_error(&mut self) -> Option<String> {
289        self.environment_error_messages.pop_front()
290    }
291}
292
293fn set_origin_marker(env: &mut HashMap<String, String>, origin: EnvironmentOrigin) {
294    env.insert(ZED_ENVIRONMENT_ORIGIN_MARKER.to_string(), origin.into());
295}
296
297const ZED_ENVIRONMENT_ORIGIN_MARKER: &str = "ZED_ENVIRONMENT";
298
299enum EnvironmentOrigin {
300    Cli,
301    WorktreeShell,
302}
303
304impl From<EnvironmentOrigin> for String {
305    fn from(val: EnvironmentOrigin) -> Self {
306        match val {
307            EnvironmentOrigin::Cli => "cli".into(),
308            EnvironmentOrigin::WorktreeShell => "worktree-shell".into(),
309        }
310    }
311}
312
313async fn load_directory_shell_environment(
314    shell: Shell,
315    abs_path: Arc<Path>,
316    load_direnv: DirenvSettings,
317    tx: mpsc::UnboundedSender<String>,
318) -> anyhow::Result<HashMap<String, String>> {
319    if let DirenvSettings::Disabled = load_direnv {
320        return Ok(HashMap::default());
321    }
322
323    let meta = smol::fs::metadata(&abs_path).await.with_context(|| {
324        tx.unbounded_send(format!("Failed to open {}", abs_path.display()))
325            .ok();
326        format!("stat {abs_path:?}")
327    })?;
328
329    let dir = if meta.is_dir() {
330        abs_path.clone()
331    } else {
332        abs_path
333            .parent()
334            .with_context(|| {
335                tx.unbounded_send(format!("Failed to open {}", abs_path.display()))
336                    .ok();
337                format!("getting parent of {abs_path:?}")
338            })?
339            .into()
340    };
341
342    let (shell, args) = shell.program_and_args();
343    let mut envs = util::shell_env::capture(shell.clone(), args, abs_path)
344        .await
345        .with_context(|| {
346            tx.unbounded_send("Failed to load environment variables".into())
347                .ok();
348            format!("capturing shell environment with {shell:?}")
349        })?;
350
351    if cfg!(target_os = "windows")
352        && let Some(path) = envs.remove("Path")
353    {
354        // windows env vars are case-insensitive, so normalize the path var
355        // so we can just assume `PATH` in other places
356        envs.insert("PATH".into(), path);
357    }
358    // If the user selects `Direct` for direnv, it would set an environment
359    // variable that later uses to know that it should not run the hook.
360    // We would include in `.envs` call so it is okay to run the hook
361    // even if direnv direct mode is enabled.
362    let direnv_environment = match load_direnv {
363        DirenvSettings::ShellHook => None,
364        DirenvSettings::Disabled => bail!("direnv integration is disabled"),
365        // Note: direnv is not available on Windows, so we skip direnv processing
366        // and just return the shell environment
367        DirenvSettings::Direct if cfg!(target_os = "windows") => None,
368        DirenvSettings::Direct => load_direnv_environment(&envs, &dir)
369            .await
370            .with_context(|| {
371                tx.unbounded_send("Failed to load direnv environment".into())
372                    .ok();
373                "load direnv environment"
374            })
375            .log_err(),
376    };
377    if let Some(direnv_environment) = direnv_environment {
378        for (key, value) in direnv_environment {
379            if let Some(value) = value {
380                envs.insert(key, value);
381            } else {
382                envs.remove(&key);
383            }
384        }
385    }
386
387    Ok(envs)
388}
389
390async fn load_direnv_environment(
391    env: &HashMap<String, String>,
392    dir: &Path,
393) -> anyhow::Result<HashMap<String, Option<String>>> {
394    let Some(direnv_path) = which::which("direnv").ok() else {
395        return Ok(HashMap::default());
396    };
397
398    let args = &["export", "json"];
399    let direnv_output = new_command(&direnv_path)
400        .args(args)
401        .envs(env)
402        .env("TERM", "dumb")
403        .current_dir(dir)
404        .output()
405        .await
406        .context("running direnv")?;
407
408    if !direnv_output.status.success() {
409        bail!(
410            "Loading direnv environment failed ({}), stderr: {}",
411            direnv_output.status,
412            String::from_utf8_lossy(&direnv_output.stderr)
413        );
414    }
415
416    let output = String::from_utf8_lossy(&direnv_output.stdout);
417    if output.is_empty() {
418        // direnv outputs nothing when it has no changes to apply to environment variables
419        return Ok(HashMap::default());
420    }
421
422    serde_json::from_str(&output).context("parsing direnv json")
423}