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