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