Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 54 additions & 43 deletions crates/atuin-desktop-runtime/src/blocks/script.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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,
Expand All @@ -730,19 +737,21 @@ 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();
Expand Down Expand Up @@ -787,8 +796,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 => {
Expand All @@ -799,25 +809,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)
}
Expand Down
87 changes: 45 additions & 42 deletions crates/atuin-desktop-runtime/src/blocks/terminal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -214,57 +214,64 @@ 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| {
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;
Expand All @@ -282,8 +289,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());
}
};
Expand Down Expand Up @@ -431,25 +439,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;
Expand Down Expand Up @@ -490,12 +494,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) => {
Expand All @@ -520,9 +524,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;
}

Expand Down
20 changes: 12 additions & 8 deletions src/components/runbooks/editor/blocks/Script/Script.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,8 @@ const ScriptBlock = ({
const blockContext = useBlockContext(script.id);
const sshParent = blockContext.sshHost;

const showSpinner = (blockExecution.isStarting || blockExecution.isStopping) && !!sshParent;

const onBlockOutput = useCallback(async (output: GenericBlockOutput<void>) => {
if (output.stdout) {
xtermRef.current?.write(output.stdout);
Expand Down Expand Up @@ -355,14 +357,16 @@ const ScriptBlock = ({
color="danger"
>
<div>
<PlayButton
eventName="runbooks.block.execute"
eventProps={{ type: "script" }}
onPlay={handlePlay}
onStop={blockExecution.cancel}
isRunning={blockExecution.isRunning}
cancellable={true}
/>
<PlayButton
eventName="runbooks.block.execute"
eventProps={{ type: "script" }}
onPlay={handlePlay}
onStop={blockExecution.cancel}
isRunning={blockExecution.isRunning}
cancellable={true}
isLoading={showSpinner}
disabled={showSpinner}
/>
</div>
</Tooltip>

Expand Down
3 changes: 2 additions & 1 deletion src/lib/blocks/terminal/component.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -328,7 +328,8 @@ export const RunBlock = ({

<div className="flex flex-row gap-2 flex-grow w-full" ref={elementRef}>
<PlayButton
isLoading={isLoading}
isLoading={isLoading || ((execution.isStarting || execution.isStopping) && !!sshParent)}
disabled={isLoading || ((execution.isStarting || execution.isStopping) && !!sshParent)}
isRunning={execution.isRunning}
cancellable={true}
onPlay={handlePlay}
Expand Down
Loading
Loading