Просмотр исходного кода

fix(desktop): restore shell path env for desktop sidecar (#15211)

Shoubhit Dash 1 месяц назад
Родитель
Сommit
bb8a1718a6
1 измененных файлов с 184 добавлено и 1 удалено
  1. 184 1
      packages/desktop/src-tauri/src/cli.rs

+ 184 - 1
packages/desktop/src-tauri/src/cli.rs

@@ -4,10 +4,13 @@ use process_wrap::tokio::CommandWrap;
 use process_wrap::tokio::ProcessGroup;
 #[cfg(windows)]
 use process_wrap::tokio::{CommandWrapper, JobObject, KillOnDrop};
+use std::collections::HashMap;
 #[cfg(unix)]
 use std::os::unix::process::ExitStatusExt;
+use std::path::Path;
+use std::process::Stdio;
 use std::sync::Arc;
-use std::{process::Stdio, time::Duration};
+use std::time::{Duration, Instant};
 use tauri::{AppHandle, Manager, path::BaseDirectory};
 use tauri_specta::Event;
 use tokio::{
@@ -39,6 +42,7 @@ impl CommandWrapper for WinCreationFlags {
 
 const CLI_INSTALL_DIR: &str = ".opencode/bin";
 const CLI_BINARY_NAME: &str = "opencode";
+const SHELL_ENV_TIMEOUT: Duration = Duration::from_secs(5);
 
 #[derive(serde::Deserialize, Debug)]
 pub struct ServerConfig {
@@ -232,6 +236,133 @@ fn shell_escape(input: &str) -> String {
     escaped
 }
 
+fn parse_shell_env(stdout: &[u8]) -> HashMap<String, String> {
+    String::from_utf8_lossy(stdout)
+        .split('\0')
+        .filter_map(|line| {
+            if line.is_empty() {
+                return None;
+            }
+
+            let (key, value) = line.split_once('=')?;
+            if key.is_empty() {
+                return None;
+            }
+
+            Some((key.to_string(), value.to_string()))
+        })
+        .collect()
+}
+
+fn command_output_with_timeout(
+    mut cmd: std::process::Command,
+    timeout: Duration,
+) -> std::io::Result<Option<std::process::Output>> {
+    let mut child = cmd.spawn()?;
+    let start = Instant::now();
+
+    loop {
+        if child.try_wait()?.is_some() {
+            return child.wait_with_output().map(Some);
+        }
+
+        if start.elapsed() >= timeout {
+            let _ = child.kill();
+            let _ = child.wait();
+            return Ok(None);
+        }
+
+        std::thread::sleep(Duration::from_millis(25));
+    }
+}
+
+enum ShellEnvProbe {
+    Loaded(HashMap<String, String>),
+    Timeout,
+    Unavailable,
+}
+
+fn probe_shell_env(shell: &str, mode: &str) -> ShellEnvProbe {
+    let mut cmd = std::process::Command::new(shell);
+    cmd.args([mode, "-c", "env -0"]);
+    cmd.stdin(Stdio::null());
+    cmd.stdout(Stdio::piped());
+    cmd.stderr(Stdio::null());
+    let output = match command_output_with_timeout(cmd, SHELL_ENV_TIMEOUT) {
+        Ok(Some(output)) => output,
+        Ok(None) => return ShellEnvProbe::Timeout,
+        Err(error) => {
+            tracing::debug!(shell, mode, ?error, "Shell env probe failed");
+            return ShellEnvProbe::Unavailable;
+        }
+    };
+    if !output.status.success() {
+        tracing::debug!(shell, mode, "Shell env probe exited with non-zero status");
+        return ShellEnvProbe::Unavailable;
+    }
+    let env = parse_shell_env(&output.stdout);
+    if env.is_empty() {
+        tracing::debug!(shell, mode, "Shell env probe returned empty env");
+        return ShellEnvProbe::Unavailable;
+    }
+
+    ShellEnvProbe::Loaded(env)
+}
+
+fn is_nushell(shell: &str) -> bool {
+    let shell_name = Path::new(shell)
+        .file_name()
+        .and_then(|name| name.to_str())
+        .unwrap_or(shell)
+        .to_ascii_lowercase();
+    shell_name == "nu" || shell_name == "nu.exe" || shell.to_ascii_lowercase().ends_with("\\nu.exe")
+}
+fn load_shell_env(shell: &str) -> Option<HashMap<String, String>> {
+    if is_nushell(shell) {
+        tracing::debug!(shell, "Skipping shell env probe for nushell");
+        return None;
+    }
+
+    match probe_shell_env(shell, "-il") {
+        ShellEnvProbe::Loaded(env) => {
+            tracing::info!(
+                shell,
+                env_count = env.len(),
+                "Loaded shell environment with -il"
+            );
+            return Some(env);
+        }
+        ShellEnvProbe::Timeout => {
+            tracing::warn!(shell, "Interactive shell env probe timed out");
+            return None;
+        }
+        ShellEnvProbe::Unavailable => {}
+    }
+
+    if let ShellEnvProbe::Loaded(env) = probe_shell_env(shell, "-l") {
+        tracing::info!(
+            shell,
+            env_count = env.len(),
+            "Loaded shell environment with -l"
+        );
+        return Some(env);
+    }
+    tracing::warn!(shell, "Falling back to app environment");
+    None
+}
+
+fn merge_shell_env(
+    shell_env: Option<HashMap<String, String>>,
+    envs: Vec<(String, String)>,
+) -> Vec<(String, String)> {
+    let mut merged = shell_env.unwrap_or_default();
+    for (key, value) in envs {
+        merged.insert(key, value);
+    }
+
+    merged.into_iter().collect()
+}
+
 pub fn spawn_command(
     app: &tauri::AppHandle,
     args: &str,
@@ -312,6 +443,7 @@ pub fn spawn_command(
     } else {
         let sidecar = get_sidecar_path(app);
         let shell = get_user_shell();
+        let envs = merge_shell_env(load_shell_env(&shell), envs);
 
         let line = if shell.ends_with("/nu") {
             format!("^\"{}\" {}", sidecar.display(), args)
@@ -556,3 +688,54 @@ async fn read_line<F: Fn(String) -> CommandEvent + Send + Copy + 'static>(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashMap;
+
+    #[test]
+    fn parse_shell_env_supports_null_delimited_pairs() {
+        let env = parse_shell_env(b"PATH=/usr/bin:/bin\0FOO=bar=baz\0\0");
+
+        assert_eq!(env.get("PATH"), Some(&"/usr/bin:/bin".to_string()));
+        assert_eq!(env.get("FOO"), Some(&"bar=baz".to_string()));
+    }
+
+    #[test]
+    fn parse_shell_env_ignores_invalid_entries() {
+        let env = parse_shell_env(b"INVALID\0=empty\0OK=1\0");
+
+        assert_eq!(env.len(), 1);
+        assert_eq!(env.get("OK"), Some(&"1".to_string()));
+    }
+
+    #[test]
+    fn merge_shell_env_keeps_explicit_overrides() {
+        let mut shell_env = HashMap::new();
+        shell_env.insert("PATH".to_string(), "/shell/path".to_string());
+        shell_env.insert("HOME".to_string(), "/tmp/home".to_string());
+
+        let merged = merge_shell_env(
+            Some(shell_env),
+            vec![
+                ("PATH".to_string(), "/desktop/path".to_string()),
+                ("OPENCODE_CLIENT".to_string(), "desktop".to_string()),
+            ],
+        )
+        .into_iter()
+        .collect::<HashMap<_, _>>();
+
+        assert_eq!(merged.get("PATH"), Some(&"/desktop/path".to_string()));
+        assert_eq!(merged.get("HOME"), Some(&"/tmp/home".to_string()));
+        assert_eq!(merged.get("OPENCODE_CLIENT"), Some(&"desktop".to_string()));
+    }
+
+    #[test]
+    fn is_nushell_handles_path_and_binary_name() {
+        assert!(is_nushell("nu"));
+        assert!(is_nushell("/opt/homebrew/bin/nu"));
+        assert!(is_nushell("C:\\Program Files\\nu.exe"));
+        assert!(!is_nushell("/bin/zsh"));
+    }
+}