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