claux 20260907.0.1

use anyhow::Result;
use async_trait::async_trait;
#[cfg(windows)]
use process_wrap::tokio::JobObject;
use process_wrap::tokio::{CommandWrap, KillOnDrop};
use serde::Deserialize;
use serde_json::{json, Value};
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncReadExt;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;

use super::{Tool, ToolOutput};
use crate::command_sandbox::CommandSandbox;

const CHILD_POLL_INTERVAL: Duration = Duration::from_millis(10);
/// Bytes of each stream retained in memory. Anything past this is drained
/// and counted, never stored, so a runaway command cannot exhaust memory
/// before the tool result is truncated for the model.
const OUTPUT_CAPTURE_LIMIT: usize = 1024 * 1024;
const OUTPUT_DRAIN_TIMEOUT: Duration = Duration::from_secs(1);
const CHILD_REAP_TIMEOUT: Duration = Duration::from_secs(1);

pub struct BashTool {
    sandbox: Arc<CommandSandbox>,
}

impl BashTool {
    pub fn new(sandbox: Arc<CommandSandbox>) -> Self {
        Self { sandbox }
    }
}

#[derive(Deserialize)]
struct Params {
    command: String,
    #[serde(default)]
    timeout: Option<u64>,
    #[serde(default)]
    #[serde(rename = "description")]
    _description: Option<String>,
}

#[async_trait]
impl Tool for BashTool {
    fn name(&self) -> &str {
        "Bash"
    }

    fn description(&self) -> &str {
        "Execute a bash command. Use for git, build tools, or other CLI operations."
    }

    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The bash command to execute"
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in milliseconds (max 600000, default 120000)"
                },
                "description": {
                    "type": "string",
                    "description": "Short description of what the command does"
                }
            },
            "required": ["command"]
        })
    }

    fn is_read_only(&self) -> bool {
        false // conservative default; could be smarter with command analysis
    }

    fn summarize(&self, input: &Value) -> String {
        let cmd = input["command"].as_str().unwrap_or("?");
        // Truncate long commands
        if cmd.len() > 80 {
            format!("{}...", crate::utils::truncate_str(cmd, 77))
        } else {
            cmd.to_string()
        }
    }

    async fn execute(&self, input: Value, cancel: CancellationToken) -> Result<ToolOutput> {
        let params: Params = serde_json::from_value(input)?;

        let timeout_ms = params.timeout.unwrap_or(120_000).min(600_000);
        let timeout = Duration::from_millis(timeout_ms);

        let mut inner = self.sandbox.command(&params.command)?;
        inner.stdout(Stdio::piped()).stderr(Stdio::piped());
        #[cfg(unix)]
        inner.process_group(0);
        let mut command = CommandWrap::from(inner);
        // Commands commonly create descendants (shell pipelines, test runners,
        // build systems). Put the whole tree in one killable unit so cancelling
        // the tool cannot leave grandchildren alive with our pipes still open.
        command.wrap(KillOnDrop);
        #[cfg(windows)]
        command.wrap(JobObject);

        let mut child = match command.spawn() {
            Ok(c) => c,
            Err(e) => {
                return Ok(ToolOutput {
                    content: format!("Failed to execute command: {e}"),
                    is_error: true,
                });
            }
        };
        let process_group = child.id();

        // Take the pipes so we can read them concurrently with wait().
        let mut stdout_pipe = child.stdout().take();
        let mut stderr_pipe = child.stderr().take();

        // Spawn readers so partial output is captured even if we get cancelled
        // or time out mid-stream. Each keeps at most OUTPUT_CAPTURE_LIMIT
        // bytes and keeps draining so the child never blocks on a full pipe.
        let mut stdout_task = tokio::spawn(async move {
            match stdout_pipe.as_mut() {
                Some(p) => read_bounded(p).await,
                None => Captured::default(),
            }
        });
        let mut stderr_task = tokio::spawn(async move {
            match stderr_pipe.as_mut() {
                Some(p) => read_bounded(p).await,
                None => Captured::default(),
            }
        });

        let outcome = tokio::select! {
            status = wait_for_parent(&mut child) => Outcome::Finished(status),
            _ = cancel.cancelled() => Outcome::Cancelled,
            _ = tokio::time::sleep(timeout) => Outcome::TimedOut,
        };

        // A shell can exit successfully while a background descendant remains
        // alive and holds stdout/stderr open. Always terminate anything left
        // in the command's process tree before draining output. Persistent
        // services must detach into their own service manager and redirect
        // their streams rather than inheriting a tool invocation's pipes.
        let residual_processes_terminated = terminate_process_tree(&mut child, process_group);

        if !matches!(outcome, Outcome::Finished(_)) {
            let _ = tokio::time::timeout(CHILD_REAP_TIMEOUT, wait_for_parent(&mut child)).await;
        }

        let (stdout, stdout_abandoned) = drain_reader(&mut stdout_task).await;
        let (stderr, stderr_abandoned) = drain_reader(&mut stderr_task).await;
        let output_abandoned = stdout_abandoned || stderr_abandoned;
        let dropped_bytes = stdout.dropped + stderr.dropped;

        let stdout_s = render_output(&stdout.bytes, "stdout");
        let stderr_s = render_output(&stderr.bytes, "stderr");

        let mut content = String::new();
        if !stdout_s.is_empty() {
            content.push_str(&stdout_s);
        }
        if !stderr_s.is_empty() {
            if !content.is_empty() {
                content.push('\n');
            }
            content.push_str(&stderr_s);
        }

        let mut is_error = match &outcome {
            Outcome::Finished(Ok(status)) => {
                if !status.success() {
                    content.push_str(&format!("\nExit code: {status}"));
                }
                !status.success()
            }
            Outcome::Finished(Err(e)) => {
                if !content.is_empty() {
                    content.push('\n');
                }
                content.push_str(&format!("wait error: {e}"));
                true
            }
            Outcome::Cancelled => {
                if !content.is_empty() {
                    content.push('\n');
                }
                content.push_str("Interrupted by user.");
                true
            }
            Outcome::TimedOut => {
                if !content.is_empty() {
                    content.push('\n');
                }
                content.push_str(&format!("Command timed out after {timeout_ms}ms"));
                true
            }
        };

        if matches!(outcome, Outcome::Finished(Ok(status)) if status.success())
            && residual_processes_terminated
        {
            if !content.is_empty() {
                content.push('\n');
            }
            content.push_str(
                "Background processes were terminated when the command exited. Use a service manager for persistent processes.",
            );
            is_error = true;
        }

        if output_abandoned {
            if !content.is_empty() {
                content.push('\n');
            }
            content.push_str(
                "A detached process kept command output open; Claux stopped waiting for its output.",
            );
            is_error = true;
        }

        if content.len() > 100_000 {
            content.truncate(100_000);
            content.push_str("\n... (output truncated)");
        }
        // After the model-facing truncation so the note itself survives.
        if dropped_bytes > 0 {
            content.push_str(&format!(
                "\n... ({dropped_bytes} bytes of output beyond the capture limit were not retained)"
            ));
        }

        Ok(ToolOutput { content, is_error })
    }
}

async fn wait_for_parent(
    child: &mut Box<dyn process_wrap::tokio::ChildWrapper>,
) -> std::io::Result<std::process::ExitStatus> {
    loop {
        if let Some(status) = child.try_wait()? {
            return Ok(status);
        }
        tokio::time::sleep(CHILD_POLL_INTERVAL).await;
    }
}

#[cfg(unix)]
fn terminate_process_tree(
    _child: &mut Box<dyn process_wrap::tokio::ChildWrapper>,
    process_group: Option<u32>,
) -> bool {
    use nix::sys::signal::{killpg, Signal};
    use nix::unistd::Pid;

    process_group
        .and_then(|pid| i32::try_from(pid).ok())
        .is_some_and(|pid| killpg(Pid::from_raw(pid), Signal::SIGKILL).is_ok())
}

#[cfg(not(unix))]
fn terminate_process_tree(
    child: &mut Box<dyn process_wrap::tokio::ChildWrapper>,
    _process_group: Option<u32>,
) -> bool {
    child.start_kill().is_ok()
}

/// Output retained from one stream plus how many bytes were drained past
/// the capture limit.
#[derive(Default)]
struct Captured {
    bytes: Vec<u8>,
    dropped: u64,
}

/// Read a stream to EOF, retaining at most OUTPUT_CAPTURE_LIMIT bytes.
async fn read_bounded<R: tokio::io::AsyncRead + Unpin>(reader: &mut R) -> Captured {
    let mut captured = Captured::default();
    let mut chunk = [0_u8; 16 * 1024];
    loop {
        let read = match reader.read(&mut chunk).await {
            Ok(0) | Err(_) => return captured,
            Ok(read) => read,
        };
        let room = OUTPUT_CAPTURE_LIMIT.saturating_sub(captured.bytes.len());
        let keep = read.min(room);
        captured.bytes.extend_from_slice(&chunk[..keep]);
        captured.dropped += (read - keep) as u64;
    }
}

async fn drain_reader(task: &mut JoinHandle<Captured>) -> (Captured, bool) {
    match tokio::time::timeout(OUTPUT_DRAIN_TIMEOUT, &mut *task).await {
        Ok(Ok(output)) => (output, false),
        Ok(Err(_)) => (Captured::default(), false),
        Err(_) => {
            task.abort();
            (Captured::default(), true)
        }
    }
}

fn render_output(bytes: &[u8], stream: &str) -> String {
    match std::str::from_utf8(bytes) {
        Ok(text) if !bytes.contains(&0) => text.to_string(),
        _ => format!("[binary {stream} suppressed: {} bytes]", bytes.len()),
    }
}

enum Outcome {
    Finished(std::io::Result<std::process::ExitStatus>),
    Cancelled,
    TimedOut,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn token() -> CancellationToken {
        CancellationToken::new()
    }

    fn tool() -> BashTool {
        BashTool::new(Arc::new(CommandSandbox::unrestricted_for_tests()))
    }

    #[tokio::test]
    async fn bash_echo() {
        let tool = tool();
        let result = tool
            .execute(json!({"command": "echo hello"}), token())
            .await
            .unwrap();
        assert!(!result.is_error);
        assert!(result.content.trim().contains("hello"));
    }

    #[tokio::test]
    async fn bash_output_beyond_the_capture_limit_is_drained_not_stored() {
        let tool = tool();
        // 8 MiB of output: far beyond the 1 MiB capture limit and the
        // 100 KB model-facing truncation. The command must finish (the pipe
        // is drained), the result must stay bounded, and the drop must be
        // reported.
        let result = tool
            .execute(
                json!({"command": "head -c 8388608 /dev/zero | tr '\\0' 'x'"}),
                token(),
            )
            .await
            .unwrap();
        assert!(
            result.content.len() <= 100_000 + 160,
            "{}",
            result.content.len()
        );
        assert!(
            result
                .content
                .contains("bytes of output beyond the capture limit were not retained"),
            "{}",
            &result.content[result.content.len().saturating_sub(200)..]
        );
        assert!(!result.content.contains("timed out"));
    }

    #[tokio::test]
    async fn bash_exit_code() {
        let tool = tool();
        let result = tool
            .execute(json!({"command": "exit 1"}), token())
            .await
            .unwrap();
        assert!(result.is_error);
        assert!(result.content.contains("Exit code"));
    }

    #[tokio::test]
    async fn bash_captures_stderr() {
        let tool = tool();
        let result = tool
            .execute(json!({"command": "echo err >&2"}), token())
            .await
            .unwrap();
        assert!(result.content.contains("err"));
    }

    #[test]
    fn preserves_utf8_output() {
        assert_eq!(render_output("héllo\n".as_bytes(), "stdout"), "héllo\n");
    }

    #[test]
    fn suppresses_invalid_utf8_output() {
        assert_eq!(
            render_output(&[0xff, 0xfe, 0xfd], "stdout"),
            "[binary stdout suppressed: 3 bytes]"
        );
    }

    #[test]
    fn suppresses_nul_containing_output() {
        assert_eq!(
            render_output(b"text\0more", "stderr"),
            "[binary stderr suppressed: 9 bytes]"
        );
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn binary_stdout_does_not_hide_text_stderr() {
        let result = tool()
            .execute(
                json!({"command": "printf '\\377'; printf 'warning' >&2"}),
                token(),
            )
            .await
            .unwrap();

        assert!(!result.is_error);
        assert_eq!(
            result.content,
            "[binary stdout suppressed: 1 bytes]\nwarning"
        );
        assert!(!result.content.contains('\u{fffd}'));
    }

    #[tokio::test]
    async fn bash_timeout() {
        let tool = tool();
        let result = tool
            .execute(json!({"command": "sleep 10", "timeout": 100}), token())
            .await
            .unwrap();
        assert!(result.is_error);
        assert!(result.content.contains("timed out"));
    }

    #[tokio::test]
    async fn bash_cancellation() {
        let tool = tool();
        let cancel = CancellationToken::new();
        let cancel_clone = cancel.clone();
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(100)).await;
            cancel_clone.cancel();
        });
        let start = std::time::Instant::now();
        let result = tool
            .execute(
                json!({
                    "command": "trap '' HUP; sleep 30 & wait",
                    "timeout": 60000
                }),
                cancel,
            )
            .await
            .unwrap();
        assert!(
            start.elapsed() < Duration::from_secs(3),
            "cancellation should kill the entire process tree (took {:?})",
            start.elapsed()
        );
        assert!(result.is_error);
        assert!(result.content.contains("Interrupted"));
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn background_process_cannot_hold_a_completed_command_open() {
        let start = std::time::Instant::now();
        let result = tool()
            .execute(
                json!({
                    "command": "nohup sleep 30 >/dev/null 2>&1 & printf ready",
                    "timeout": 60000
                }),
                token(),
            )
            .await
            .unwrap();

        assert!(
            start.elapsed() < Duration::from_secs(3),
            "background descendants must not hold the tool open (took {:?})",
            start.elapsed()
        );
        assert!(result.is_error);
        assert!(result.content.contains("ready"));
        assert!(result
            .content
            .contains("Background processes were terminated"));
    }

    #[cfg(target_os = "linux")]
    #[tokio::test]
    async fn detached_process_cannot_hold_output_capture_open() {
        let start = std::time::Instant::now();
        let result = tool()
            .execute(
                json!({
                    "command": "setsid sh -c 'sleep 2' & printf ready",
                    "timeout": 60000
                }),
                token(),
            )
            .await
            .unwrap();

        assert!(
            start.elapsed() < Duration::from_secs(3),
            "escaped descendants must not hold the tool open (took {:?})",
            start.elapsed()
        );
        assert!(result.is_error);
        assert!(result.content.contains("stopped waiting"));
    }
}