use crate::apr_bin::apr_binary;
use crate::types::{ContentBlock, ToolCallResult};
use std::ffi::OsStr;
use std::io::{BufRead, BufReader, Read};
use std::process::{Command, Stdio};
use std::sync::mpsc::{Receiver, TryRecvError};
use std::time::{Duration, Instant};
pub const CANCEL_GRACE_MS: u64 = 30_000;
const POLL_INTERVAL: Duration = Duration::from_millis(10);
fn quote_arg(arg: &str) -> String {
let safe = !arg.is_empty()
&& arg
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b"@%+=:,./-_".contains(&b));
if safe {
arg.to_string()
} else {
format!("'{}'", arg.replace('\'', r"'\''"))
}
}
fn failure_result(cmd_display: &str, code: i32, stdout: &str, stderr: &str) -> ToolCallResult {
let summary = if stderr.trim().is_empty() {
stdout.to_string()
} else {
stderr.to_string()
};
let mut content = vec![ContentBlock::text(format!(
"`{cmd_display}` failed (exit {code}): {summary}"
))];
if !stderr.trim().is_empty() && !stdout.trim().is_empty() {
content.push(ContentBlock::text(stdout.to_string()));
}
ToolCallResult {
content,
is_error: Some(true),
}
}
#[must_use]
pub fn run_apr(args: &[&str]) -> ToolCallResult {
run_program(apr_binary(), args)
}
#[must_use]
pub fn run_program<P: AsRef<OsStr>>(program: P, args: &[&str]) -> ToolCallResult {
let program = program.as_ref();
let cmd_display = display_cmd(program, args);
let output = match Command::new(program).args(args).output() {
Ok(o) => o,
Err(e) => {
return ToolCallResult::error(format!("Failed to spawn `{cmd_display}`: {e}"));
}
};
let stdout = String::from_utf8_lossy(&output.stdout).into_owned();
let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
if output.status.success() {
if stdout.trim().is_empty() {
ToolCallResult::error(format!("`{cmd_display}` produced no output"))
} else {
ToolCallResult::success(stdout)
}
} else {
let code = output.status.code().unwrap_or(-1);
failure_result(&cmd_display, code, &stdout, &stderr)
}
}
fn display_cmd(program: &OsStr, args: &[&str]) -> String {
let mut out = quote_arg(&program.to_string_lossy());
for a in args {
out.push(' ');
out.push_str("e_arg(a));
}
out
}
#[must_use]
pub fn run_apr_cancellable(
args: &[&str],
cancel_rx: &Receiver<()>,
grace_ms: u64,
) -> ToolCallResult {
spawn_cancellable(apr_binary(), args, cancel_rx, grace_ms)
}
#[must_use]
pub fn spawn_cancellable<P: AsRef<OsStr>>(
program: P,
args: &[&str],
cancel_rx: &Receiver<()>,
grace_ms: u64,
) -> ToolCallResult {
let program = program.as_ref();
let cmd_display = display_cmd(program, args);
let mut child = match Command::new(program)
.args(args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(c) => c,
Err(e) => {
return ToolCallResult::error(format!("Failed to spawn `{cmd_display}`: {e}"));
}
};
let pid = child.id();
let wait_status = loop {
match child.try_wait() {
Ok(Some(status)) => break Ok(status),
Ok(None) => {}
Err(e) => {
return ToolCallResult::error(format!("Failed to poll `{cmd_display}`: {e}"));
}
}
match cancel_rx.try_recv() {
Ok(()) => break Err(CancelReason::Signalled),
Err(TryRecvError::Empty) => {}
Err(TryRecvError::Disconnected) => {
}
}
std::thread::sleep(POLL_INTERVAL);
};
match wait_status {
Ok(status) => {
let stdout = drain(&mut child.stdout.take());
let stderr = drain(&mut child.stderr.take());
if status.success() {
if stdout.trim().is_empty() {
ToolCallResult::error(format!("`{cmd_display}` produced no output"))
} else {
ToolCallResult::success(stdout)
}
} else {
let code = status.code().unwrap_or(-1);
failure_result(&cmd_display, code, &stdout, &stderr)
}
}
Err(CancelReason::Signalled) => {
send_sigterm(pid);
let deadline = Instant::now() + Duration::from_millis(grace_ms);
let mut escalated = false;
loop {
match child.try_wait() {
Ok(Some(_)) => break,
Ok(None) => {}
Err(_) => break,
}
if Instant::now() >= deadline {
if !escalated {
let _ = child.kill();
escalated = true;
} else {
break;
}
}
std::thread::sleep(POLL_INTERVAL);
}
let _ = child.wait();
let stdout = drain(&mut child.stdout.take());
let preview = truncate_for_preview(&stdout);
ToolCallResult::error(format!(
"Cancelled: `{cmd_display}` terminated by notifications/cancelled; partial stdout: {preview}"
))
}
}
}
enum CancelReason {
Signalled,
}
fn drain<R: Read>(reader: &mut Option<R>) -> String {
let mut buf = String::new();
if let Some(r) = reader.as_mut() {
let _ = r.read_to_string(&mut buf);
}
buf
}
fn truncate_for_preview(s: &str) -> String {
const MAX: usize = 512;
if s.len() <= MAX {
s.to_string()
} else {
let truncated: String = s.chars().take(MAX).collect();
format!("{truncated}… (truncated)")
}
}
#[cfg(unix)]
fn send_sigterm(pid: u32) {
use nix::sys::signal::{kill, Signal};
use nix::unistd::Pid;
#[allow(clippy::cast_possible_wrap)]
let raw = pid as i32;
let _ = kill(Pid::from_raw(raw), Signal::SIGTERM);
}
#[cfg(not(unix))]
fn send_sigterm(_pid: u32) {
}
#[must_use]
pub fn run_apr_streaming<F>(args: &[&str], on_line: F) -> ToolCallResult
where
F: FnMut(&str),
{
spawn_streaming(apr_binary(), args, on_line)
}
#[must_use]
pub fn spawn_streaming<P: AsRef<OsStr>, F>(
program: P,
args: &[&str],
mut on_line: F,
) -> ToolCallResult
where
F: FnMut(&str),
{
let program = program.as_ref();
let cmd_display = display_cmd(program, args);
let mut child = match Command::new(program)
.args(args)
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(c) => c,
Err(e) => {
return ToolCallResult::error(format!("Failed to spawn `{cmd_display}`: {e}"));
}
};
let stdout_pipe = match child.stdout.take() {
Some(p) => p,
None => {
let _ = child.wait();
return ToolCallResult::error(format!("Failed to capture stdout of `{cmd_display}`"));
}
};
let mut accumulated = String::new();
let reader = BufReader::new(stdout_pipe);
for line in reader.lines() {
match line {
Ok(text) => {
on_line(&text);
accumulated.push_str(&text);
accumulated.push('\n');
}
Err(e) => {
let _ = child.wait();
return ToolCallResult::error(format!(
"Failed to read stdout of `{cmd_display}`: {e}"
));
}
}
}
let status = match child.wait() {
Ok(s) => s,
Err(e) => {
return ToolCallResult::error(format!("Failed to reap `{cmd_display}`: {e}"));
}
};
let stderr = drain(&mut child.stderr.take());
if status.success() {
if accumulated.trim().is_empty() {
ToolCallResult::error(format!("`{cmd_display}` produced no output"))
} else {
ToolCallResult::success(accumulated)
}
} else {
let code = status.code().unwrap_or(-1);
let detail = if stderr.trim().is_empty() {
accumulated
} else {
stderr
};
ToolCallResult::error(format!("`{cmd_display}` failed (exit {code}): {detail}"))
}
}
#[cfg(test)]
#[allow(clippy::disallowed_methods)] mod tests {
use super::*;
use std::sync::mpsc;
use std::thread;
#[test]
fn failure_keeps_the_stdout_report_when_stderr_also_spoke() {
let report = r#"{"passed":false,"gates":[{"name":"ollama_parity","passed":false}]}"#;
let result = failure_result(
"apr qa m.gguf --json",
5,
report,
"error: Validation failed",
);
assert_eq!(result.is_error, Some(true));
let whole: String = result
.content
.iter()
.map(|b| b.text.as_str())
.collect::<Vec<_>>()
.join("\n");
assert!(
whole.contains("ollama_parity"),
"gate report must reach the client, got: {whole}"
);
assert!(
whole.contains("failed (exit 5)"),
"summary line must survive too, got: {whole}"
);
}
#[test]
fn cancellable_failure_carries_both_streams() {
let (_tx, rx) = mpsc::channel::<()>();
let result = spawn_cancellable(
"sh",
&[
"-c",
"printf '{\"gates\":\"REP\"}\\nORT\\n'; echo SUMMARY >&2; exit 5",
],
&rx,
CANCEL_GRACE_MS,
);
assert_eq!(result.is_error, Some(true));
assert!(
result.content[0].text.contains("SUMMARY"),
"stderr dropped: {}",
result.content[0].text
);
assert_eq!(
result.content.len(),
2,
"stdout report dropped, only got: {:?}",
result.content
);
assert!(
result.content[1].text.contains("{\"gates\":\"REP\"}\nORT"),
"stdout report mangled: {}",
result.content[1].text
);
}
#[test]
fn failure_with_empty_stderr_reports_stdout_once() {
let result = failure_result("apr qa m.gguf", 1, "only-stdout", " \n");
assert_eq!(result.content.len(), 1);
assert!(result.content[0].text.contains("only-stdout"));
}
#[test]
fn echoed_command_is_shell_quoted() {
let cmd = display_cmd(
OsStr::new("apr"),
&["run", "m.gguf", "--prompt", "What is 2+2?"],
);
assert_eq!(cmd, "apr run m.gguf --prompt 'What is 2+2?'");
}
#[test]
fn quoting_leaves_safe_args_alone_and_escapes_quotes() {
assert_eq!(quote_arg("--max-tokens"), "--max-tokens");
assert_eq!(
quote_arg("/home/noah/models/a.gguf"),
"/home/noah/models/a.gguf"
);
assert_eq!(quote_arg(""), "''");
assert_eq!(quote_arg("it's"), r"'it'\''s'");
}
#[test]
fn spawn_failure_maps_to_tool_error() {
let result = run_apr(&["this-subcommand-does-not-exist"]);
assert_eq!(result.is_error, Some(true));
}
#[test]
#[cfg(unix)]
fn falsify_2384_run_apr_executes_the_resolved_binary() {
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
let dir =
std::env::temp_dir().join(format!("aprender-mcp-2384-run-apr-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).expect("mkdir scratch");
let shim = dir.join("apr");
{
let mut f = std::fs::File::create(&shim).expect("create shim");
writeln!(f, "#!/bin/sh").expect("shebang");
writeln!(f, "if [ \"$1\" = \"validate\" ]; then").expect("if");
writeln!(f, " echo '{{\"marker\":\"APR-BIN-RESOLVED-SHIM\"}}'").expect("body");
writeln!(f, " exit 0").expect("ok");
writeln!(f, "fi").expect("fi");
writeln!(f, "exit 2").expect("unknown subcommand");
f.sync_all().expect("sync");
}
let mut perms = std::fs::metadata(&shim).expect("stat").permissions();
perms.set_mode(0o755);
std::fs::set_permissions(&shim, perms).expect("chmod");
std::env::set_var(crate::apr_bin::APR_BIN_ENV, &shim);
let result = run_apr(&["validate", "/dev/null", "--json"]);
std::env::remove_var(crate::apr_bin::APR_BIN_ENV);
assert!(
result.is_error.is_none(),
"resolved shim should succeed, got: {}",
result.content[0].text
);
assert!(
result.content[0].text.contains("APR-BIN-RESOLVED-SHIM"),
"run_apr must execute the resolved binary; got: {}",
result.content[0].text
);
}
#[test]
fn cancellable_natural_exit_matches_run_apr() {
let (_tx, rx) = mpsc::channel::<()>();
let result = spawn_cancellable("echo", &["hello"], &rx, CANCEL_GRACE_MS);
assert!(result.is_error.is_none(), "echo should succeed");
assert!(result.content[0].text.contains("hello"));
}
#[test]
fn cancellable_disconnected_channel_is_noop() {
let (tx, rx) = mpsc::channel::<()>();
drop(tx);
let result = spawn_cancellable("echo", &["world"], &rx, CANCEL_GRACE_MS);
assert!(result.is_error.is_none());
assert!(result.content[0].text.contains("world"));
}
#[test]
fn cancellable_spawn_failure_maps_to_error() {
let (_tx, rx) = mpsc::channel::<()>();
let result = spawn_cancellable(
"/this/binary/does/not/exist/apr-mcp-test",
&[],
&rx,
CANCEL_GRACE_MS,
);
assert_eq!(result.is_error, Some(true));
assert!(result.content[0].text.contains("Failed to spawn"));
}
#[test]
fn streaming_invokes_callback_per_line() {
let lines = std::sync::Mutex::new(Vec::<String>::new());
let result = spawn_streaming("printf", &["line1\nline2\nline3\n"], |line| {
lines
.lock()
.expect("test mutex not poisoned")
.push(line.to_string());
});
assert!(result.is_error.is_none(), "printf should succeed");
let captured = lines.lock().expect("mutex").clone();
assert_eq!(captured, vec!["line1", "line2", "line3"]);
assert!(result.content[0].text.contains("line1"));
assert!(result.content[0].text.contains("line3"));
}
#[test]
fn streaming_spawn_failure_does_not_call_callback() {
let called = std::sync::Mutex::new(false);
let result = spawn_streaming(
"/this/binary/does/not/exist/apr-mcp-streaming-test",
&[],
|_| {
*called.lock().expect("mutex") = true;
},
);
assert_eq!(result.is_error, Some(true));
assert!(!*called.lock().expect("mutex"));
assert!(result.content[0].text.contains("Failed to spawn"));
}
#[test]
#[cfg(unix)]
fn streaming_nonzero_exit_is_error() {
let result = spawn_streaming("sh", &["-c", "echo partial; exit 3"], |_| {});
assert_eq!(result.is_error, Some(true));
assert!(
result.content[0].text.contains("exit 3"),
"message should include exit code: {}",
result.content[0].text
);
}
#[test]
#[cfg(unix)]
fn cancellable_stops_long_running_subprocess_within_grace() {
let (tx, rx) = mpsc::channel::<()>();
let handle = thread::spawn(move || {
thread::sleep(Duration::from_millis(100));
let _ = tx.send(());
});
let t0 = Instant::now();
let result = spawn_cancellable("sleep", &["60"], &rx, 2_000);
let elapsed = t0.elapsed();
handle.join().expect("cancel-sender thread joins");
assert_eq!(result.is_error, Some(true), "cancelled calls are errors");
assert!(
result.content[0].text.starts_with("Cancelled:"),
"message should indicate cancellation, got: {}",
result.content[0].text
);
assert!(
elapsed < Duration::from_millis(2_500),
"cancel should finish within grace + slack, took {elapsed:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"cancelled call must return far before sleep 60's natural exit"
);
}
}