1use 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
22const AUTO_BACKGROUND_TIMEOUT: u64 = 45;
24
25#[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
37pub 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#[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#[derive(Debug, Clone, Default, Serialize, Deserialize)]
115pub struct ExecToolArgs {
116 #[serde(default)]
117 pub action: String,
118 pub command: Option<Value>,
120 pub cwd: Option<String>,
122 pub workdir: Option<String>,
124 pub env: Option<HashMap<String, String>>,
126 pub timeout: Option<u64>,
128 pub shell: Option<String>,
130 pub proc_id: Option<String>,
132 pub timeout_ms: Option<u64>,
134 pub signal: Option<String>,
136 pub tail: Option<usize>,
138 pub filter: Option<String>,
140}
141
142pub 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 if let Ok(shell) = std::env::var("HANZO_MCP_FORCE_SHELL") {
159 return shell;
160 }
161
162 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 let cmd_str = match command {
202 Value::String(s) => s,
203 Value::Array(arr) => {
204 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 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 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 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 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; let default_timeout_ms = 600_000u64; 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 !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 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 if let Some(ref filter_id) = args.proc_id {
351 if id != *filter_id {
352 continue;
353 }
354 }
355
356 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 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, };
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 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#[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}