Skip to main content

hematite/tools/
shell.rs

1use serde_json::Value;
2use std::path::Path;
3use std::time::Duration;
4use tokio::io::AsyncBufReadExt;
5use tokio::sync::mpsc;
6
7const DEFAULT_TIMEOUT_SECS: u64 = 60;
8const MAX_OUTPUT_BYTES: usize = 65_536; // 64 KB cap (higher for professional mode)
9
10/// Execute a shell command and return its stdout/stderr combined as a String.
11///
12/// Unified Shell Adapter:
13/// - Windows: Tries `pwsh` first, then `powershell.exe`, then `cmd /C`.
14/// - Unix: Tries `sh -c`.
15pub async fn execute(args: &Value) -> Result<String, String> {
16    let mut command = args
17        .get("command")
18        .and_then(|v| v.as_str())
19        .ok_or_else(|| "Missing required argument: 'command'".to_string())?
20        .to_string();
21
22    // Expand @path/to/file into the absolute workspace path before execution.
23    if command.contains('@') {
24        let root = crate::tools::file_ops::workspace_root();
25        let root_str = root.to_string_lossy().to_string().replace("\\", "/");
26        command = command.replace('@', &format!("{}/", root_str.trim_end_matches('/')));
27    }
28
29    let timeout_ms = args
30        .get("timeout_ms")
31        .and_then(|v| v.as_u64())
32        .or_else(|| {
33            args.get("timeout_secs")
34                .and_then(|v| v.as_u64())
35                .map(|s| s * 1000)
36        })
37        .unwrap_or(DEFAULT_TIMEOUT_SECS * 1000);
38
39    let run_in_background = args
40        .get("run_in_background")
41        .and_then(|v| v.as_bool())
42        .unwrap_or(false);
43
44    let cwd =
45        std::env::current_dir().map_err(|e| format!("Failed to get working directory: {e}"))?;
46
47    execute_command_in_dir(&command, &cwd, timeout_ms, run_in_background).await
48}
49
50/// Like `execute`, but streams each stdout/stderr line to the TUI as a
51/// `ShellLine` event while the command runs, so the operator sees live
52/// progress instead of a blank screen until completion.
53///
54/// Falls back to plain `execute` for background tasks.
55pub async fn execute_streaming(
56    args: &Value,
57    tx: mpsc::Sender<crate::agent::inference::InferenceEvent>,
58) -> Result<String, String> {
59    // Background tasks don't benefit from streaming — delegate to execute().
60    if args
61        .get("run_in_background")
62        .and_then(|v| v.as_bool())
63        .unwrap_or(false)
64    {
65        return execute(args).await;
66    }
67
68    let mut command = args
69        .get("command")
70        .and_then(|v| v.as_str())
71        .ok_or_else(|| "Missing required argument: 'command'".to_string())?
72        .to_string();
73
74    if command.contains('@') {
75        let root = crate::tools::file_ops::workspace_root();
76        let root_str = root.to_string_lossy().replace("\\", "/").to_string();
77        command = command.replace('@', &format!("{}/", root_str.trim_end_matches('/')));
78    }
79
80    let timeout_ms = args
81        .get("timeout_ms")
82        .and_then(|v| v.as_u64())
83        .or_else(|| {
84            args.get("timeout_secs")
85                .and_then(|v| v.as_u64())
86                .map(|s| s * 1000)
87        })
88        .unwrap_or(DEFAULT_TIMEOUT_SECS * 1000);
89
90    crate::tools::guard::bash_is_safe(&command)?;
91
92    let cwd =
93        std::env::current_dir().map_err(|e| format!("Failed to get working directory: {e}"))?;
94
95    let mut tokio_cmd = build_command(&command).await;
96    tokio_cmd
97        .current_dir(&cwd)
98        .stdout(std::process::Stdio::piped())
99        .stderr(std::process::Stdio::piped());
100
101    let sandbox_root = cwd.join(".hematite").join("sandbox");
102    let _ = std::fs::create_dir_all(&sandbox_root);
103    tokio_cmd.env("HOME", &sandbox_root);
104    tokio_cmd.env("TMPDIR", &sandbox_root);
105
106    let mut child = tokio_cmd
107        .spawn()
108        .map_err(|e| format!("Failed to spawn process: {e}"))?;
109
110    let stdout = child.stdout.take().expect("stdout was piped");
111    let stderr = child.stderr.take().expect("stderr was piped");
112
113    let mut stdout_lines = tokio::io::BufReader::new(stdout).lines();
114    let mut stderr_lines = tokio::io::BufReader::new(stderr).lines();
115
116    let mut out_buf = String::new();
117    let mut err_buf = String::new();
118    let mut stdout_done = false;
119    let mut stderr_done = false;
120
121    let deadline = tokio::time::Instant::now() + Duration::from_millis(timeout_ms);
122
123    loop {
124        if stdout_done && stderr_done {
125            break;
126        }
127        tokio::select! {
128            _ = tokio::time::sleep_until(deadline) => {
129                let _ = child.kill().await;
130                return Err(format!("Command timed out after {} ms: {}", timeout_ms, command));
131            }
132            line = stdout_lines.next_line(), if !stdout_done => {
133                match line {
134                    Ok(Some(l)) => {
135                        let clean = l.trim_end_matches('\r').to_string();
136                        let _ = tx
137                            .send(crate::agent::inference::InferenceEvent::ShellLine(clean.clone()))
138                            .await;
139                        out_buf.push_str(&clean);
140                        out_buf.push('\n');
141                    }
142                    _ => stdout_done = true,
143                }
144            }
145            line = stderr_lines.next_line(), if !stderr_done => {
146                match line {
147                    Ok(Some(l)) => {
148                        let clean = l.trim_end_matches('\r').to_string();
149                        let _ = tx
150                            .send(crate::agent::inference::InferenceEvent::ShellLine(
151                                format!("[err] {}", clean),
152                            ))
153                            .await;
154                        err_buf.push_str(&clean);
155                        err_buf.push('\n');
156                    }
157                    _ => stderr_done = true,
158                }
159            }
160        }
161    }
162
163    // Wait for exit, with a short grace period for cleanup.
164    let status = tokio::time::timeout(Duration::from_millis(5_000), child.wait())
165        .await
166        .map_err(|_| "Process cleanup timed out".to_string())?
167        .map_err(|e| format!("Failed to wait for process: {e}"))?;
168
169    let stdout_capped = cap_bytes(out_buf.as_bytes(), MAX_OUTPUT_BYTES / 2);
170    let stderr_capped = cap_bytes(err_buf.as_bytes(), MAX_OUTPUT_BYTES / 2);
171
172    let exit_info = match status.code() {
173        Some(0) => String::new(),
174        Some(code) => format!("\n[exit code: {code}]"),
175        None => "\n[process terminated by signal]".to_string(),
176    };
177
178    let mut result = String::new();
179    if !stdout_capped.is_empty() {
180        result.push_str(&stdout_capped);
181    }
182    if !stderr_capped.is_empty() {
183        if !result.is_empty() {
184            result.push('\n');
185        }
186        result.push_str("[stderr]\n");
187        result.push_str(&stderr_capped);
188    }
189    if result.is_empty() {
190        result.push_str("(no output)");
191    }
192    result.push_str(&exit_info);
193
194    Ok(crate::agent::utils::strip_ansi(&result))
195}
196
197pub async fn execute_command_in_dir(
198    command: &str,
199    cwd: &Path,
200    timeout_ms: u64,
201    run_in_background: bool,
202) -> Result<String, String> {
203    crate::tools::guard::bash_is_safe(command)?;
204
205    let mut tokio_cmd = build_command(command).await;
206    tokio_cmd
207        .current_dir(cwd)
208        .stdout(std::process::Stdio::piped())
209        .stderr(std::process::Stdio::piped());
210
211    let sandbox_root = cwd.join(".hematite").join("sandbox");
212    let _ = std::fs::create_dir_all(&sandbox_root);
213    tokio_cmd.env("HOME", &sandbox_root);
214    tokio_cmd.env("TMPDIR", &sandbox_root);
215
216    if run_in_background {
217        let _child = tokio_cmd
218            .spawn()
219            .map_err(|e| format!("Failed to spawn background process: {e}"))?;
220        return Ok(
221            "[background_task_id: spawned]\nCommand started in background. Use `ps` or `jobs` to monitor if available."
222                .into(),
223        );
224    }
225
226    let child_future = tokio_cmd.output();
227
228    let output = match tokio::time::timeout(Duration::from_millis(timeout_ms), child_future).await {
229        Ok(Ok(output)) => output,
230        Ok(Err(e)) => return Err(format!("Failed to execution process: {e}")),
231        Err(_) => {
232            return Err(format!(
233                "Command timed out after {} ms: {}",
234                timeout_ms, command
235            ))
236        }
237    };
238
239    let stdout = cap_bytes(&output.stdout, MAX_OUTPUT_BYTES / 2);
240    let stderr = cap_bytes(&output.stderr, MAX_OUTPUT_BYTES / 2);
241
242    let exit_info = match output.status.code() {
243        Some(0) => String::new(),
244        Some(code) => format!("\n[exit code: {code}]"),
245        None => "\n[process terminated by signal]".to_string(),
246    };
247
248    let mut result = String::new();
249    if !stdout.is_empty() {
250        result.push_str(&stdout);
251    }
252    if !stderr.is_empty() {
253        if !result.is_empty() {
254            result.push('\n');
255        }
256        result.push_str("[stderr]\n");
257        result.push_str(&stderr);
258    }
259    if result.is_empty() {
260        result.push_str("(no output)");
261    }
262    result.push_str(&exit_info);
263
264    Ok(crate::agent::utils::strip_ansi(&result))
265}
266
267/// Build the platform-appropriate shell invocation.
268async fn build_command(command: &str) -> tokio::process::Command {
269    #[cfg(target_os = "windows")]
270    {
271        let normalized = command
272            .replace("/dev/null", "$null")
273            .replace("1>/dev/null", "2>$null")
274            .replace("2>/dev/null", "2>$null");
275
276        if which("pwsh").await {
277            let mut cmd = tokio::process::Command::new("pwsh");
278            cmd.args(["-NoProfile", "-NonInteractive", "-Command", &normalized]);
279            cmd
280        } else {
281            let mut cmd = tokio::process::Command::new("powershell");
282            cmd.args(["-NoProfile", "-NonInteractive", "-Command", &normalized]);
283            cmd
284        }
285    }
286    #[cfg(not(target_os = "windows"))]
287    {
288        let mut cmd = tokio::process::Command::new("sh");
289        cmd.args(["-c", command]);
290        cmd
291    }
292}
293
294#[allow(dead_code)]
295async fn which(name: &str) -> bool {
296    #[cfg(target_os = "windows")]
297    let check = format!("{}.exe", name);
298    #[cfg(not(target_os = "windows"))]
299    let check = name;
300
301    tokio::process::Command::new("where")
302        .arg(check)
303        .stdout(std::process::Stdio::null())
304        .stderr(std::process::Stdio::null())
305        .status()
306        .await
307        .map(|s| s.success())
308        .unwrap_or(false)
309}
310
311fn cap_bytes(bytes: &[u8], max: usize) -> String {
312    if bytes.len() <= max {
313        String::from_utf8_lossy(bytes).into_owned()
314    } else {
315        let mut s = String::from_utf8_lossy(&bytes[..max]).into_owned();
316        s.push_str(&format!("\n... [truncated - {} bytes total]", bytes.len()));
317        s
318    }
319}