rumdl_lib/code_block_tools/
executor.rs1use super::config::ToolDefinition;
7use std::collections::HashMap;
8use std::io::{Read, Write};
9use std::process::{Command, Stdio};
10use std::sync::{Arc, Mutex};
11use std::thread;
12use std::time::{Duration, Instant};
13
14#[derive(Debug, Clone)]
16pub struct ToolOutput {
17 pub stdout: String,
19 pub stderr: String,
21 pub exit_code: i32,
23 pub success: bool,
25}
26
27#[derive(Debug, Clone)]
29pub enum ExecutorError {
30 ToolNotFound { tool: String },
32 ExecutionFailed { tool: String, message: String },
34 Timeout { tool: String, timeout_ms: u64 },
36 IoError { message: String },
38}
39
40impl std::fmt::Display for ExecutorError {
41 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42 match self {
43 Self::ToolNotFound { tool } => {
44 write!(f, "Tool '{tool}' not found in PATH")
45 }
46 Self::ExecutionFailed { tool, message } => {
47 write!(f, "Tool '{tool}' failed: {message}")
48 }
49 Self::Timeout { tool, timeout_ms } => {
50 write!(f, "Tool '{tool}' timed out after {timeout_ms}ms")
51 }
52 Self::IoError { message } => {
53 write!(f, "I/O error: {message}")
54 }
55 }
56 }
57}
58
59impl std::error::Error for ExecutorError {}
60
61pub struct ToolExecutor {
65 tool_cache: Arc<Mutex<HashMap<String, bool>>>,
67 default_timeout_ms: u64,
69}
70
71impl ToolExecutor {
72 pub fn new(default_timeout_ms: u64) -> Self {
74 Self {
75 tool_cache: Arc::new(Mutex::new(HashMap::new())),
76 default_timeout_ms,
77 }
78 }
79
80 pub fn is_tool_available(&self, tool_name: &str) -> bool {
82 {
84 let cache = self.tool_cache.lock().unwrap();
85 if let Some(&available) = cache.get(tool_name) {
86 return available;
87 }
88 }
89
90 let available = self.check_tool_exists(tool_name);
92
93 {
95 let mut cache = self.tool_cache.lock().unwrap();
96 cache.insert(tool_name.to_string(), available);
97 }
98
99 available
100 }
101
102 fn check_tool_exists(&self, tool_name: &str) -> bool {
104 #[cfg(unix)]
105 {
106 Command::new("which")
107 .arg(tool_name)
108 .stdout(Stdio::null())
109 .stderr(Stdio::null())
110 .status()
111 .is_ok_and(|s| s.success())
112 }
113
114 #[cfg(windows)]
115 {
116 Command::new("where")
117 .arg(tool_name)
118 .stdout(Stdio::null())
119 .stderr(Stdio::null())
120 .status()
121 .is_ok_and(|s| s.success())
122 }
123
124 #[cfg(not(any(unix, windows)))]
125 {
126 let _ = tool_name;
128 false
129 }
130 }
131
132 pub fn execute(
143 &self,
144 tool_def: &ToolDefinition,
145 input: &str,
146 is_format_mode: bool,
147 timeout_ms: Option<u64>,
148 ) -> Result<ToolOutput, ExecutorError> {
149 if tool_def.command.is_empty() {
150 return Err(ExecutorError::ExecutionFailed {
151 tool: "unknown".to_string(),
152 message: "Empty command".to_string(),
153 });
154 }
155
156 let tool_name = &tool_def.command[0];
157
158 if !self.is_tool_available(tool_name) {
160 return Err(ExecutorError::ToolNotFound {
161 tool: tool_name.clone(),
162 });
163 }
164
165 let mut cmd = Command::new(tool_name);
167
168 if tool_def.command.len() > 1 {
170 cmd.args(&tool_def.command[1..]);
171 }
172
173 let extra_args = if is_format_mode {
175 &tool_def.format_args
176 } else {
177 &tool_def.lint_args
178 };
179 if !extra_args.is_empty() {
180 cmd.args(extra_args);
181 }
182
183 if tool_def.stdin {
185 cmd.stdin(Stdio::piped());
186 }
187 cmd.stdout(Stdio::piped());
188 cmd.stderr(Stdio::piped());
189
190 let mut child = cmd.spawn().map_err(|e| ExecutorError::IoError {
192 message: format!("Failed to spawn '{tool_name}': {e}"),
193 })?;
194
195 let mut stdout_handle = child
196 .stdout
197 .take()
198 .map(|stdout| thread::spawn(move || read_pipe_to_string(stdout)));
199 let mut stderr_handle = child
200 .stderr
201 .take()
202 .map(|stderr| thread::spawn(move || read_pipe_to_string(stderr)));
203
204 if tool_def.stdin
208 && let Some(mut stdin) = child.stdin.take()
209 && let Err(e) = stdin.write_all(input.as_bytes())
210 && e.kind() != std::io::ErrorKind::BrokenPipe
211 {
212 return Err(ExecutorError::IoError {
213 message: format!("Failed to write to stdin: {e}"),
214 });
215 }
216
217 let timeout = Duration::from_millis(timeout_ms.unwrap_or(self.default_timeout_ms));
219 let status = if timeout.is_zero() {
220 child.wait().map_err(|e| ExecutorError::IoError {
221 message: format!("Failed to wait for '{tool_name}': {e}"),
222 })?
223 } else {
224 let start = Instant::now();
225 loop {
226 if let Some(status) = child.try_wait().map_err(|e| ExecutorError::IoError {
227 message: format!("Failed to poll '{tool_name}': {e}"),
228 })? {
229 break status;
230 }
231 if start.elapsed() >= timeout {
232 let _ = child.kill();
233 let _ = child.wait();
234 let _ = join_reader(stdout_handle.take());
235 let _ = join_reader(stderr_handle.take());
236 return Err(ExecutorError::Timeout {
237 tool: tool_name.clone(),
238 timeout_ms: timeout.as_millis() as u64,
239 });
240 }
241 thread::sleep(Duration::from_millis(10));
242 }
243 };
244
245 let stdout = join_reader(stdout_handle.take()).map_err(|e| ExecutorError::IoError { message: e })?;
246 let stderr = join_reader(stderr_handle.take()).map_err(|e| ExecutorError::IoError { message: e })?;
247 let exit_code = status.code().unwrap_or(-1);
248
249 Ok(ToolOutput {
250 stdout,
251 stderr,
252 exit_code,
253 success: status.success(),
254 })
255 }
256
257 pub fn format(
259 &self,
260 tool_def: &ToolDefinition,
261 input: &str,
262 timeout_ms: Option<u64>,
263 ) -> Result<String, ExecutorError> {
264 let output = self.execute(tool_def, input, true, timeout_ms)?;
265
266 if output.success && tool_def.stdout {
267 Ok(output.stdout)
268 } else if !output.success {
269 let exit_code = output.exit_code;
270 let stderr = &output.stderr;
271 Err(ExecutorError::ExecutionFailed {
272 tool: tool_def.command.first().cloned().unwrap_or_default(),
273 message: format!("Exit code {exit_code}: {stderr}"),
274 })
275 } else {
276 Err(ExecutorError::ExecutionFailed {
278 tool: tool_def.command.first().cloned().unwrap_or_default(),
279 message: "Formatter doesn't output to stdout".to_string(),
280 })
281 }
282 }
283
284 pub fn lint(
286 &self,
287 tool_def: &ToolDefinition,
288 input: &str,
289 timeout_ms: Option<u64>,
290 ) -> Result<ToolOutput, ExecutorError> {
291 self.execute(tool_def, input, false, timeout_ms)
292 }
293}
294
295fn read_pipe_to_string<R: Read>(mut pipe: R) -> std::io::Result<String> {
296 let mut buf = Vec::new();
297 pipe.read_to_end(&mut buf)?;
298 Ok(String::from_utf8_lossy(&buf).to_string())
299}
300
301fn join_reader(handle: Option<thread::JoinHandle<std::io::Result<String>>>) -> Result<String, String> {
302 match handle {
303 Some(handle) => match handle.join() {
304 Ok(res) => res.map_err(|e| format!("Failed to read output: {e}")),
305 Err(_) => Err("Output reader thread panicked".to_string()),
306 },
307 None => Ok(String::new()),
308 }
309}
310
311impl Default for ToolExecutor {
312 fn default() -> Self {
313 Self::new(30_000) }
315}
316
317#[cfg(test)]
318mod tests {
319 use super::*;
320
321 #[test]
322 fn test_executor_creation() {
323 let executor = ToolExecutor::new(10_000);
324 assert_eq!(executor.default_timeout_ms, 10_000);
326 }
327
328 #[test]
329 fn test_tool_not_found() {
330 let executor = ToolExecutor::default();
331 let tool_def = ToolDefinition {
332 command: vec!["nonexistent-tool-xyz123".to_string()],
333 stdin: true,
334 stdout: true,
335 lint_args: vec![],
336 format_args: vec![],
337 };
338
339 let result = executor.execute(&tool_def, "test", false, None);
340 assert!(matches!(result, Err(ExecutorError::ToolNotFound { .. })));
341 }
342
343 #[test]
344 fn test_empty_command() {
345 let executor = ToolExecutor::default();
346 let tool_def = ToolDefinition {
347 command: vec![],
348 stdin: true,
349 stdout: true,
350 lint_args: vec![],
351 format_args: vec![],
352 };
353
354 let result = executor.execute(&tool_def, "test", false, None);
355 assert!(matches!(result, Err(ExecutorError::ExecutionFailed { .. })));
356 }
357
358 #[test]
362 #[ignore = "requires 'cat' to be available"]
363 fn test_execute_cat() {
364 let executor = ToolExecutor::default();
365 let tool_def = ToolDefinition {
366 command: vec!["cat".to_string()],
367 stdin: true,
368 stdout: true,
369 lint_args: vec![],
370 format_args: vec![],
371 };
372
373 let result = executor.execute(&tool_def, "hello world", false, None);
374 let output = result.expect("cat should succeed");
375 assert!(output.success);
376 assert_eq!(output.stdout.trim(), "hello world");
377 }
378
379 #[test]
380 #[cfg(unix)]
381 #[ignore = "requires 'sleep' to be available"]
382 fn test_timeout() {
383 let executor = ToolExecutor::new(5);
384 let tool_def = ToolDefinition {
385 command: vec!["sleep".to_string(), "1".to_string()],
386 stdin: false,
387 stdout: true,
388 lint_args: vec![],
389 format_args: vec![],
390 };
391
392 let result = executor.execute(&tool_def, "", false, Some(5));
393 assert!(matches!(result, Err(ExecutorError::Timeout { .. })));
394 }
395}