fix: env blocks not setting variable for SSH script and terminal blocks (#377)

Author: BinaryMuse

crates/atuin-desktop-runtime/src/blocks/mod.rs

@@ -27,6 +27,7 @@ pub(crate) mod script;
 pub(crate) mod sql_block;
 pub(crate) mod sqlite;
 pub(crate) mod ssh_connect;
+
 pub(crate) mod sub_runbook;
 pub(crate) mod terminal;
 pub(crate) mod var;

crates/atuin-desktop-runtime/src/blocks/script.rs

@@ -17,7 +17,7 @@ use crate::events::GCEvent;
 use crate::execution::{
     CancellationToken, ExecutionContext, ExecutionHandle, ExecutionStatus, StreamingBlockOutput,
 };
-use crate::ssh::{OutputLine as SessionOutputLine, SshWarning};
+use crate::ssh::{build_env_exports, OutputLine as SessionOutputLine, SshWarning};
 
 use super::FromDocument;
 
@@ -697,8 +697,15 @@ impl Script {
             None
         };
 
+        let env_exports = build_env_exports(context.context_resolver.env_vars());
+
         let code_to_run = if let Some(ref path) = remote_temp_path {
-            format!("export ATUIN_OUTPUT_VARS='{}'\n{}", path, code)
+            format!(
+                "{}export ATUIN_OUTPUT_VARS='{}'\n{}",
+                env_exports, path, code
+            )
+        } else if !env_exports.is_empty() {
+            format!("{}{}", env_exports, code)
         } else {
             code.to_string()
         };

crates/atuin-desktop-runtime/src/blocks/terminal.rs

@@ -14,7 +14,7 @@ use crate::execution::{
     CancellationToken, ExecutionContext, ExecutionHandle, ExecutionStatus, StreamingBlockOutput,
 };
 use crate::pty::{Pty, PtyLike};
-use crate::ssh::{SshPty, SshWarning};
+use crate::ssh::{build_env_exports, SshPty, SshWarning};
 
 /// Output structure for Terminal blocks that implements BlockExecutionOutput
 /// for template access to terminal output.
@@ -491,6 +491,15 @@ impl Terminal {
             .emit_gc_event(GCEvent::PtyOpened(metadata.clone()))
             .await;
 
+        if ssh_host.is_some() {
+            let env_exports = build_env_exports(context.context_resolver.env_vars());
+            if !env_exports.is_empty() {
+                if let Err(e) = pty_store.write_pty(self.id, env_exports.into()).await {
+                    tracing::warn!("Failed to write env exports to SSH PTY: {}", e);
+                }
+            }
+        }
+
         if let Some(ref remote_path) = remote_var_path {
             let export_cmd = format!("export ATUIN_OUTPUT_VARS='{}'\n", remote_path);
             if let Err(e) = pty_store.write_pty(self.id, export_cmd.into()).await {

crates/atuin-desktop-runtime/src/ssh/mod.rs

@@ -22,11 +22,13 @@
 
 mod pool;
 mod session;
+mod ssh_env;
 mod ssh_pool;
 
 #[cfg(test)]
 mod integration_tests;
 
 pub use pool::Pool;
 pub use session::{Authentication, CommandResult, OutputLine, Session, SshConfig, SshWarning};
+pub use ssh_env::build_env_exports;
 pub use ssh_pool::{SshPoolHandle, SshPty};

crates/atuin-desktop-runtime/src/ssh/ssh_env.rs

@@ -0,0 +1,102 @@
+use std::collections::HashMap;
+
+/// Escape a value for safe inclusion in a POSIX single-quoted string.
+///
+/// Replaces each `'` with `'\''` (end quote, escaped literal quote, start new quote)
+/// and wraps the result in single quotes. POSIX single-quoted strings treat all
+/// characters as literal, so this handles `$`, backticks, newlines, etc.
+fn shell_escape_value(value: &str) -> String {
+    let mut escaped = String::with_capacity(value.len() + 2);
+    escaped.push('\'');
+    for ch in value.chars() {
+        if ch == '\'' {
+            escaped.push_str("'\\''");
+        } else {
+            escaped.push(ch);
+        }
+    }
+    escaped.push('\'');
+    escaped
+}
+
+/// Build `export K='V'\n` lines for all env vars.
+///
+/// Returns an empty string if the map is empty.
+pub fn build_env_exports(env_vars: &HashMap<String, String>) -> String {
+    if env_vars.is_empty() {
+        return String::new();
+    }
+
+    let mut out = String::new();
+    for (key, value) in env_vars {
+        out.push_str("export ");
+        out.push_str(key);
+        out.push('=');
+        out.push_str(&shell_escape_value(value));
+        out.push('\n');
+    }
+    out
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_shell_escape_simple_value() {
+        assert_eq!(shell_escape_value("hello"), "'hello'");
+    }
+
+    #[test]
+    fn test_shell_escape_empty_value() {
+        assert_eq!(shell_escape_value(""), "''");
+    }
+
+    #[test]
+    fn test_shell_escape_value_with_single_quotes() {
+        assert_eq!(shell_escape_value("it's"), "'it'\\''s'");
+    }
+
+    #[test]
+    fn test_shell_escape_value_with_special_chars() {
+        assert_eq!(shell_escape_value("$HOME"), "'$HOME'");
+        assert_eq!(shell_escape_value("`cmd`"), "'`cmd`'");
+        assert_eq!(shell_escape_value("a\nb"), "'a\nb'");
+        assert_eq!(shell_escape_value("a b"), "'a b'");
+        assert_eq!(shell_escape_value("a\"b"), "'a\"b'");
+    }
+
+    #[test]
+    fn test_build_env_exports_empty_map() {
+        let map = HashMap::new();
+        assert_eq!(build_env_exports(&map), "");
+    }
+
+    #[test]
+    fn test_build_env_exports_single_var() {
+        let mut map = HashMap::new();
+        map.insert("FOO".to_string(), "bar".to_string());
+        assert_eq!(build_env_exports(&map), "export FOO='bar'\n");
+    }
+
+    #[test]
+    fn test_build_env_exports_special_value() {
+        let mut map = HashMap::new();
+        map.insert("VAR".to_string(), "it's $complex".to_string());
+        assert_eq!(build_env_exports(&map), "export VAR='it'\\''s $complex'\n");
+    }
+
+    #[test]
+    fn test_build_env_exports_multiple_vars() {
+        let mut map = HashMap::new();
+        map.insert("A".to_string(), "1".to_string());
+        map.insert("B".to_string(), "2".to_string());
+
+        let result = build_env_exports(&map);
+        // HashMap order is not guaranteed, so check both lines are present
+        assert!(result.contains("export A='1'\n"));
+        assert!(result.contains("export B='2'\n"));
+        // Should have exactly 2 lines
+        assert_eq!(result.lines().count(), 2);
+    }
+}