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