Load shell environment once at startup

src/agent.rs

@@ -13,44 +13,168 @@ use crate::settings::{EngineSettings, ReasoningMode, TurnSettings};
 use serde::Deserialize;
 use serde_json::{Map, Value};
 use std::collections::HashMap;
-use std::ffi::OsString;
+use std::ffi::{OsStr, OsString};
 use std::fs::{self, File};
+use std::io::Read;
+use std::os::unix::ffi::OsStringExt;
 use std::os::unix::process::CommandExt;
 use std::path::{Path, PathBuf};
 use std::process::{Child, Command, Stdio};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender, TryRecvError};
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, OnceLock};
 use std::thread;
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
 const MAX_FILE_BYTES: u64 = 16 * 1024 * 1024;
-const SHELL_ENV_ALLOWLIST: &[&str] = &[
+const SHELL_ENV_TIMEOUT: Duration = Duration::from_secs(5);
-    "PATH",
+const SHELL_ENV_SENTINEL: &[u8] = b"\0DS4_ENV\0";
-    "HOME",
+static USER_SHELL_ENVIRONMENT: OnceLock<Vec<(OsString, OsString)>> = OnceLock::new();
-    "USER",
-    "LOGNAME",
-    "SHELL",
-    "TMPDIR",
-    "LANG",
-    "LC_ALL",
-    "TERM",
-    "DEVELOPER_DIR",
-    "SDKROOT",
-    "MACOSX_DEPLOYMENT_TARGET",
-    "RUSTUP_TOOLCHAIN",
-];
 
 fn user_shell() -> OsString {
     std::env::var_os("SHELL").unwrap_or_else(|| OsString::from("/bin/sh"))
 }
 
-fn shell_process(shell: &std::ffi::OsStr, command: &str) -> Command {
+enum ShellEnvironmentProbe {
+    Loaded(Vec<(OsString, OsString)>),
+    Timeout,
+    Unavailable,
+}
+
+pub(crate) fn initialize_shell_environment() {
+    let shell = user_shell();
+    let environment = match load_shell_environment(
+        &shell,
+        std::env::var_os("HOME").as_deref(),
+        std::env::var_os("ZDOTDIR").as_deref(),
+    ) {
+        Some(environment) => environment,
+        None => {
+            eprintln!(
+                "DS4Server: could not load the environment from {}; using the app environment",
+                shell.to_string_lossy()
+            );
+            std::env::vars_os().collect()
+        }
+    };
+    let _ = USER_SHELL_ENVIRONMENT.set(environment);
+}
+
+fn load_shell_environment(
+    shell: &OsStr,
+    home: Option<&OsStr>,
+    zdotdir: Option<&OsStr>,
+) -> Option<Vec<(OsString, OsString)>> {
+    match probe_shell_environment(shell, "-il", home, zdotdir, SHELL_ENV_TIMEOUT) {
+        ShellEnvironmentProbe::Loaded(environment) => Some(environment),
+        ShellEnvironmentProbe::Timeout => None,
+        ShellEnvironmentProbe::Unavailable => {
+            match probe_shell_environment(shell, "-l", home, zdotdir, SHELL_ENV_TIMEOUT) {
+                ShellEnvironmentProbe::Loaded(environment) => Some(environment),
+                ShellEnvironmentProbe::Timeout | ShellEnvironmentProbe::Unavailable => None,
+            }
+        }
+    }
+}
+
+fn probe_shell_environment(
+    shell: &OsStr,
+    mode: &str,
+    home: Option<&OsStr>,
+    zdotdir: Option<&OsStr>,
+    timeout: Duration,
+) -> ShellEnvironmentProbe {
+    let mut process = Command::new(shell);
+    process
+        .arg(mode)
+        .arg("-c")
+        .arg("printf '\\0DS4_ENV\\0'; /usr/bin/env -0")
+        .stdin(Stdio::null())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::null())
+        .process_group(0);
+    if let Some(home) = home {
+        process.env("HOME", home);
+    }
+    if let Some(zdotdir) = zdotdir {
+        process.env("ZDOTDIR", zdotdir);
+    }
+    let Ok(mut child) = process.spawn() else {
+        return ShellEnvironmentProbe::Unavailable;
+    };
+    let Some(mut stdout) = child.stdout.take() else {
+        stop_process_group(&mut child);
+        return ShellEnvironmentProbe::Unavailable;
+    };
+    let (output_sender, output_receiver) = mpsc::channel();
+    thread::spawn(move || {
+        let mut output = Vec::new();
+        let _ = stdout.read_to_end(&mut output);
+        let _ = output_sender.send(output);
+    });
+    let deadline = Instant::now() + timeout;
+    loop {
+        match child.try_wait() {
+            Ok(Some(status)) => {
+                let Ok(output) = output_receiver
+                    .recv_timeout(deadline.saturating_duration_since(Instant::now()))
+                else {
+                    signal_process_group(child.id(), "-KILL");
+                    return ShellEnvironmentProbe::Timeout;
+                };
+                return if status.success() {
+                    parse_shell_environment(output).map_or(
+                        ShellEnvironmentProbe::Unavailable,
+                        ShellEnvironmentProbe::Loaded,
+                    )
+                } else {
+                    ShellEnvironmentProbe::Unavailable
+                };
+            }
+            Err(_) => {
+                stop_process_group(&mut child);
+                return ShellEnvironmentProbe::Unavailable;
+            }
+            Ok(None) if Instant::now() < deadline => thread::sleep(Duration::from_millis(10)),
+            Ok(None) => {
+                stop_process_group(&mut child);
+                return ShellEnvironmentProbe::Timeout;
+            }
+        }
+    }
+}
+
+fn parse_shell_environment(output: Vec<u8>) -> Option<Vec<(OsString, OsString)>> {
+    let start = output
+        .windows(SHELL_ENV_SENTINEL.len())
+        .rposition(|window| window == SHELL_ENV_SENTINEL)?
+        + SHELL_ENV_SENTINEL.len();
+    let environment = output[start..]
+        .split(|byte| *byte == 0)
+        .filter_map(|entry| {
+            let equals = entry.iter().position(|byte| *byte == b'=')?;
+            Some((
+                OsString::from_vec(entry[..equals].to_vec()),
+                OsString::from_vec(entry[equals + 1..].to_vec()),
+            ))
+        })
+        .collect::<Vec<_>>();
+    (!environment.is_empty()).then_some(environment)
+}
+
+fn shell_environment() -> Vec<(OsString, OsString)> {
+    USER_SHELL_ENVIRONMENT
+        .get()
+        .cloned()
+        .unwrap_or_else(|| std::env::vars_os().collect())
+}
+
+fn shell_process(shell: &OsStr, command: &str) -> Command {
     let mut process = Command::new("/bin/sh");
     process
         .arg("-c")
         .arg(format!(
-            "ulimit -f {}; exec \"$1\" -l -i -c \"$2\"",
+            "ulimit -f {}; exec \"$1\" -c \"$2\"",
             MAX_FILE_BYTES / 512
         ))
         .arg("ds4-agent")
@@ -829,12 +953,8 @@ impl Tools {
             .stdout(stdout)
             .stderr(stderr)
             .process_group(0)
-            .env_clear();
+            .env_clear()
-        for name in SHELL_ENV_ALLOWLIST {
+            .envs(shell_environment());
-            if let Some(value) = std::env::var_os(name) {
-                process.env(name, value);
-            }
-        }
         process.env("PWD", &self.root);
         let child = process
             .spawn()
@@ -1743,30 +1863,33 @@ fn unique_match(data: &str, needle: &str, label: &str) -> Result<usize, String>
 }
 
 fn stop_job(job: &mut BashJob) {
-    if job.child.try_wait().ok().flatten().is_some() {
+    stop_process_group(&mut job.child);
+}
+
+fn stop_process_group(child: &mut Child) {
+    if child.try_wait().ok().flatten().is_some() {
         return;
     }
-    let pid = job.child.id();
+    let pid = child.id();
-    let _ = Command::new("/bin/kill")
+    signal_process_group(pid, "-TERM");
-        .arg("-TERM")
-        .arg(format!("-{pid}"))
-        .stdout(Stdio::null())
-        .stderr(Stdio::null())
-        .status();
     let deadline = Instant::now() + Duration::from_secs(1);
     while Instant::now() < deadline {
-        if job.child.try_wait().ok().flatten().is_some() {
+        if child.try_wait().ok().flatten().is_some() {
             return;
         }
         thread::sleep(Duration::from_millis(20));
     }
+    signal_process_group(pid, "-KILL");
+    let _ = child.wait();
+}
+
+fn signal_process_group(pid: u32, signal: &str) {
     let _ = Command::new("/bin/kill")
-        .arg("-KILL")
+        .arg(signal)
         .arg(format!("-{pid}"))
         .stdout(Stdio::null())
         .stderr(Stdio::null())
         .status();
-    let _ = job.child.wait();
 }
 
 fn string<'a>(call: &'a ToolCall, name: &str) -> Option<&'a str> {
@@ -2007,14 +2130,11 @@ mod tests {
         ] {
             assert!(risky_shell_reason(command, root).is_some(), "{command}");
         }
-        assert!(!SHELL_ENV_ALLOWLIST.contains(&"GITHUB_TOKEN"));
-        assert!(!SHELL_ENV_ALLOWLIST.contains(&"AWS_SECRET_ACCESS_KEY"));
-        assert!(!SHELL_ENV_ALLOWLIST.contains(&"SSH_AUTH_SOCK"));
     }
 
     #[cfg(target_os = "macos")]
     #[test]
-    fn shell_commands_load_the_user_login_and_interactive_environment() {
+    fn shell_environment_is_loaded_once_from_login_and_interactive_startup_files() {
         let directory = std::env::temp_dir().join(format!(
             "ds4-agent-shell-{}",
             SystemTime::now()
@@ -2025,18 +2145,88 @@ mod tests {
         fs::create_dir_all(&directory).unwrap();
         fs::write(
             directory.join(".zprofile"),
-            "export PATH=\"$HOME/homebrew/bin:$PATH\"\nexport DS4_LOGIN_PROFILE=loaded\n",
+            "export PATH=\"$HOME/login-bin:$PATH\"\nexport DS4_LOGIN=loaded\n",
         )
         .unwrap();
         fs::write(
             directory.join(".zshrc"),
-            "export DS4_INTERACTIVE_PROFILE=loaded\n",
+            "export PATH=\"$HOME/interactive-bin:$PATH\"\nexport DS4_INTERACTIVE=loaded\n",
         )
         .unwrap();
 
+        let environment = load_shell_environment(
+            OsStr::new("/bin/zsh"),
+            Some(directory.as_os_str()),
+            Some(directory.as_os_str()),
+        )
+        .unwrap()
+        .into_iter()
+        .collect::<HashMap<_, _>>();
+        assert_eq!(
+            environment.get(OsStr::new("DS4_LOGIN")),
+            Some(&OsString::from("loaded"))
+        );
+        assert_eq!(
+            environment.get(OsStr::new("DS4_INTERACTIVE")),
+            Some(&OsString::from("loaded"))
+        );
+        let path = environment
+            .get(OsStr::new("PATH"))
+            .unwrap()
+            .to_string_lossy();
+        assert!(path.starts_with(&format!(
+            "{}/interactive-bin:{}/login-bin:",
+            directory.display(),
+            directory.display()
+        )));
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[cfg(target_os = "macos")]
+    #[test]
+    fn shell_environment_probe_stops_blocked_startup_files() {
+        let directory = std::env::temp_dir().join(format!(
+            "ds4-agent-shell-timeout-{}",
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        fs::create_dir_all(&directory).unwrap();
+        fs::write(directory.join(".zshrc"), "sleep 10\n").unwrap();
+
+        let started = Instant::now();
+        assert!(matches!(
+            probe_shell_environment(
+                OsStr::new("/bin/zsh"),
+                "-il",
+                Some(directory.as_os_str()),
+                Some(directory.as_os_str()),
+                Duration::from_millis(50),
+            ),
+            ShellEnvironmentProbe::Timeout
+        ));
+        assert!(started.elapsed() < Duration::from_secs(2));
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[cfg(target_os = "macos")]
+    #[test]
+    fn shell_commands_skip_user_startup_files() {
+        let directory = std::env::temp_dir().join(format!(
+            "ds4-agent-shell-startup-{}",
+            SystemTime::now()
+                .duration_since(UNIX_EPOCH)
+                .unwrap()
+                .as_nanos()
+        ));
+        fs::create_dir_all(&directory).unwrap();
+        fs::write(directory.join(".zprofile"), "export DS4_PROFILE=login\n").unwrap();
+        fs::write(directory.join(".zshrc"), "export DS4_RC=interactive\n").unwrap();
+
         let mut process = shell_process(
-            std::ffi::OsStr::new("/bin/zsh"),
+            OsStr::new("/bin/zsh"),
-            "printf '%s|%s|%s' \"$DS4_LOGIN_PROFILE\" \"$DS4_INTERACTIVE_PROFILE\" \"$PATH\"",
+            "printf '%s|%s' \"$DS4_PROFILE\" \"$DS4_RC\"",
         );
         let output = process
             .env_clear()
@@ -2048,11 +2238,7 @@ mod tests {
             .unwrap();
         assert!(output.status.success());
         let stdout = String::from_utf8(output.stdout).unwrap();
-        assert!(stdout.starts_with("loaded|loaded|"), "{stdout}");
+        assert_eq!(stdout, "|");
-        assert!(
-            stdout.contains(&format!("{}/homebrew/bin", directory.display())),
-            "{stdout}"
-        );
         fs::remove_dir_all(directory).unwrap();
     }
 

src/main.rs

@@ -36,6 +36,7 @@ fn main() -> iced::Result {
         }
         return Ok(());
     }
+    agent::initialize_shell_environment();
     #[cfg(target_os = "macos")]
     if let Err(error) = engine::configure_metal_sources() {
         eprintln!("DS4Server: {error}");