feat: only set ATUIN_OUTPUT_VARS over ssh if used

This feature requires several extra calls to setup the file + download
it before closing the terminal. This adds some latency

If it's not being used, don't pay the latency cost.

Author: ellie

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

@@ -676,21 +676,29 @@ impl Script {
             }
         };
 
-        // Create remote temp file for variable output
-        let remote_temp_path = match ssh_pool
-            .create_temp_file(&hostname, username.as_deref(), "atuin-desktop-vars")
-            .await
-        {
-            Ok(path) => path,
-            Err(e) => {
-                let error_msg = format!("Failed to create remote temp file: {}", e);
-                let _ = context.block_failed(error_msg.clone()).await;
-                return (Err(error_msg.into()), Vec::new(), None);
+        let uses_output_vars = code.contains("ATUIN_OUTPUT_VARS");
+
+        let remote_temp_path: Option<String> = if uses_output_vars {
+            match ssh_pool
+                .create_temp_file(&hostname, username.as_deref(), "atuin-desktop-vars")
+                .await
+            {
+                Ok(path) => Some(path),
+                Err(e) => {
+                    let error_msg = format!("Failed to create remote temp file: {}", e);
+                    let _ = context.block_failed(error_msg.clone()).await;
+                    return (Err(error_msg.into()), Vec::new(), None);
+                }
             }
+        } else {
+            None
         };
 
-        // Prepend environment variable export to the code
-        let code_with_vars = format!("export ATUIN_OUTPUT_VARS='{}'\n{}", remote_temp_path, code);
+        let code_to_run = if let Some(ref path) = remote_temp_path {
+            format!("export ATUIN_OUTPUT_VARS='{}'\n{}", path, code)
+        } else {
+            code.to_string()
+        };
 
         let channel_id = self.id.to_string();
         let (output_sender, mut output_receiver) = mpsc::channel::<SessionOutputLine>(100);
@@ -699,27 +707,26 @@ impl Script {
         let captured_output = Arc::new(RwLock::new(Vec::new()));
         let captured_output_clone = captured_output.clone();
 
-        // Take the cancellation receiver once at the start
         let mut cancel_rx = match cancellation_token.take_receiver() {
             Some(rx) => rx,
             None => {
                 let error_msg = "Cancellation receiver already taken";
                 let _ = context.block_failed(error_msg.to_string()).await;
-                // Cleanup remote temp file
-                let _ = ssh_pool
-                    .delete_file(&hostname, username.as_deref(), &remote_temp_path)
-                    .await;
+                if let Some(ref path) = remote_temp_path {
+                    let _ = ssh_pool
+                        .delete_file(&hostname, username.as_deref(), path)
+                        .await;
+                }
                 return (Err(error_msg.into()), Vec::new(), None);
             }
         };
 
-        // Execute SSH command with cancellation support
         let exec_result = tokio::select! {
             result = ssh_pool.exec(
                 &hostname,
                 username.as_deref(),
                 &self.interpreter,
-                &code_with_vars,
+                &code_to_run,
                 &channel_id,
                 output_sender,
                 result_tx,
@@ -730,19 +737,19 @@ impl Script {
                 tracing::trace!("Sending cancel to SSH execution for channel {channel_id}");
                 let _ = ssh_pool.exec_cancel(&channel_id).await;
                 let _ = context.block_cancelled().await;
-                // Cleanup remote temp file
-                let _ = ssh_pool.delete_file(&hostname, username.as_deref(), &remote_temp_path).await;
+                if let Some(ref path) = remote_temp_path {
+                    let _ = ssh_pool.delete_file(&hostname, username.as_deref(), path).await;
+                }
                 return (Err("SSH script execution cancelled before start".into()), Vec::new(), None);
             }
         };
 
         if let Err(e) = exec_result {
             let error_msg = format!("Failed to start SSH execution: {}", e);
             let _ = context.block_failed(error_msg.to_string()).await;
-            // Cleanup remote temp file
-            let _ = ssh_pool
-                .delete_file(&hostname, username.as_deref(), &remote_temp_path)
-                .await;
+            if let Some(ref path) = remote_temp_path {
+                let _ = ssh_pool.delete_file(&hostname, username.as_deref(), path).await;
+            }
             return (Err(error_msg.into()), Vec::new(), None);
         }
         let context_clone = context.clone();
@@ -787,8 +794,9 @@ impl Script {
                 let captured = captured_output.read().await.clone();
 
                 let _ = context.block_cancelled().await;
-                // Cleanup remote temp file
-                let _ = ssh_pool.delete_file(&hostname, username.as_deref(), &remote_temp_path).await;
+                if let Some(ref path) = remote_temp_path {
+                    let _ = ssh_pool.delete_file(&hostname, username.as_deref(), path).await;
+                }
                 return (Err("SSH script execution cancelled".into()), captured, None);
             }
             _ = result_rx => {
@@ -799,25 +807,26 @@ impl Script {
         let _ = output_task.await;
         let captured = captured_output.read().await.clone();
 
-        // Read variables from remote temp file
-        let vars = match ssh_pool
-            .read_file(&hostname, username.as_deref(), &remote_temp_path)
-            .await
-        {
-            Ok(contents) => {
-                // Parse the file contents using fs_var::parse_vars
-                Some(fs_var::parse_vars(&contents))
-            }
-            Err(e) => {
-                tracing::warn!("Failed to read remote temp file for variables: {}", e);
-                None
+        let vars = if let Some(ref path) = remote_temp_path {
+            match ssh_pool
+                .read_file(&hostname, username.as_deref(), path)
+                .await
+            {
+                Ok(contents) => Some(fs_var::parse_vars(&contents)),
+                Err(e) => {
+                    tracing::warn!("Failed to read remote temp file for variables: {}", e);
+                    None
+                }
             }
+        } else {
+            None
         };
 
-        // Cleanup remote temp file
-        let _ = ssh_pool
-            .delete_file(&hostname, username.as_deref(), &remote_temp_path)
-            .await;
+        if let Some(ref path) = remote_temp_path {
+            let _ = ssh_pool
+                .delete_file(&hostname, username.as_deref(), path)
+                .await;
+        }
 
         (Ok(exit_code), captured, vars)
     }

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

@@ -214,57 +214,65 @@ impl Terminal {
             .take_receiver()
             .ok_or("Cancellation receiver already taken")?;
 
-        // Accumulator for terminal output (shared between reader tasks)
         let output_accumulator: Arc<RwLock<Vec<u8>>> = Arc::new(RwLock::new(Vec::new()));
 
-        // Setup fs_var for local terminal or remote path for SSH
+        let templated_code = if !self.code.is_empty() {
+            context.context_resolver.resolve_template(&self.code)?
+        } else {
+            String::new()
+        };
+
+        let uses_output_vars = templated_code.contains("ATUIN_OUTPUT_VARS");
+
+        let ssh_host = context.context_resolver.ssh_host().cloned();
         let fs_var_handle: Option<crate::context::fs_var::FsVarHandle>;
-        let remote_var_path: Option<(String, Option<String>, String)>; // (hostname, username, path)
+        let remote_var_path: Option<String>;
 
-        if let Some(ssh_host) = context.context_resolver.ssh_host() {
+        if let Some(ref host) = ssh_host {
             fs_var_handle = None;
 
-            // For SSH, create remote temp file
-            let (username, hostname) = Self::parse_ssh_host(ssh_host);
-            let ssh_pool = context
-                .ssh_pool
-                .clone()
-                .ok_or("SSH pool not available in execution context")?;
+            if uses_output_vars {
+                let (username, hostname) = Self::parse_ssh_host(host);
+                let ssh_pool = context
+                    .ssh_pool
+                    .clone()
+                    .ok_or("SSH pool not available in execution context")?;
 
-            let remote_path = ssh_pool
-                .create_temp_file(&hostname, username.as_deref(), "atuin-desktop-vars")
-                .await
-                .map_err(|e| format!("Failed to create remote temp file: {}", e))?;
+                let remote_path = ssh_pool
+                    .create_temp_file(&hostname, username.as_deref(), "atuin-desktop-vars")
+                    .await
+                    .map_err(|e| format!("Failed to create remote temp file: {}", e))?;
 
-            remote_var_path = Some((hostname, username, remote_path));
+                remote_var_path = Some(remote_path);
+            } else {
+                remote_var_path = None;
+            }
         } else {
-            fs_var_handle =
-                Some(crate::context::fs_var::setup().map_err(|e| {
-                    format!("Failed to setup temp file for output variables: {}", e)
-                })?);
+            if uses_output_vars {
+                fs_var_handle =
+                    Some(crate::context::fs_var::setup().map_err(|e| {
+                        format!("Failed to setup temp file for output variables: {}", e)
+                    })?);
+            } else {
+                fs_var_handle = None;
+            }
             remote_var_path = None;
         }
 
-        // Open PTY based on context (local or SSH)
         let cancellation_token_clone = cancellation_token.clone();
-        let pty: Box<dyn PtyLike + Send> = if let Some((
-            ref hostname,
-            ref username,
-            ref remote_path,
-        )) = remote_var_path
-        {
-            // Get SSH pool from context
+        let pty: Box<dyn PtyLike + Send> = if let Some(ref host) = ssh_host {
+            let (username, hostname) = Self::parse_ssh_host(host);
             let ssh_pool = context
                 .ssh_pool
                 .clone()
                 .ok_or("SSH pool not available in execution context")?;
 
-            // Create SSH PTY with cancellation support
             let (output_sender, mut output_receiver) = tokio::sync::mpsc::channel(100);
             let hostname_clone = hostname.clone();
             let username_clone = username.clone();
             let pty_id_str = self.id.to_string();
             let ssh_pool_clone = ssh_pool.clone();
+            let remote_path_clone = remote_var_path.clone();
 
             let initial_cols = self.cols;
             let initial_rows = self.rows;
@@ -282,8 +290,9 @@ impl Terminal {
                 _ = &mut cancel_rx => {
                     let _ = ssh_pool_clone.close_pty(&pty_id_str).await;
                     let _ = context.block_cancelled().await;
-                    // Cleanup remote temp file if it was created
-                    let _ = ssh_pool_clone.delete_file(&hostname_clone, username_clone.as_deref(), remote_path).await;
+                    if let Some(ref path) = remote_path_clone {
+                        let _ = ssh_pool_clone.delete_file(&hostname_clone, username_clone.as_deref(), path).await;
+                    }
                     return Err("SSH PTY connection cancelled".into());
                 }
             };
@@ -431,25 +440,21 @@ impl Terminal {
             .emit_gc_event(GCEvent::PtyOpened(metadata.clone()))
             .await;
 
-        // For SSH terminals, export ATUIN_OUTPUT_VARS first
-        if let Some((_, _, ref remote_path)) = remote_var_path {
+        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 {
                 tracing::warn!("Failed to write export command to SSH PTY: {}", e);
             }
         }
 
-        // Write the command to the PTY after started event
-        if !self.code.is_empty() {
-            let command = context.context_resolver.resolve_template(&self.code)?;
-            let command = if command.ends_with('\n') {
-                command
+        if !templated_code.is_empty() {
+            let command = if templated_code.ends_with('\n') {
+                templated_code.clone()
             } else {
-                format!("{}\n", command)
+                format!("{}\n", templated_code)
             };
 
             if let Err(e) = pty_store.write_pty(self.id, command.into()).await {
-                // Send error event if command writing fails
                 let _ = context
                     .block_failed(format!("Failed to write command to PTY: {}", e))
                     .await;
@@ -490,12 +495,12 @@ impl Terminal {
                     tracing::warn!("Failed to read terminal output variables: {}", e);
                 }
             }
-        } else if let Some((hostname, username, remote_path)) = remote_var_path {
-            // SSH terminal
+        } else if let (Some(ref host), Some(ref remote_path)) = (&ssh_host, &remote_var_path) {
+            let (username, hostname) = Self::parse_ssh_host(host);
             let ssh_pool = context.ssh_pool.clone().unwrap();
 
             match ssh_pool
-                .read_file(&hostname, username.as_deref(), &remote_path)
+                .read_file(&hostname, username.as_deref(), remote_path)
                 .await
             {
                 Ok(contents) => {
@@ -520,9 +525,8 @@ impl Terminal {
                 }
             }
 
-            // Cleanup remote temp file
             let _ = ssh_pool
-                .delete_file(&hostname, username.as_deref(), &remote_path)
+                .delete_file(&hostname, username.as_deref(), remote_path)
                 .await;
         }