diff --git a/crates/atuin-desktop-runtime/src/blocks/mod.rs b/crates/atuin-desktop-runtime/src/blocks/mod.rs index 150f4910..ade08a51 100644 --- a/crates/atuin-desktop-runtime/src/blocks/mod.rs +++ b/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; diff --git a/crates/atuin-desktop-runtime/src/blocks/script.rs b/crates/atuin-desktop-runtime/src/blocks/script.rs index c43b10fb..feb2ea35 100644 --- a/crates/atuin-desktop-runtime/src/blocks/script.rs +++ b/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() }; diff --git a/crates/atuin-desktop-runtime/src/blocks/terminal.rs b/crates/atuin-desktop-runtime/src/blocks/terminal.rs index 52f589d4..346479e7 100644 --- a/crates/atuin-desktop-runtime/src/blocks/terminal.rs +++ b/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 { diff --git a/crates/atuin-desktop-runtime/src/ssh/mod.rs b/crates/atuin-desktop-runtime/src/ssh/mod.rs index cafac459..bf608d47 100644 --- a/crates/atuin-desktop-runtime/src/ssh/mod.rs +++ b/crates/atuin-desktop-runtime/src/ssh/mod.rs @@ -22,6 +22,7 @@ mod pool; mod session; +mod ssh_env; mod ssh_pool; #[cfg(test)] @@ -29,4 +30,5 @@ 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}; diff --git a/crates/atuin-desktop-runtime/src/ssh/ssh_env.rs b/crates/atuin-desktop-runtime/src/ssh/ssh_env.rs new file mode 100644 index 00000000..acd1b571 --- /dev/null +++ b/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 { + 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); + } +}