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