Skip to main content

hanzo_mcp/tools/
exec_tool.rs

1/// Unified process execution tool (HIP-0300)
2///
3/// Handles all process operations:
4/// - exec: Execute commands (the ONE execution primitive)
5/// - wait: Wait for background process
6/// - ps: List processes
7/// - kill: Kill process
8/// - logs: Get process logs
9
10use anyhow::{anyhow, Result};
11use serde::{Deserialize, Serialize};
12use serde_json::{json, Value};
13use std::collections::HashMap;
14use std::path::PathBuf;
15use std::process::Stdio;
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18use tokio::io::{AsyncBufReadExt, BufReader};
19use tokio::process::Command;
20use tokio::sync::RwLock;
21
22/// Auto-background timeout in seconds
23const AUTO_BACKGROUND_TIMEOUT: u64 = 45;
24
25/// Process info tracked by the manager
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct ProcessInfo {
28    pub proc_id: String,
29    pub pid: Option<u32>,
30    pub command: String,
31    pub running: bool,
32    pub exit_code: Option<i32>,
33    pub started: String,
34    pub log_file: Option<PathBuf>,
35}
36
37/// Process manager singleton
38pub struct ProcessManager {
39    processes: Arc<RwLock<HashMap<String, ProcessInfo>>>,
40    counter: Arc<RwLock<u64>>,
41}
42
43impl ProcessManager {
44    pub fn new() -> Self {
45        Self {
46            processes: Arc::new(RwLock::new(HashMap::new())),
47            counter: Arc::new(RwLock::new(0)),
48        }
49    }
50
51    async fn next_id(&self) -> String {
52        let mut counter = self.counter.write().await;
53        *counter += 1;
54        format!("proc_{}", *counter)
55    }
56
57    async fn register(&self, info: ProcessInfo) {
58        let mut procs = self.processes.write().await;
59        procs.insert(info.proc_id.clone(), info);
60    }
61
62    async fn update(&self, proc_id: &str, exit_code: i32) {
63        let mut procs = self.processes.write().await;
64        if let Some(info) = procs.get_mut(proc_id) {
65            info.running = false;
66            info.exit_code = Some(exit_code);
67        }
68    }
69
70    pub async fn list(&self) -> HashMap<String, ProcessInfo> {
71        self.processes.read().await.clone()
72    }
73
74    pub async fn get(&self, proc_id: &str) -> Option<ProcessInfo> {
75        self.processes.read().await.get(proc_id).cloned()
76    }
77}
78
79/// Actions for the proc tool
80#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
81#[serde(rename_all = "snake_case")]
82pub enum ProcAction {
83    Exec,
84    Wait,
85    Ps,
86    Kill,
87    Logs,
88    Help,
89}
90
91impl Default for ProcAction {
92    fn default() -> Self {
93        Self::Help
94    }
95}
96
97impl std::str::FromStr for ProcAction {
98    type Err = anyhow::Error;
99
100    fn from_str(s: &str) -> Result<Self> {
101        match s.to_lowercase().as_str() {
102            "exec" => Ok(Self::Exec),
103            "wait" => Ok(Self::Wait),
104            "ps" | "list" => Ok(Self::Ps),
105            "kill" => Ok(Self::Kill),
106            "logs" | "log" => Ok(Self::Logs),
107            "help" | "" => Ok(Self::Help),
108            _ => Err(anyhow!("Unknown action: {}", s)),
109        }
110    }
111}
112
113/// Arguments for proc tool
114#[derive(Debug, Clone, Default, Serialize, Deserialize)]
115pub struct ExecToolArgs {
116    #[serde(default)]
117    pub action: String,
118    /// Command to execute (string or array for Rust parity)
119    pub command: Option<Value>,
120    /// Working directory
121    pub cwd: Option<String>,
122    /// Alias for cwd (Rust parity)
123    pub workdir: Option<String>,
124    /// Environment variables
125    pub env: Option<HashMap<String, String>>,
126    /// Timeout in seconds
127    pub timeout: Option<u64>,
128    /// Shell to use
129    pub shell: Option<String>,
130    /// Process ID for wait/kill/logs
131    pub proc_id: Option<String>,
132    /// Timeout in milliseconds for wait
133    pub timeout_ms: Option<u64>,
134    /// Signal for kill
135    pub signal: Option<String>,
136    /// Number of lines for logs
137    pub tail: Option<usize>,
138    /// Filter for ps
139    pub filter: Option<String>,
140}
141
142/// Shell execution tool
143pub struct ExecTool {
144    manager: Arc<ProcessManager>,
145    shell: String,
146}
147
148impl ExecTool {
149    pub fn new() -> Self {
150        Self {
151            manager: Arc::new(ProcessManager::new()),
152            shell: Self::resolve_shell(),
153        }
154    }
155
156    fn resolve_shell() -> String {
157        // Check environment override
158        if let Ok(shell) = std::env::var("HANZO_MCP_FORCE_SHELL") {
159            return shell;
160        }
161
162        // Prefer zsh, fallback to others
163        for shell in ["zsh", "bash", "fish", "dash", "sh"] {
164            for prefix in ["/opt/homebrew/bin", "/usr/local/bin", "/bin", "/usr/bin"] {
165                let path = format!("{}/{}", prefix, shell);
166                if std::path::Path::new(&path).exists() {
167                    return path;
168                }
169            }
170            if let Ok(found) = which::which(shell) {
171                return found.to_string_lossy().to_string();
172            }
173        }
174
175        "sh".to_string()
176    }
177
178    pub async fn execute(&self, args: ExecToolArgs) -> Result<String> {
179        let action: ProcAction = if args.action.is_empty() {
180            ProcAction::Help
181        } else {
182            args.action.parse()?
183        };
184
185        let result = match action {
186            ProcAction::Exec => self.exec(args).await?,
187            ProcAction::Wait => self.wait(args).await?,
188            ProcAction::Ps => self.ps(args).await?,
189            ProcAction::Kill => self.kill(args).await?,
190            ProcAction::Logs => self.logs(args).await?,
191            ProcAction::Help => self.help()?,
192        };
193
194        Ok(serde_json::to_string(&result)?)
195    }
196
197    async fn exec(&self, args: ExecToolArgs) -> Result<Value> {
198        let command = args.command.ok_or_else(|| anyhow!("command required"))?;
199
200        // Support both string and array format
201        let cmd_str = match command {
202            Value::String(s) => s,
203            Value::Array(arr) => {
204                // Join array into shell command
205                arr.iter()
206                    .filter_map(|v| v.as_str())
207                    .map(|s| shell_escape::escape(s.into()).to_string())
208                    .collect::<Vec<_>>()
209                    .join(" ")
210            }
211            _ => return Err(anyhow!("command must be string or array")),
212        };
213
214        let cwd = args.workdir.or(args.cwd);
215        let timeout = args.timeout.unwrap_or(AUTO_BACKGROUND_TIMEOUT);
216        let shell = args.shell.unwrap_or_else(|| self.shell.clone());
217
218        let proc_id = self.manager.next_id().await;
219        let started = chrono::Utc::now().to_rfc3339();
220
221        // Build command
222        let mut cmd = Command::new(&shell);
223        cmd.arg("-c").arg(&cmd_str);
224        cmd.stdout(Stdio::piped());
225        cmd.stderr(Stdio::piped());
226
227        if let Some(ref dir) = cwd {
228            cmd.current_dir(dir);
229        }
230
231        if let Some(ref env_vars) = args.env {
232            for (k, v) in env_vars {
233                cmd.env(k, v);
234            }
235        }
236
237        let start = Instant::now();
238        let mut child = cmd.spawn()?;
239        let pid = child.id();
240
241        // Register process
242        self.manager.register(ProcessInfo {
243            proc_id: proc_id.clone(),
244            pid,
245            command: cmd_str.clone(),
246            running: true,
247            exit_code: None,
248            started: started.clone(),
249            log_file: None,
250        }).await;
251
252        // Wait with timeout
253        let timeout_duration = Duration::from_secs(timeout);
254        let result = tokio::time::timeout(timeout_duration, child.wait_with_output()).await;
255
256        match result {
257            Ok(Ok(output)) => {
258                let exit_code = output.status.code().unwrap_or(-1);
259                let duration_ms = start.elapsed().as_millis() as u64;
260
261                self.manager.update(&proc_id, exit_code).await;
262
263                let stdout = String::from_utf8_lossy(&output.stdout).to_string();
264                let stderr = String::from_utf8_lossy(&output.stderr).to_string();
265
266                Ok(json!({
267                    "proc_id": proc_id,
268                    "exit_code": exit_code,
269                    "stdout": stdout,
270                    "stderr": stderr,
271                    "duration_ms": duration_ms,
272                    "status": if exit_code == 0 { "success" } else { "failed" }
273                }))
274            }
275            Ok(Err(e)) => Err(anyhow!("Process failed: {}", e)),
276            Err(_) => {
277                // Timeout - process is backgrounded
278                Ok(json!({
279                    "proc_id": proc_id,
280                    "exit_code": null,
281                    "stdout_ref": format!("proc:{}:stdout", proc_id),
282                    "stderr_ref": format!("proc:{}:stderr", proc_id),
283                    "status": "running",
284                    "message": format!("Command backgrounded after {}s. Use proc(action='logs', proc_id='{}') to view output.", timeout, proc_id)
285                }))
286            }
287        }
288    }
289
290    async fn wait(&self, args: ExecToolArgs) -> Result<Value> {
291        let proc_id = args.proc_id.ok_or_else(|| anyhow!("proc_id required"))?;
292
293        let max_timeout_ms = 3_600_000u64; // 1 hour
294        let default_timeout_ms = 600_000u64; // 10 minutes
295        let timeout_ms = args.timeout_ms.unwrap_or(default_timeout_ms).min(max_timeout_ms);
296        let timeout_sec = timeout_ms as f64 / 1000.0;
297
298        let info = self.manager.get(&proc_id).await
299            .ok_or_else(|| anyhow!("Process not found: {}", proc_id))?;
300
301        // If already completed, return immediately
302        if !info.running {
303            return Ok(json!({
304                "proc_id": proc_id,
305                "exit_code": info.exit_code,
306                "output": "",
307                "status": "completed"
308            }));
309        }
310
311        // Poll until complete or timeout
312        let start = Instant::now();
313        let poll_interval = Duration::from_millis(500);
314
315        loop {
316            if start.elapsed().as_secs_f64() >= timeout_sec {
317                return Ok(json!({
318                    "proc_id": proc_id,
319                    "exit_code": null,
320                    "output": "",
321                    "status": "timeout",
322                    "message": format!("Timed out after {}ms", timeout_ms)
323                }));
324            }
325
326            if let Some(info) = self.manager.get(&proc_id).await {
327                if !info.running {
328                    return Ok(json!({
329                        "proc_id": proc_id,
330                        "exit_code": info.exit_code,
331                        "output": "",
332                        "status": "completed",
333                        "duration_ms": start.elapsed().as_millis() as u64
334                    }));
335                }
336            } else {
337                return Err(anyhow!("Process disappeared: {}", proc_id));
338            }
339
340            tokio::time::sleep(poll_interval).await;
341        }
342    }
343
344    async fn ps(&self, args: ExecToolArgs) -> Result<Value> {
345        let processes = self.manager.list().await;
346        let mut results = Vec::new();
347
348        for (id, info) in processes {
349            // Filter by proc_id
350            if let Some(ref filter_id) = args.proc_id {
351                if id != *filter_id {
352                    continue;
353                }
354            }
355
356            // Filter by command pattern
357            if let Some(ref filter) = args.filter {
358                if !info.command.to_lowercase().contains(&filter.to_lowercase()) {
359                    continue;
360                }
361            }
362
363            results.push(json!({
364                "proc_id": info.proc_id,
365                "pid": info.pid,
366                "command": info.command,
367                "running": info.running,
368                "exit_code": info.exit_code,
369                "started": info.started
370            }));
371        }
372
373        Ok(json!({
374            "processes": results,
375            "total": results.len()
376        }))
377    }
378
379    async fn kill(&self, args: ExecToolArgs) -> Result<Value> {
380        let proc_id = args.proc_id.ok_or_else(|| anyhow!("proc_id required"))?;
381
382        let info = self.manager.get(&proc_id).await
383            .ok_or_else(|| anyhow!("Process not found: {}", proc_id))?;
384
385        let pid = info.pid.ok_or_else(|| anyhow!("Process has no PID"))?;
386
387        // Resolve signal
388        let sig = match args.signal.as_deref() {
389            Some("KILL") | Some("9") => 9,
390            Some("INT") | Some("2") => 2,
391            Some("HUP") | Some("1") => 1,
392            Some("QUIT") | Some("3") => 3,
393            _ => 15, // TERM
394        };
395
396        #[cfg(unix)]
397        {
398            use nix::sys::signal::{kill, Signal};
399            use nix::unistd::Pid;
400
401            let signal = Signal::try_from(sig).unwrap_or(Signal::SIGTERM);
402            match kill(Pid::from_raw(pid as i32), signal) {
403                Ok(_) => Ok(json!({
404                    "proc_id": proc_id,
405                    "pid": pid,
406                    "signal": sig,
407                    "killed": true
408                })),
409                Err(nix::errno::Errno::ESRCH) => Ok(json!({
410                    "proc_id": proc_id,
411                    "pid": pid,
412                    "signal": sig,
413                    "killed": false,
414                    "message": "Process already terminated"
415                })),
416                Err(e) => Err(anyhow!("Cannot kill process: {}", e)),
417            }
418        }
419
420        #[cfg(not(unix))]
421        {
422            Err(anyhow!("kill not supported on this platform"))
423        }
424    }
425
426    async fn logs(&self, args: ExecToolArgs) -> Result<Value> {
427        let proc_id = args.proc_id.ok_or_else(|| anyhow!("proc_id required"))?;
428
429        let info = self.manager.get(&proc_id).await
430            .ok_or_else(|| anyhow!("Process not found: {}", proc_id))?;
431
432        // If log file exists, read it
433        if let Some(ref log_file) = info.log_file {
434            if log_file.exists() {
435                let content = tokio::fs::read_to_string(log_file).await?;
436                let lines: Vec<&str> = content.lines().collect();
437                let total_lines = lines.len();
438                let tail = args.tail.unwrap_or(100);
439                let output = if total_lines > tail {
440                    lines[total_lines - tail..].join("\n")
441                } else {
442                    content
443                };
444
445                return Ok(json!({
446                    "proc_id": proc_id,
447                    "output": output,
448                    "running": info.running,
449                    "exit_code": info.exit_code,
450                    "total_lines": total_lines
451                }));
452            }
453        }
454
455        Ok(json!({
456            "proc_id": proc_id,
457            "stdout": "",
458            "stderr": "",
459            "message": "No log file available"
460        }))
461    }
462
463    fn help(&self) -> Result<Value> {
464        let shell_name = std::path::Path::new(&self.shell)
465            .file_name()
466            .and_then(|n| n.to_str())
467            .unwrap_or("sh");
468
469        Ok(json!({
470            "name": "exec",
471            "version": "0.12.0",
472            "description": format!("Unified process execution tool (HIP-0300). Shell: {}", shell_name),
473            "actions": {
474                "exec": "Execute command (the ONE execution primitive)",
475                "wait": "Wait for background process to complete",
476                "ps": "List processes",
477                "kill": "Kill process",
478                "logs": "Get process logs"
479            },
480            "returns": "proc_id, exit_code, stdout, stderr",
481            "auto_background": format!("{}s", AUTO_BACKGROUND_TIMEOUT)
482        }))
483    }
484}
485
486/// MCP Tool Definition
487#[derive(Debug, Serialize, Deserialize)]
488pub struct ExecToolDefinition {
489    pub name: String,
490    pub description: String,
491    pub input_schema: Value,
492}
493
494impl ExecToolDefinition {
495    pub fn new() -> Self {
496        Self {
497            name: "exec".to_string(),
498            description: format!(
499                r#"Unified process execution tool (HIP-0300).
500
501Actions:
502- exec: Execute command (the ONE primitive)
503- wait: Wait for background process
504- ps: List processes
505- kill: Kill process
506- logs: Get process logs
507
508Returns: {{proc_id, exit_code, stdout, stderr}}
509Auto-backgrounds commands after {}s."#,
510                AUTO_BACKGROUND_TIMEOUT
511            ),
512            input_schema: json!({
513                "type": "object",
514                "properties": {
515                    "action": {
516                        "type": "string",
517                        "enum": ["exec", "wait", "ps", "kill", "logs", "help"],
518                        "default": "help",
519                        "description": "Action to perform"
520                    },
521                    "command": {
522                        "oneOf": [
523                            {"type": "string"},
524                            {"type": "array", "items": {"type": "string"}}
525                        ],
526                        "description": "Command to execute (string or array)"
527                    },
528                    "cwd": {"type": "string", "description": "Working directory"},
529                    "workdir": {"type": "string", "description": "Alias for cwd (Rust parity)"},
530                    "env": {
531                        "type": "object",
532                        "additionalProperties": {"type": "string"},
533                        "description": "Environment variables"
534                    },
535                    "timeout": {"type": "integer", "description": "Timeout in seconds"},
536                    "shell": {"type": "string", "description": "Shell to use"},
537                    "proc_id": {"type": "string", "description": "Process ID"},
538                    "timeout_ms": {"type": "integer", "description": "Wait timeout in milliseconds"},
539                    "signal": {"type": "string", "description": "Kill signal"},
540                    "tail": {"type": "integer", "description": "Number of log lines"},
541                    "filter": {"type": "string", "description": "Filter for ps"}
542                }
543            }),
544        }
545    }
546}
547
548#[cfg(test)]
549mod tests {
550    use super::*;
551
552    #[tokio::test]
553    async fn test_exec_simple() {
554        let tool = ExecTool::new();
555        let args = ExecToolArgs {
556            action: "exec".to_string(),
557            command: Some(Value::String("echo hello".to_string())),
558            ..Default::default()
559        };
560
561        let result = tool.execute(args).await;
562        assert!(result.is_ok());
563        let output = result.unwrap();
564        assert!(output.contains("hello"));
565    }
566
567    #[tokio::test]
568    async fn test_exec_array_command() {
569        let tool = ExecTool::new();
570        let args = ExecToolArgs {
571            action: "exec".to_string(),
572            command: Some(Value::Array(vec![
573                Value::String("echo".to_string()),
574                Value::String("hello world".to_string()),
575            ])),
576            ..Default::default()
577        };
578
579        let result = tool.execute(args).await;
580        assert!(result.is_ok());
581    }
582
583    #[tokio::test]
584    async fn test_ps() {
585        let tool = ExecTool::new();
586        let args = ExecToolArgs {
587            action: "ps".to_string(),
588            ..Default::default()
589        };
590
591        let result = tool.execute(args).await;
592        assert!(result.is_ok());
593        let output = result.unwrap();
594        assert!(output.contains("processes"));
595    }
596
597    #[tokio::test]
598    async fn test_help() {
599        let tool = ExecTool::new();
600        let args = ExecToolArgs {
601            action: "help".to_string(),
602            ..Default::default()
603        };
604
605        let result = tool.execute(args).await;
606        assert!(result.is_ok());
607        let output = result.unwrap();
608        assert!(output.contains("exec"));
609        assert!(output.contains("exec"));
610    }
611}