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