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