1use anyhow::{anyhow, Context as _, Result};
2use futures::{future::Shared, FutureExt};
3use std::{
4 path::{Path, PathBuf},
5 sync::Arc,
6};
7use util::{parse_env_output, ResultExt};
8
9use collections::HashMap;
10use gpui::{AppContext, Context, Model, ModelContext, Task};
11use settings::Settings as _;
12use worktree::WorktreeId;
13
14use crate::{
15 project_settings::{DirenvSettings, ProjectSettings},
16 worktree_store::{WorktreeStore, WorktreeStoreEvent},
17};
18
19pub struct ProjectEnvironment {
20 cli_environment: Option<HashMap<String, String>>,
21 get_environment_task: Option<Shared<Task<Option<HashMap<String, String>>>>>,
22 cached_shell_environments: HashMap<WorktreeId, HashMap<String, String>>,
23}
24
25impl ProjectEnvironment {
26 pub fn new(
27 worktree_store: &Model<WorktreeStore>,
28 cli_environment: Option<HashMap<String, String>>,
29 cx: &mut AppContext,
30 ) -> Model<Self> {
31 cx.new_model(|cx| {
32 cx.subscribe(worktree_store, |this: &mut Self, _, event, _| {
33 if let WorktreeStoreEvent::WorktreeRemoved(_, id) = event {
34 this.remove_worktree_environment(*id);
35 }
36 })
37 .detach();
38
39 Self {
40 cli_environment,
41 get_environment_task: None,
42 cached_shell_environments: Default::default(),
43 }
44 })
45 }
46
47 #[cfg(any(test, feature = "test-support"))]
48 pub(crate) fn set_cached(
49 &mut self,
50 shell_environments: &[(WorktreeId, HashMap<String, String>)],
51 ) {
52 self.cached_shell_environments = shell_environments
53 .iter()
54 .cloned()
55 .collect::<HashMap<_, _>>();
56 }
57
58 pub(crate) fn remove_worktree_environment(&mut self, worktree_id: WorktreeId) {
59 self.cached_shell_environments.remove(&worktree_id);
60 }
61
62 /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
63 pub(crate) fn get_cli_environment(&self) -> Option<HashMap<String, String>> {
64 if let Some(mut env) = self.cli_environment.clone() {
65 set_origin_marker(&mut env, EnvironmentOrigin::Cli);
66 Some(env)
67 } else {
68 None
69 }
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: &ModelContext<Self>,
82 ) -> Shared<Task<Option<HashMap<String, String>>>> {
83 if let Some(task) = self.get_environment_task.as_ref() {
84 task.clone()
85 } else {
86 let task = self
87 .build_environment_task(worktree_id, worktree_abs_path, cx)
88 .shared();
89
90 self.get_environment_task = Some(task.clone());
91 task
92 }
93 }
94
95 fn build_environment_task(
96 &mut self,
97 worktree_id: Option<WorktreeId>,
98 worktree_abs_path: Option<Arc<Path>>,
99 cx: &ModelContext<Self>,
100 ) -> Task<Option<HashMap<String, String>>> {
101 let worktree = worktree_id.zip(worktree_abs_path);
102
103 let cli_environment = self.get_cli_environment();
104 if cli_environment.is_some() {
105 Task::ready(cli_environment)
106 } else if let Some((worktree_id, worktree_abs_path)) = worktree {
107 self.get_worktree_env(worktree_id, worktree_abs_path, cx)
108 } else {
109 Task::ready(None)
110 }
111 }
112
113 fn get_worktree_env(
114 &mut self,
115 worktree_id: WorktreeId,
116 worktree_abs_path: Arc<Path>,
117 cx: &ModelContext<Self>,
118 ) -> Task<Option<HashMap<String, String>>> {
119 let cached_env = self.cached_shell_environments.get(&worktree_id).cloned();
120 if let Some(env) = cached_env {
121 Task::ready(Some(env))
122 } else {
123 let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
124
125 cx.spawn(|this, mut cx| async move {
126 let mut shell_env = cx
127 .background_executor()
128 .spawn({
129 let cwd = worktree_abs_path.clone();
130 async move { load_shell_environment(&cwd, &load_direnv).await }
131 })
132 .await
133 .ok();
134
135 if let Some(shell_env) = shell_env.as_mut() {
136 this.update(&mut cx, |this, _| {
137 this.cached_shell_environments
138 .insert(worktree_id, shell_env.clone())
139 })
140 .log_err();
141
142 set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
143 }
144
145 shell_env
146 })
147 }
148 }
149}
150
151fn set_origin_marker(env: &mut HashMap<String, String>, origin: EnvironmentOrigin) {
152 env.insert(ZED_ENVIRONMENT_ORIGIN_MARKER.to_string(), origin.into());
153}
154
155const ZED_ENVIRONMENT_ORIGIN_MARKER: &str = "ZED_ENVIRONMENT";
156
157enum EnvironmentOrigin {
158 Cli,
159 WorktreeShell,
160}
161
162impl From<EnvironmentOrigin> for String {
163 fn from(val: EnvironmentOrigin) -> Self {
164 match val {
165 EnvironmentOrigin::Cli => "cli".into(),
166 EnvironmentOrigin::WorktreeShell => "worktree-shell".into(),
167 }
168 }
169}
170
171async fn load_shell_environment(
172 dir: &Path,
173 load_direnv: &DirenvSettings,
174) -> Result<HashMap<String, String>> {
175 let direnv_environment = match load_direnv {
176 DirenvSettings::ShellHook => None,
177 DirenvSettings::Direct => load_direnv_environment(dir).await?,
178 }
179 .unwrap_or(HashMap::default());
180
181 let marker = "ZED_SHELL_START";
182 let shell = std::env::var("SHELL").context(
183 "SHELL environment variable is not assigned so we can't source login environment variables",
184 )?;
185
186 // What we're doing here is to spawn a shell and then `cd` into
187 // the project directory to get the env in there as if the user
188 // `cd`'d into it. We do that because tools like direnv, asdf, ...
189 // hook into `cd` and only set up the env after that.
190 //
191 // If the user selects `Direct` for direnv, it would set an environment
192 // variable that later uses to know that it should not run the hook.
193 // We would include in `.envs` call so it is okay to run the hook
194 // even if direnv direct mode is enabled.
195 //
196 // In certain shells we need to execute additional_command in order to
197 // trigger the behavior of direnv, etc.
198 //
199 //
200 // The `exit 0` is the result of hours of debugging, trying to find out
201 // why running this command here, without `exit 0`, would mess
202 // up signal process for our process so that `ctrl-c` doesn't work
203 // anymore.
204 //
205 // We still don't know why `$SHELL -l -i -c '/usr/bin/env -0'` would
206 // do that, but it does, and `exit 0` helps.
207 let additional_command = PathBuf::from(&shell)
208 .file_name()
209 .and_then(|f| f.to_str())
210 .and_then(|shell| match shell {
211 "fish" => Some("emit fish_prompt;"),
212 _ => None,
213 });
214
215 let command = format!(
216 "cd '{}';{} printf '%s' {marker}; /usr/bin/env; exit 0;",
217 dir.display(),
218 additional_command.unwrap_or("")
219 );
220
221 let output = smol::process::Command::new(&shell)
222 .args(["-l", "-i", "-c", &command])
223 .envs(direnv_environment)
224 .output()
225 .await
226 .context("failed to spawn login shell to source login environment variables")?;
227
228 anyhow::ensure!(
229 output.status.success(),
230 "login shell exited with error {:?}",
231 output.status
232 );
233
234 let stdout = String::from_utf8_lossy(&output.stdout);
235 let env_output_start = stdout.find(marker).ok_or_else(|| {
236 anyhow!(
237 "failed to parse output of `env` command in login shell: {}",
238 stdout
239 )
240 })?;
241
242 let mut parsed_env = HashMap::default();
243 let env_output = &stdout[env_output_start + marker.len()..];
244
245 parse_env_output(env_output, |key, value| {
246 parsed_env.insert(key, value);
247 });
248
249 Ok(parsed_env)
250}
251
252async fn load_direnv_environment(dir: &Path) -> Result<Option<HashMap<String, String>>> {
253 let Ok(direnv_path) = which::which("direnv") else {
254 return Ok(None);
255 };
256
257 let direnv_output = smol::process::Command::new(direnv_path)
258 .args(["export", "json"])
259 .current_dir(dir)
260 .output()
261 .await
262 .context("failed to spawn direnv to get local environment variables")?;
263
264 anyhow::ensure!(
265 direnv_output.status.success(),
266 "direnv exited with error {:?}",
267 direnv_output.status
268 );
269
270 let output = String::from_utf8_lossy(&direnv_output.stdout);
271 if output.is_empty() {
272 return Ok(None);
273 }
274
275 Ok(Some(
276 serde_json::from_str(&output).context("failed to parse direnv output")?,
277 ))
278}