terminal: Do not auto close shell terminals if they error out (#38182)

Lukas Wirth created

Closes https://github.com/zed-industries/zed/issues/38134

This also reduces an annoying level of shell nesting

Release Notes:

- N/A

Change summary

crates/remote/src/remote_client.rs |  26 +--
crates/remote/src/transport/ssh.rs | 212 +++++++++++++++++++++++--------
2 files changed, 165 insertions(+), 73 deletions(-)

Detailed changes

crates/remote/src/remote_client.rs 🔗

@@ -769,13 +769,11 @@ impl RemoteClient {
     }
 
     pub fn shell(&self) -> Option<String> {
-        Some(self.state.as_ref()?.remote_connection()?.shell())
+        Some(self.remote_connection()?.shell())
     }
 
     pub fn shares_network_interface(&self) -> bool {
-        self.state
-            .as_ref()
-            .and_then(|state| state.remote_connection())
+        self.remote_connection()
             .map_or(false, |connection| connection.shares_network_interface())
     }
 
@@ -787,12 +785,8 @@ impl RemoteClient {
         working_dir: Option<String>,
         port_forward: Option<(u16, String, u16)>,
     ) -> Result<CommandTemplate> {
-        let Some(connection) = self
-            .state
-            .as_ref()
-            .and_then(|state| state.remote_connection())
-        else {
-            return Err(anyhow!("no connection"));
+        let Some(connection) = self.remote_connection() else {
+            return Err(anyhow!("no ssh connection"));
         };
         connection.build_command(program, args, env, working_dir, port_forward)
     }
@@ -803,11 +797,7 @@ impl RemoteClient {
         dest_path: RemotePathBuf,
         cx: &App,
     ) -> Task<Result<()>> {
-        let Some(connection) = self
-            .state
-            .as_ref()
-            .and_then(|state| state.remote_connection())
-        else {
+        let Some(connection) = self.remote_connection() else {
             return Task::ready(Err(anyhow!("no ssh connection")));
         };
         connection.upload_directory(src_path, dest_path, cx)
@@ -916,6 +906,12 @@ impl RemoteClient {
             .unwrap()
             .unwrap()
     }
+
+    fn remote_connection(&self) -> Option<Arc<dyn RemoteConnection>> {
+        self.state
+            .as_ref()
+            .and_then(|state| state.remote_connection())
+    }
 }
 
 enum ConnectionPoolEntry {

crates/remote/src/transport/ssh.rs 🔗

@@ -113,64 +113,24 @@ impl RemoteConnection for SshRemoteConnection {
         working_dir: Option<String>,
         port_forward: Option<(u16, String, u16)>,
     ) -> Result<CommandTemplate> {
-        use std::fmt::Write as _;
-
-        let mut script = String::new();
-        if let Some(working_dir) = working_dir {
-            let working_dir =
-                RemotePathBuf::new(working_dir.into(), self.ssh_path_style).to_string();
-
-            // shlex will wrap the command in single quotes (''), disabling ~ expansion,
-            // replace ith with something that works
-            const TILDE_PREFIX: &'static str = "~/";
-            let working_dir = if working_dir.starts_with(TILDE_PREFIX) {
-                let working_dir = working_dir.trim_start_matches("~").trim_start_matches("/");
-                format!("$HOME/{working_dir}")
-            } else {
-                working_dir
-            };
-            write!(&mut script, "cd \"{working_dir}\"; ",).unwrap();
-        } else {
-            write!(&mut script, "cd; ").unwrap();
-        };
-
-        for (k, v) in input_env.iter() {
-            if let Some((k, v)) = shlex::try_quote(k).ok().zip(shlex::try_quote(v).ok()) {
-                write!(&mut script, "{}={} ", k, v).unwrap();
-            }
-        }
-
-        let shell = &self.ssh_shell;
-
-        if let Some(input_program) = input_program {
-            let command = shlex::try_quote(&input_program)?;
-            script.push_str(&command);
-            for arg in input_args {
-                let arg = shlex::try_quote(&arg)?;
-                script.push_str(" ");
-                script.push_str(&arg);
-            }
-        } else {
-            write!(&mut script, "exec {shell} -l").unwrap();
-        };
-
-        let shell_invocation = format!("{shell} -c {}", shlex::try_quote(&script).unwrap());
-
-        let mut args = Vec::new();
-        args.extend(self.socket.ssh_args());
-
-        if let Some((local_port, host, remote_port)) = port_forward {
-            args.push("-L".into());
-            args.push(format!("{local_port}:{host}:{remote_port}"));
-        }
-
-        args.push("-t".into());
-        args.push(shell_invocation);
-        Ok(CommandTemplate {
-            program: "ssh".into(),
-            args,
-            env: self.socket.envs.clone(),
-        })
+        let Self {
+            ssh_path_style,
+            socket,
+            ssh_shell,
+            ..
+        } = self;
+        let env = socket.envs.clone();
+        build_command(
+            input_program,
+            input_args,
+            input_env,
+            working_dir,
+            port_forward,
+            env,
+            *ssh_path_style,
+            ssh_shell,
+            socket.ssh_args(),
+        )
     }
 
     fn upload_directory(
@@ -1037,3 +997,139 @@ impl SshConnectionOptions {
         }
     }
 }
+
+fn build_command(
+    input_program: Option<String>,
+    input_args: &[String],
+    input_env: &HashMap<String, String>,
+    working_dir: Option<String>,
+    port_forward: Option<(u16, String, u16)>,
+    ssh_env: HashMap<String, String>,
+    ssh_path_style: PathStyle,
+    ssh_shell: &str,
+    ssh_args: Vec<String>,
+) -> Result<CommandTemplate> {
+    use std::fmt::Write as _;
+
+    let mut exec = String::from("exec env -C ");
+    if let Some(working_dir) = working_dir {
+        let working_dir = RemotePathBuf::new(working_dir.into(), ssh_path_style).to_string();
+
+        // shlex will wrap the command in single quotes (''), disabling ~ expansion,
+        // replace with with something that works
+        const TILDE_PREFIX: &'static str = "~/";
+        if working_dir.starts_with(TILDE_PREFIX) {
+            let working_dir = working_dir.trim_start_matches("~").trim_start_matches("/");
+            write!(exec, "\"$HOME/{working_dir}\" ",).unwrap();
+        } else {
+            write!(exec, "\"{working_dir}\" ",).unwrap();
+        }
+    } else {
+        write!(exec, "\"$HOME\" ").unwrap();
+    };
+
+    for (k, v) in input_env.iter() {
+        if let Some((k, v)) = shlex::try_quote(k).ok().zip(shlex::try_quote(v).ok()) {
+            write!(exec, "{}={} ", k, v).unwrap();
+        }
+    }
+
+    write!(exec, "{ssh_shell} ").unwrap();
+    if let Some(input_program) = input_program {
+        let mut script = shlex::try_quote(&input_program)?.into_owned();
+        for arg in input_args {
+            let arg = shlex::try_quote(&arg)?;
+            script.push_str(" ");
+            script.push_str(&arg);
+        }
+        write!(exec, "-c {}", shlex::try_quote(&script).unwrap()).unwrap();
+    } else {
+        write!(exec, "-l").unwrap();
+    };
+
+    let mut args = Vec::new();
+    args.extend(ssh_args);
+
+    if let Some((local_port, host, remote_port)) = port_forward {
+        args.push("-L".into());
+        args.push(format!("{local_port}:{host}:{remote_port}"));
+    }
+
+    args.push("-t".into());
+    args.push(exec);
+    Ok(CommandTemplate {
+        program: "ssh".into(),
+        args,
+        env: ssh_env,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_build_command() -> Result<()> {
+        let mut input_env = HashMap::default();
+        input_env.insert("INPUT_VA".to_string(), "val".to_string());
+        let mut env = HashMap::default();
+        env.insert("SSH_VAR".to_string(), "ssh-val".to_string());
+
+        let command = build_command(
+            Some("remote_program".to_string()),
+            &["arg1".to_string(), "arg2".to_string()],
+            &input_env,
+            Some("~/work".to_string()),
+            None,
+            env.clone(),
+            PathStyle::Posix,
+            "/bin/fish",
+            vec!["-p".to_string(), "2222".to_string()],
+        )?;
+
+        assert_eq!(command.program, "ssh");
+        assert_eq!(
+            command.args.iter().map(String::as_str).collect::<Vec<_>>(),
+            [
+                "-p",
+                "2222",
+                "-t",
+                "exec env -C \"$HOME/work\" INPUT_VA=val /bin/fish -c 'remote_program arg1 arg2'"
+            ]
+        );
+        assert_eq!(command.env, env);
+
+        let mut input_env = HashMap::default();
+        input_env.insert("INPUT_VA".to_string(), "val".to_string());
+        let mut env = HashMap::default();
+        env.insert("SSH_VAR".to_string(), "ssh-val".to_string());
+
+        let command = build_command(
+            None,
+            &["arg1".to_string(), "arg2".to_string()],
+            &input_env,
+            None,
+            Some((1, "foo".to_owned(), 2)),
+            env.clone(),
+            PathStyle::Posix,
+            "/bin/fish",
+            vec!["-p".to_string(), "2222".to_string()],
+        )?;
+
+        assert_eq!(command.program, "ssh");
+        assert_eq!(
+            command.args.iter().map(String::as_str).collect::<Vec<_>>(),
+            [
+                "-p",
+                "2222",
+                "-L",
+                "1:foo:2",
+                "-t",
+                "exec env -C \"$HOME\" INPUT_VA=val /bin/fish -l"
+            ]
+        );
+        assert_eq!(command.env, env);
+
+        Ok(())
+    }
+}