1use anyhow::{Context as _, bail};
2use futures::{FutureExt, StreamExt as _, channel::mpsc, future::Shared};
3use language::Buffer;
4use remote::RemoteClient;
5use rpc::proto::{self, REMOTE_SERVER_PROJECT_ID};
6use std::{collections::VecDeque, path::Path, sync::Arc};
7use task::{Shell, shell_to_proto};
8use terminal::terminal_settings::TerminalSettings;
9use util::{ResultExt, rel_path::RelPath};
10use worktree::Worktree;
11
12use collections::HashMap;
13use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Task, WeakEntity};
14use settings::Settings as _;
15
16use crate::{
17 project_settings::{DirenvSettings, ProjectSettings},
18 worktree_store::WorktreeStore,
19};
20
21pub struct ProjectEnvironment {
22 cli_environment: Option<HashMap<String, String>>,
23 local_environments: HashMap<(Shell, Arc<Path>), Shared<Task<Option<HashMap<String, String>>>>>,
24 remote_environments: HashMap<(Shell, Arc<Path>), Shared<Task<Option<HashMap<String, String>>>>>,
25 environment_error_messages: VecDeque<String>,
26 environment_error_messages_tx: mpsc::UnboundedSender<String>,
27 worktree_store: WeakEntity<WorktreeStore>,
28 remote_client: Option<WeakEntity<RemoteClient>>,
29 is_remote_project: bool,
30 _tasks: Vec<Task<()>>,
31}
32
33pub enum ProjectEnvironmentEvent {
34 ErrorsUpdated,
35}
36
37impl EventEmitter<ProjectEnvironmentEvent> for ProjectEnvironment {}
38
39impl ProjectEnvironment {
40 pub fn new(
41 cli_environment: Option<HashMap<String, String>>,
42 worktree_store: WeakEntity<WorktreeStore>,
43 remote_client: Option<WeakEntity<RemoteClient>>,
44 is_remote_project: bool,
45 cx: &mut Context<Self>,
46 ) -> Self {
47 let (tx, mut rx) = mpsc::unbounded();
48 let task = cx.spawn(async move |this, cx| {
49 while let Some(message) = rx.next().await {
50 this.update(cx, |this, cx| {
51 this.environment_error_messages.push_back(message);
52 cx.emit(ProjectEnvironmentEvent::ErrorsUpdated);
53 })
54 .ok();
55 }
56 });
57 Self {
58 cli_environment,
59 local_environments: Default::default(),
60 remote_environments: Default::default(),
61 environment_error_messages: Default::default(),
62 environment_error_messages_tx: tx,
63 worktree_store,
64 remote_client,
65 is_remote_project,
66 _tasks: vec![task],
67 }
68 }
69
70 /// Returns the inherited CLI environment, if this project was opened from the Zed CLI.
71 pub(crate) fn get_cli_environment(&self) -> Option<HashMap<String, String>> {
72 if cfg!(any(test, feature = "test-support")) {
73 return Some(HashMap::default());
74 }
75 if let Some(mut env) = self.cli_environment.clone() {
76 set_origin_marker(&mut env, EnvironmentOrigin::Cli);
77 Some(env)
78 } else {
79 None
80 }
81 }
82
83 pub fn buffer_environment(
84 &mut self,
85 buffer: &Entity<Buffer>,
86 worktree_store: &Entity<WorktreeStore>,
87 cx: &mut Context<Self>,
88 ) -> Shared<Task<Option<HashMap<String, String>>>> {
89 if let Some(cli_environment) = self.get_cli_environment() {
90 log::debug!("using project environment variables from CLI");
91 return Task::ready(Some(cli_environment)).shared();
92 }
93
94 let Some(worktree) = buffer
95 .read(cx)
96 .file()
97 .map(|f| f.worktree_id(cx))
98 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
99 else {
100 return Task::ready(None).shared();
101 };
102 self.worktree_environment(worktree, cx)
103 }
104
105 pub fn worktree_environment(
106 &mut self,
107 worktree: Entity<Worktree>,
108 cx: &mut App,
109 ) -> Shared<Task<Option<HashMap<String, String>>>> {
110 if let Some(cli_environment) = self.get_cli_environment() {
111 log::debug!("using project environment variables from CLI");
112 return Task::ready(Some(cli_environment)).shared();
113 }
114
115 let worktree = worktree.read(cx);
116 let mut abs_path = worktree.abs_path();
117 if worktree.is_single_file() {
118 let Some(parent) = abs_path.parent() else {
119 return Task::ready(None).shared();
120 };
121 abs_path = parent.into();
122 }
123
124 let remote_client = self.remote_client.as_ref().and_then(|it| it.upgrade());
125 match remote_client {
126 Some(remote_client) => remote_client.clone().read(cx).shell().map(|shell| {
127 self.remote_directory_environment(
128 &Shell::Program(shell),
129 abs_path,
130 remote_client,
131 cx,
132 )
133 }),
134 None if self.is_remote_project => {
135 Some(self.local_directory_environment(&Shell::System, abs_path, cx))
136 }
137 None => Some({
138 let shell = TerminalSettings::get(
139 Some(settings::SettingsLocation {
140 worktree_id: worktree.id(),
141 path: RelPath::empty(),
142 }),
143 cx,
144 )
145 .shell
146 .clone();
147
148 self.local_directory_environment(&shell, abs_path, cx)
149 }),
150 }
151 .unwrap_or_else(|| Task::ready(None).shared())
152 }
153
154 pub fn directory_environment(
155 &mut self,
156 abs_path: Arc<Path>,
157 cx: &mut App,
158 ) -> Shared<Task<Option<HashMap<String, String>>>> {
159 let remote_client = self.remote_client.as_ref().and_then(|it| it.upgrade());
160 match remote_client {
161 Some(remote_client) => remote_client.clone().read(cx).shell().map(|shell| {
162 self.remote_directory_environment(
163 &Shell::Program(shell),
164 abs_path,
165 remote_client,
166 cx,
167 )
168 }),
169 None if self.is_remote_project => {
170 Some(self.local_directory_environment(&Shell::System, abs_path, cx))
171 }
172 None => self
173 .worktree_store
174 .read_with(cx, |worktree_store, cx| {
175 worktree_store.find_worktree(&abs_path, cx)
176 })
177 .ok()
178 .map(|worktree| {
179 let shell = terminal::terminal_settings::TerminalSettings::get(
180 worktree
181 .as_ref()
182 .map(|(worktree, path)| settings::SettingsLocation {
183 worktree_id: worktree.read(cx).id(),
184 path: &path,
185 }),
186 cx,
187 )
188 .shell
189 .clone();
190
191 self.local_directory_environment(&shell, abs_path, cx)
192 }),
193 }
194 .unwrap_or_else(|| Task::ready(None).shared())
195 }
196
197 /// Returns the project environment, if possible.
198 /// If the project was opened from the CLI, then the inherited CLI environment is returned.
199 /// If it wasn't opened from the CLI, and an absolute path is given, then a shell is spawned in
200 /// that directory, to get environment variables as if the user has `cd`'d there.
201 pub fn local_directory_environment(
202 &mut self,
203 shell: &Shell,
204 abs_path: Arc<Path>,
205 cx: &mut App,
206 ) -> Shared<Task<Option<HashMap<String, String>>>> {
207 if let Some(cli_environment) = self.get_cli_environment() {
208 log::debug!("using project environment variables from CLI");
209 return Task::ready(Some(cli_environment)).shared();
210 }
211
212 self.local_environments
213 .entry((shell.clone(), abs_path.clone()))
214 .or_insert_with(|| {
215 let load_direnv = ProjectSettings::get_global(cx).load_direnv.clone();
216 let shell = shell.clone();
217 let tx = self.environment_error_messages_tx.clone();
218 cx.spawn(async move |cx| {
219 let mut shell_env = cx
220 .background_spawn(load_directory_shell_environment(
221 shell,
222 abs_path.clone(),
223 load_direnv,
224 tx,
225 ))
226 .await
227 .log_err();
228
229 if let Some(shell_env) = shell_env.as_mut() {
230 let path = shell_env
231 .get("PATH")
232 .map(|path| path.as_str())
233 .unwrap_or_default();
234 log::debug!(
235 "using project environment variables shell launched in {:?}. PATH={:?}",
236 abs_path,
237 path
238 );
239
240 set_origin_marker(shell_env, EnvironmentOrigin::WorktreeShell);
241 }
242
243 shell_env
244 })
245 .shared()
246 })
247 .clone()
248 }
249
250 pub fn remote_directory_environment(
251 &mut self,
252 shell: &Shell,
253 abs_path: Arc<Path>,
254 remote_client: Entity<RemoteClient>,
255 cx: &mut App,
256 ) -> Shared<Task<Option<HashMap<String, String>>>> {
257 if cfg!(any(test, feature = "test-support")) {
258 return Task::ready(Some(HashMap::default())).shared();
259 }
260
261 self.remote_environments
262 .entry((shell.clone(), abs_path.clone()))
263 .or_insert_with(|| {
264 let response =
265 remote_client
266 .read(cx)
267 .proto_client()
268 .request(proto::GetDirectoryEnvironment {
269 project_id: REMOTE_SERVER_PROJECT_ID,
270 shell: Some(shell_to_proto(shell.clone())),
271 directory: abs_path.to_string_lossy().to_string(),
272 });
273 cx.background_spawn(async move {
274 let environment = response.await.log_err()?;
275 Some(environment.environment.into_iter().collect())
276 })
277 .shared()
278 })
279 .clone()
280 }
281
282 pub fn peek_environment_error(&self) -> Option<&String> {
283 self.environment_error_messages.front()
284 }
285
286 pub fn pop_environment_error(&mut self) -> Option<String> {
287 self.environment_error_messages.pop_front()
288 }
289}
290
291fn set_origin_marker(env: &mut HashMap<String, String>, origin: EnvironmentOrigin) {
292 env.insert(ZED_ENVIRONMENT_ORIGIN_MARKER.to_string(), origin.into());
293}
294
295const ZED_ENVIRONMENT_ORIGIN_MARKER: &str = "ZED_ENVIRONMENT";
296
297enum EnvironmentOrigin {
298 Cli,
299 WorktreeShell,
300}
301
302impl From<EnvironmentOrigin> for String {
303 fn from(val: EnvironmentOrigin) -> Self {
304 match val {
305 EnvironmentOrigin::Cli => "cli".into(),
306 EnvironmentOrigin::WorktreeShell => "worktree-shell".into(),
307 }
308 }
309}
310
311async fn load_directory_shell_environment(
312 shell: Shell,
313 abs_path: Arc<Path>,
314 load_direnv: DirenvSettings,
315 tx: mpsc::UnboundedSender<String>,
316) -> anyhow::Result<HashMap<String, String>> {
317 let meta = smol::fs::metadata(&abs_path).await.with_context(|| {
318 tx.unbounded_send(format!("Failed to open {}", abs_path.display()))
319 .ok();
320 format!("stat {abs_path:?}")
321 })?;
322
323 let dir = if meta.is_dir() {
324 abs_path.clone()
325 } else {
326 abs_path
327 .parent()
328 .with_context(|| {
329 tx.unbounded_send(format!("Failed to open {}", abs_path.display()))
330 .ok();
331 format!("getting parent of {abs_path:?}")
332 })?
333 .into()
334 };
335
336 let (shell, args) = shell.program_and_args();
337 let mut envs = util::shell_env::capture(shell.clone(), args, abs_path)
338 .await
339 .with_context(|| {
340 tx.unbounded_send("Failed to load environment variables".into())
341 .ok();
342 format!("capturing shell environment with {shell:?}")
343 })?;
344
345 if cfg!(target_os = "windows")
346 && let Some(path) = envs.remove("Path")
347 {
348 // windows env vars are case-insensitive, so normalize the path var
349 // so we can just assume `PATH` in other places
350 envs.insert("PATH".into(), path);
351 }
352 // If the user selects `Direct` for direnv, it would set an environment
353 // variable that later uses to know that it should not run the hook.
354 // We would include in `.envs` call so it is okay to run the hook
355 // even if direnv direct mode is enabled.
356 let direnv_environment = match load_direnv {
357 DirenvSettings::ShellHook => None,
358 // Note: direnv is not available on Windows, so we skip direnv processing
359 // and just return the shell environment
360 DirenvSettings::Direct if cfg!(target_os = "windows") => None,
361 DirenvSettings::Direct => load_direnv_environment(&envs, &dir)
362 .await
363 .with_context(|| {
364 tx.unbounded_send("Failed to load direnv environment".into())
365 .ok();
366 "load direnv environment"
367 })
368 .log_err(),
369 };
370 if let Some(direnv_environment) = direnv_environment {
371 for (key, value) in direnv_environment {
372 if let Some(value) = value {
373 envs.insert(key, value);
374 } else {
375 envs.remove(&key);
376 }
377 }
378 }
379
380 Ok(envs)
381}
382
383async fn load_direnv_environment(
384 env: &HashMap<String, String>,
385 dir: &Path,
386) -> anyhow::Result<HashMap<String, Option<String>>> {
387 let Some(direnv_path) = which::which("direnv").ok() else {
388 return Ok(HashMap::default());
389 };
390
391 let args = &["export", "json"];
392 let direnv_output = smol::process::Command::new(&direnv_path)
393 .args(args)
394 .envs(env)
395 .env("TERM", "dumb")
396 .current_dir(dir)
397 .output()
398 .await
399 .context("running direnv")?;
400
401 if !direnv_output.status.success() {
402 bail!(
403 "Loading direnv environment failed ({}), stderr: {}",
404 direnv_output.status,
405 String::from_utf8_lossy(&direnv_output.stderr)
406 );
407 }
408
409 let output = String::from_utf8_lossy(&direnv_output.stdout);
410 if output.is_empty() {
411 // direnv outputs nothing when it has no changes to apply to environment variables
412 return Ok(HashMap::default());
413 }
414
415 serde_json::from_str(&output).context("parsing direnv json")
416}