1use futures::{FutureExt, future::Shared};
2use language::Buffer;
3use std::{path::Path, sync::Arc};
4use task::Shell;
5use util::ResultExt;
6use worktree::Worktree;
7
8use collections::HashMap;
9use gpui::{AppContext as _, Context, Entity, EventEmitter, Task};
10use settings::Settings as _;
11
12use crate::{
13 project_settings::{DirenvSettings, ProjectSettings},
14 worktree_store::WorktreeStore,
15};
16
17pub struct ProjectEnvironment {
18 cli_environment: Option<HashMap<String, String>>,
19 environments: HashMap<Arc<Path>, Shared<Task<Option<HashMap<String, String>>>>>,
20 shell_based_environments:
21 HashMap<(Shell, Arc<Path>), Shared<Task<Option<HashMap<String, String>>>>>,
22 environment_error_messages: HashMap<Arc<Path>, EnvironmentErrorMessage>,
23}
24
25pub enum ProjectEnvironmentEvent {
26 ErrorsUpdated,
27}
28
29impl EventEmitter<ProjectEnvironmentEvent> for ProjectEnvironment {}
30
31impl ProjectEnvironment {
32 pub fn new(cli_environment: Option<HashMap<String, String>>) -> Self {
33 Self {
34 cli_environment,
35 environments: Default::default(),
36 shell_based_environments: Default::default(),
37 environment_error_messages: Default::default(),
38 }
39 }
40
41 /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
42 pub(crate) fn get_cli_environment(&self) -> Option<HashMap<String, String>> {
43 if let Some(mut env) = self.cli_environment.clone() {
44 set_origin_marker(&mut env, EnvironmentOrigin::Cli);
45 Some(env)
46 } else {
47 None
48 }
49 }
50
51 /// Returns an iterator over all pairs `(abs_path, error_message)` of
52 /// environment errors associated with this project environment.
53 pub(crate) fn environment_errors(
54 &self,
55 ) -> impl Iterator<Item = (&Arc<Path>, &EnvironmentErrorMessage)> {
56 self.environment_error_messages.iter()
57 }
58
59 pub(crate) fn remove_environment_error(&mut self, abs_path: &Path, cx: &mut Context<Self>) {
60 self.environment_error_messages.remove(abs_path);
61 cx.emit(ProjectEnvironmentEvent::ErrorsUpdated);
62 }
63
64 pub(crate) fn get_buffer_environment(
65 &mut self,
66 buffer: &Entity<Buffer>,
67 worktree_store: &Entity<WorktreeStore>,
68 cx: &mut Context<Self>,
69 ) -> Shared<Task<Option<HashMap<String, String>>>> {
70 if cfg!(any(test, feature = "test-support")) {
71 return Task::ready(Some(HashMap::default())).shared();
72 }
73
74 if let Some(cli_environment) = self.get_cli_environment() {
75 log::debug!("using project environment variables from CLI");
76 return Task::ready(Some(cli_environment)).shared();
77 }
78
79 let Some(worktree) = buffer
80 .read(cx)
81 .file()
82 .map(|f| f.worktree_id(cx))
83 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
84 else {
85 return Task::ready(None).shared();
86 };
87
88 self.get_worktree_environment(worktree, cx)
89 }
90
91 pub fn get_worktree_environment(
92 &mut self,
93 worktree: Entity<Worktree>,
94 cx: &mut Context<Self>,
95 ) -> Shared<Task<Option<HashMap<String, String>>>> {
96 if cfg!(any(test, feature = "test-support")) {
97 return Task::ready(Some(HashMap::default())).shared();
98 }
99
100 if let Some(cli_environment) = self.get_cli_environment() {
101 log::debug!("using project environment variables from CLI");
102 return Task::ready(Some(cli_environment)).shared();
103 }
104
105 let mut abs_path = worktree.read(cx).abs_path();
106 if !worktree.read(cx).is_local() {
107 log::error!(
108 "attempted to get project environment for a non-local worktree at {abs_path:?}"
109 );
110 return Task::ready(None).shared();
111 } else if worktree.read(cx).is_single_file() {
112 let Some(parent) = abs_path.parent() else {
113 return Task::ready(None).shared();
114 };
115 abs_path = parent.into();
116 }
117
118 self.get_directory_environment(abs_path, cx)
119 }
120
121 /// Returns the project environment, if possible.
122 /// If the project was opened from the CLI, then the inherited CLI environment is returned.
123 /// If it wasn't opened from the CLI, and an absolute path is given, then a shell is spawned in
124 /// that directory, to get environment variables as if the user has `cd`'d there.
125 pub fn get_directory_environment(
126 &mut self,
127 abs_path: Arc<Path>,
128 cx: &mut Context<Self>,
129 ) -> Shared<Task<Option<HashMap<String, String>>>> {
130 if cfg!(any(test, feature = "test-support")) {
131 return Task::ready(Some(HashMap::default())).shared();
132 }
133
134 if let Some(cli_environment) = self.get_cli_environment() {
135 log::debug!("using project environment variables from CLI");
136 return Task::ready(Some(cli_environment)).shared();
137 }
138
139 self.environments
140 .entry(abs_path.clone())
141 .or_insert_with(|| {
142 get_directory_env_impl(&Shell::System, abs_path.clone(), cx).shared()
143 })
144 .clone()
145 }
146
147 /// Returns the project environment, if possible, with the given shell.
148 pub fn get_directory_environment_for_shell(
149 &mut self,
150 shell: &Shell,
151 abs_path: Arc<Path>,
152 cx: &mut Context<Self>,
153 ) -> Shared<Task<Option<HashMap<String, String>>>> {
154 self.shell_based_environments
155 .entry((shell.clone(), abs_path.clone()))
156 .or_insert_with(|| get_directory_env_impl(shell, abs_path.clone(), cx).shared())
157 .clone()
158 }
159}
160
161fn set_origin_marker(env: &mut HashMap<String, String>, origin: EnvironmentOrigin) {
162 env.insert(ZED_ENVIRONMENT_ORIGIN_MARKER.to_string(), origin.into());
163}
164
165const ZED_ENVIRONMENT_ORIGIN_MARKER: &str = "ZED_ENVIRONMENT";
166
167enum EnvironmentOrigin {
168 Cli,
169 WorktreeShell,
170}
171
172impl From<EnvironmentOrigin> for String {
173 fn from(val: EnvironmentOrigin) -> Self {
174 match val {
175 EnvironmentOrigin::Cli => "cli".into(),
176 EnvironmentOrigin::WorktreeShell => "worktree-shell".into(),
177 }
178 }
179}
180
181#[derive(Debug)]
182pub struct EnvironmentErrorMessage(pub String);
183
184impl std::fmt::Display for EnvironmentErrorMessage {
185 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186 write!(f, "{}", self.0)
187 }
188}
189
190impl EnvironmentErrorMessage {
191 #[allow(dead_code)]
192 fn from_str(s: &str) -> Self {
193 Self(String::from(s))
194 }
195}
196
197async fn load_directory_shell_environment(
198 shell: &Shell,
199 abs_path: &Path,
200 load_direnv: &DirenvSettings,
201) -> (
202 Option<HashMap<String, String>>,
203 Option<EnvironmentErrorMessage>,
204) {
205 match smol::fs::metadata(abs_path).await {
206 Ok(meta) => {
207 let dir = if meta.is_dir() {
208 abs_path
209 } else if let Some(parent) = abs_path.parent() {
210 parent
211 } else {
212 return (
213 None,
214 Some(EnvironmentErrorMessage(format!(
215 "Failed to load shell environment in {}: not a directory",
216 abs_path.display()
217 ))),
218 );
219 };
220
221 load_shell_environment(shell, dir, load_direnv).await
222 }
223 Err(err) => (
224 None,
225 Some(EnvironmentErrorMessage(format!(
226 "Failed to load shell environment in {}: {}",
227 abs_path.display(),
228 err
229 ))),
230 ),
231 }
232}
233
234async fn load_shell_environment(
235 shell: &Shell,
236 dir: &Path,
237 load_direnv: &DirenvSettings,
238) -> (
239 Option<HashMap<String, String>>,
240 Option<EnvironmentErrorMessage>,
241) {
242 use crate::direnv::load_direnv_environment;
243 use util::shell_env;
244
245 if cfg!(any(test, feature = "test-support")) {
246 let fake_env = [("ZED_FAKE_TEST_ENV".into(), "true".into())]
247 .into_iter()
248 .collect();
249 (Some(fake_env), None)
250 } else if cfg!(target_os = "windows",) {
251 let (shell, args) = shell.program_and_args();
252 let envs = match shell_env::capture(shell, args, dir).await {
253 Ok(envs) => envs,
254 Err(err) => {
255 util::log_err(&err);
256 return (
257 None,
258 Some(EnvironmentErrorMessage(format!(
259 "Failed to load environment variables: {}",
260 err
261 ))),
262 );
263 }
264 };
265
266 // Note: direnv is not available on Windows, so we skip direnv processing
267 // and just return the shell environment
268 (Some(envs), None)
269 } else {
270 let dir_ = dir.to_owned();
271 let (shell, args) = shell.program_and_args();
272 let mut envs = match shell_env::capture(shell, args, &dir_).await {
273 Ok(envs) => envs,
274 Err(err) => {
275 util::log_err(&err);
276 return (
277 None,
278 Some(EnvironmentErrorMessage::from_str(
279 "Failed to load environment variables. See log for details",
280 )),
281 );
282 }
283 };
284
285 // If the user selects `Direct` for direnv, it would set an environment
286 // variable that later uses to know that it should not run the hook.
287 // We would include in `.envs` call so it is okay to run the hook
288 // even if direnv direct mode is enabled.
289 let (direnv_environment, direnv_error) = match load_direnv {
290 DirenvSettings::ShellHook => (None, None),
291 DirenvSettings::Direct => match load_direnv_environment(&envs, dir).await {
292 Ok(env) => (Some(env), None),
293 Err(err) => (None, err.into()),
294 },
295 };
296 if let Some(direnv_environment) = direnv_environment {
297 for (key, value) in direnv_environment {
298 if let Some(value) = value {
299 envs.insert(key, value);
300 } else {
301 envs.remove(&key);
302 }
303 }
304 }
305
306 (Some(envs), direnv_error)
307 }
308}
309
310fn get_directory_env_impl(
311 shell: &Shell,
312 abs_path: Arc<Path>,
313 cx: &Context<ProjectEnvironment>,
314) -> Task<Option<HashMap<String, String>>> {
315 let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
316
317 let shell = shell.clone();
318 cx.spawn(async move |this, cx| {
319 let (mut shell_env, error_message) = cx
320 .background_spawn({
321 let abs_path = abs_path.clone();
322 async move {
323 load_directory_shell_environment(&shell, &abs_path, &load_direnv).await
324 }
325 })
326 .await;
327
328 if let Some(shell_env) = shell_env.as_mut() {
329 let path = shell_env
330 .get("PATH")
331 .map(|path| path.as_str())
332 .unwrap_or_default();
333 log::info!(
334 "using project environment variables shell launched in {:?}. PATH={:?}",
335 abs_path,
336 path
337 );
338
339 set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
340 }
341
342 if let Some(error) = error_message {
343 this.update(cx, |this, cx| {
344 log::error!("{error}",);
345 this.environment_error_messages.insert(abs_path, error);
346 cx.emit(ProjectEnvironmentEvent::ErrorsUpdated)
347 })
348 .log_err();
349 }
350
351 shell_env
352 })
353}