heartbit-core 2026.507.3

use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Mutex;

use serde_json::json;

use crate::error::Error;
use crate::llm::types::ToolDefinition;
use crate::tool::{Tool, ToolOutput};

const DEFAULT_TIMEOUT_MS: u64 = 120_000;
const MAX_TIMEOUT_MS: u64 = 600_000;
const MAX_OUTPUT_CHARS: usize = 30_000;
const HEAD_TAIL_SIZE: usize = 14_000;

/// Builtin tool that executes shell commands in a persistent working directory.
///
/// Each agent session gets one `BashTool` instance; `cd` commands update an
/// internal `cwd` that survives across calls within the session. On Linux with
/// the `sandbox` feature the tool can be restricted to a Landlock filesystem
/// policy via `with_sandbox_policy`; the `workspace` field additionally jails
/// the working directory to a project root. Output is capped and head-tailed
/// when it exceeds `MAX_OUTPUT_CHARS` so the LLM context is not overwhelmed.
pub struct BashTool {
    /// Tracked working directory that persists across calls.
    cwd: Mutex<PathBuf>,
    /// Optional workspace root. When set, cwd is jailed to this directory.
    workspace: Option<PathBuf>,
    /// Controls which environment variables are visible to bash subprocesses.
    env_policy: crate::workspace::EnvPolicy,
    /// Landlock filesystem sandbox policy (Linux only, feature-gated).
    #[cfg(all(target_os = "linux", feature = "sandbox"))]
    sandbox_policy: Option<crate::sandbox::SandboxPolicy>,
    /// App-layer path policy. Applied to the working directory before spawning.
    /// Compose with `with_sandbox_policy` for kernel-level enforcement on Linux.
    path_policy: Option<std::sync::Arc<crate::sandbox::CorePathPolicy>>,
}

/// Emit a one-time warning when BashTool is constructed without a kernel
/// sandbox. Helps operators notice the absence of Landlock on macOS / Windows
/// or when the `sandbox` feature was disabled.
fn warn_kernel_sandbox_status() {
    use std::sync::atomic::{AtomicBool, Ordering};
    static WARNED: AtomicBool = AtomicBool::new(false);
    if WARNED
        .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
        .is_err()
    {
        return;
    }
    #[cfg(all(target_os = "linux", feature = "sandbox"))]
    {
        tracing::debug!("BashTool: Linux Landlock kernel sandbox available");
    }
    #[cfg(all(target_os = "linux", not(feature = "sandbox")))]
    {
        tracing::warn!(
            "BashTool: built without `sandbox` feature on Linux. \
             The kernel-level Landlock filesystem isolation is OFF; \
             agents can read/write the entire filesystem under this process's \
             identity. Rebuild with --features sandbox to enable. (F-FS-5)"
        );
    }
    #[cfg(not(target_os = "linux"))]
    {
        tracing::warn!(
            "BashTool: kernel-level filesystem sandbox is unavailable on this \
             platform. Only application-layer `path_policy` enforcement is \
             active. For multi-tenant deployments, prefer Linux + the \
             `sandbox` feature. (F-FS-5)"
        );
    }
}

impl BashTool {
    pub fn new() -> Self {
        warn_kernel_sandbox_status();
        let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/"));
        Self {
            cwd: Mutex::new(cwd),
            workspace: None,
            env_policy: crate::workspace::EnvPolicy::Inherit,
            #[cfg(all(target_os = "linux", feature = "sandbox"))]
            sandbox_policy: None,
            path_policy: None,
        }
    }

    /// Create a BashTool with workspace jailing and optional env variable filtering.
    pub fn with_sandbox(workspace: PathBuf, env_policy: crate::workspace::EnvPolicy) -> Self {
        warn_kernel_sandbox_status();
        Self {
            cwd: Mutex::new(workspace.clone()),
            workspace: Some(workspace),
            env_policy,
            #[cfg(all(target_os = "linux", feature = "sandbox"))]
            sandbox_policy: None,
            path_policy: None,
        }
    }

    /// Set a Landlock filesystem sandbox policy (Linux only).
    #[cfg(all(target_os = "linux", feature = "sandbox"))]
    pub fn with_sandbox_policy(mut self, policy: crate::sandbox::SandboxPolicy) -> Self {
        self.sandbox_policy = Some(policy);
        self
    }

    /// Set a `CorePathPolicy` that restricts the bash subprocess's working directory.
    /// Application-layer; complements `with_sandbox_policy` (Linux Landlock).
    pub fn with_path_policy(
        mut self,
        policy: std::sync::Arc<crate::sandbox::CorePathPolicy>,
    ) -> Self {
        self.path_policy = Some(policy);
        self
    }
}

impl Tool for BashTool {
    fn definition(&self) -> ToolDefinition {
        ToolDefinition {
            name: "bash".into(),
            description: "Execute a bash command. Working directory persists between calls. \
                          Captures stdout and stderr. Default timeout: 120s, max: 600s."
                .into(),
            input_schema: json!({
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The bash command to execute"
                    },
                    "timeout": {
                        "type": "number",
                        "description": "Timeout in milliseconds (default 120000, max 600000)"
                    }
                },
                "required": ["command"]
            }),
        }
    }

    fn execute(
        &self,
        _ctx: &crate::ExecutionContext,
        input: serde_json::Value,
    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, Error>> + Send + '_>> {
        Box::pin(async move {
            let command = input
                .get("command")
                .and_then(|v| v.as_str())
                .ok_or_else(|| Error::Agent("command is required".into()))?;

            let timeout_ms = input
                .get("timeout")
                .and_then(|v| v.as_u64())
                .unwrap_or(DEFAULT_TIMEOUT_MS)
                .min(MAX_TIMEOUT_MS);

            let cwd = {
                let guard = self.cwd.lock().expect("bash cwd lock poisoned");
                guard.clone()
            };

            if let Some(policy) = &self.path_policy
                && let Err(e) = policy.check_path(&cwd)
            {
                return Ok(ToolOutput::error(format!("path policy: {e}")));
            }

            // Build the command: cd to tracked cwd, run user command, then print pwd.
            // Detect trailing `&` to avoid `{ cmd &; }` syntax error
            // (`&` is a command terminator — `&;` is invalid bash).
            let wrapped = if command.trim_end().ends_with('&') {
                format!("{{ {} }}", command)
            } else {
                format!("{{ {}; }}", command)
            };
            // SECURITY (F-FS-8): use a per-spawn nonce in the cwd marker so
            // the running command cannot forge it via `echo`. The previous
            // fixed `__HEARTBIT_CWD__=` was emit-able from any echo and let
            // a hostile command mutate the tracked cwd to e.g. `/etc`.
            //
            // PERF (P-TOOL-10): generate the random base **once** per
            // process via `LazyLock` and combine it with an atomic
            // counter for each spawn. Per-call cost is one atomic
            // increment + one short `format!()` instead of a fresh
            // UUID v4 generation (~5–50 µs saved per bash call). The
            // nonce is still unforgeable: the attacker (the bash command
            // itself) cannot observe the per-process random base.
            static CWD_MARKER_BASE: std::sync::LazyLock<String> = std::sync::LazyLock::new(|| {
                let nonce = uuid::Uuid::new_v4().simple();
                format!("__HEARTBIT_CWD_{nonce}_")
            });
            static CWD_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
            let counter = CWD_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            let cwd_marker = format!("{}{counter:x}__", *CWD_MARKER_BASE);
            let full_command = format!(
                "cd {} && {}; __exit_code=$?; echo; echo \"{cwd_marker}=$(pwd)\"; exit $__exit_code",
                shell_escape(&cwd.display().to_string()),
                wrapped
            );

            let mut cmd = tokio::process::Command::new("bash");
            cmd.arg("-c")
                .arg(&full_command)
                .stdout(std::process::Stdio::piped())
                .stderr(std::process::Stdio::piped())
                .kill_on_drop(true);

            match &self.env_policy {
                crate::workspace::EnvPolicy::Inherit => {}
                crate::workspace::EnvPolicy::Allowlist(allowed) => {
                    cmd.env_clear();
                    for key in allowed {
                        if let Ok(val) = std::env::var(key) {
                            cmd.env(key, &val);
                        }
                    }
                }
            }

            // Apply Landlock sandbox if configured (Linux only).
            #[cfg(all(target_os = "linux", feature = "sandbox"))]
            if let Some(ref policy) = self.sandbox_policy {
                use std::os::unix::process::CommandExt as _;
                let pre_exec_fn = policy
                    .clone()
                    .into_pre_exec()
                    .map_err(|e| Error::Agent(format!("Sandbox setup failed: {e}")))?;
                // SAFETY: pre_exec runs between fork() and exec() in the child.
                // The closure only calls Landlock syscalls — no allocations.
                unsafe {
                    cmd.as_std_mut().pre_exec(pre_exec_fn);
                }
            }

            let child = cmd
                .spawn()
                .map_err(|e| Error::Agent(format!("Failed to spawn bash: {e}")))?;

            let timeout_duration = std::time::Duration::from_millis(timeout_ms);

            let output =
                match tokio::time::timeout(timeout_duration, child.wait_with_output()).await {
                    Ok(Ok(output)) => output,
                    Ok(Err(e)) => return Ok(ToolOutput::error(format!("Command failed: {e}"))),
                    Err(_) => {
                        // kill_on_drop ensures cleanup when `child` is dropped here
                        return Ok(ToolOutput::error(format!(
                            "Command timed out after {timeout_ms}ms"
                        )));
                    }
                };

            let exit_code = output.status.code().unwrap_or(-1);
            let stdout = String::from_utf8_lossy(&output.stdout).to_string();
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();

            // Extract new cwd from the output (using the per-spawn nonce —
            // F-FS-8 — so a hostile echo cannot inject a fake marker).
            let (user_stdout, new_cwd) = extract_cwd_marker(&stdout, &cwd_marker);
            if let Some(new_dir) = new_cwd {
                let new_path = PathBuf::from(&new_dir);
                let mut guard = self.cwd.lock().expect("bash cwd lock poisoned");

                if let Some(ref ws) = self.workspace {
                    // Jail check: pwd output is already absolute/canonical,
                    // so a direct starts_with is sufficient (no extra I/O).
                    if new_path.starts_with(ws) {
                        *guard = new_path;
                    }
                    // If outside workspace, silently keep old cwd
                } else {
                    *guard = new_path;
                }
            }

            // Combine output
            let mut combined = String::new();
            if !user_stdout.is_empty() {
                combined.push_str(&user_stdout);
            }
            if !stderr.is_empty() {
                if !combined.is_empty() {
                    combined.push('\n');
                }
                combined.push_str(&stderr);
            }

            // Truncate if needed
            let combined = truncate_middle(&combined, MAX_OUTPUT_CHARS);

            let exit_info = format!("\n\n(exit code: {exit_code})");
            let output_text = format!("{combined}{exit_info}");

            if exit_code == 0 {
                Ok(ToolOutput::success(output_text))
            } else {
                Ok(ToolOutput::error(output_text))
            }
        })
    }
}

/// Extract the cwd marker from stdout, returning (user output, optional new cwd).
///
/// SECURITY (F-FS-8): pass the per-spawn nonce-bearing marker so a user
/// command's `echo __HEARTBIT_CWD_xxx__=/etc` cannot forge a hijack of the
/// tracked cwd; only the marker generated at spawn time matches.
fn extract_cwd_marker(stdout: &str, marker: &str) -> (String, Option<String>) {
    let needle = format!("{marker}=");
    if let Some(marker_pos) = stdout.rfind(&needle) {
        let user_output = stdout[..marker_pos].trim_end().to_string();
        let cwd_line = &stdout[marker_pos + needle.len()..];
        let cwd = cwd_line.trim().to_string();
        if cwd.is_empty() {
            (user_output, None)
        } else {
            (user_output, Some(cwd))
        }
    } else {
        (stdout.to_string(), None)
    }
}

/// Legacy extract — kept for tests, uses the unguarded fixed marker.
#[cfg(test)]
fn extract_cwd(stdout: &str) -> (String, Option<String>) {
    extract_cwd_marker(stdout, "__HEARTBIT_CWD__")
}

/// Truncate output with middle omission if it exceeds max bytes.
fn truncate_middle(text: &str, max_bytes: usize) -> String {
    if text.len() <= max_bytes {
        return text.to_string();
    }

    let total_lines = text.lines().count();

    // Find char boundaries for head/tail slices
    let head_end = super::floor_char_boundary(text, HEAD_TAIL_SIZE.min(text.len()));
    let tail_start = ceil_char_boundary(text, text.len().saturating_sub(HEAD_TAIL_SIZE));
    let head = &text[..head_end];
    let tail = &text[tail_start..];

    // Count omitted lines (approximate)
    let head_lines = head.lines().count();
    let tail_lines = tail.lines().count();
    let omitted = total_lines.saturating_sub(head_lines + tail_lines);

    format!("{head}\n\n... [{omitted} lines truncated] ...\n\n{tail}")
}

/// Find the smallest char boundary >= target.
fn ceil_char_boundary(text: &str, target: usize) -> usize {
    let mut pos = target.min(text.len());
    while pos < text.len() && !text.is_char_boundary(pos) {
        pos += 1;
    }
    pos
}

fn shell_escape(s: &str) -> String {
    format!("'{}'", s.replace('\'', "'\\''"))
}

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

    #[test]
    fn definition_has_correct_name() {
        let tool = BashTool::new();
        assert_eq!(tool.definition().name, "bash");
    }

    #[tokio::test]
    async fn bash_echo() {
        let tool = BashTool::new();
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "echo hello"}),
            )
            .await
            .unwrap();
        assert!(!result.is_error, "got error: {}", result.content);
        assert!(result.content.contains("hello"));
        assert!(result.content.contains("exit code: 0"));
    }

    #[tokio::test]
    async fn bash_failing_command() {
        let tool = BashTool::new();
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "exit 42"}),
            )
            .await
            .unwrap();
        assert!(result.is_error);
        assert!(result.content.contains("exit code: 42"));
    }

    #[tokio::test]
    async fn bash_preserves_cwd() {
        let dir = tempfile::tempdir().unwrap();
        let tool = BashTool::new();

        // Change directory
        tool.execute(
            &crate::ExecutionContext::default(),
            json!({"command": format!("cd {}", dir.path().display())}),
        )
        .await
        .unwrap();

        // Verify pwd shows the new directory
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "pwd"}),
            )
            .await
            .unwrap();
        assert!(
            result.content.contains(&dir.path().display().to_string()),
            "expected cwd to be {}, got: {}",
            dir.path().display(),
            result.content
        );
    }

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

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

    #[test]
    fn extract_cwd_parses_marker() {
        let stdout = "some output\n__HEARTBIT_CWD__=/home/user\n";
        let (user, cwd) = extract_cwd(stdout);
        assert_eq!(user, "some output");
        assert_eq!(cwd, Some("/home/user".into()));
    }

    #[test]
    fn extract_cwd_no_marker() {
        let stdout = "just output\n";
        let (user, cwd) = extract_cwd(stdout);
        assert_eq!(user, "just output\n");
        assert!(cwd.is_none());
    }

    #[test]
    fn truncate_middle_short_text() {
        let text = "hello world";
        assert_eq!(truncate_middle(text, 100), text);
    }

    #[test]
    fn truncate_middle_long_text() {
        let text = "a\n".repeat(20_000);
        let result = truncate_middle(&text, MAX_OUTPUT_CHARS);
        assert!(result.contains("truncated"));
        assert!(result.len() < text.len());
    }

    #[test]
    fn truncate_middle_multibyte_utf8() {
        // Each '🦀' is 4 bytes. Build a string that must be truncated.
        let text = "🦀".repeat(10_000); // 40,000 bytes
        let result = truncate_middle(&text, MAX_OUTPUT_CHARS);
        // Should not panic and should be valid UTF-8
        assert!(result.contains("truncated"));
        // Verify it's valid UTF-8 (String is UTF-8 by construction, but let's check boundaries)
        assert!(result.len() < text.len());
    }

    #[test]
    fn shell_escape_simple_path() {
        assert_eq!(shell_escape("/tmp/foo"), "'/tmp/foo'");
    }

    #[test]
    fn shell_escape_path_with_single_quote() {
        // dir's should become 'dir'\''s'
        assert_eq!(shell_escape("dir's"), "'dir'\\''s'");
    }

    #[test]
    fn shell_escape_path_with_spaces() {
        assert_eq!(shell_escape("/tmp/my dir"), "'/tmp/my dir'");
    }

    #[test]
    fn shell_escape_empty_string() {
        assert_eq!(shell_escape(""), "''");
    }

    #[tokio::test]
    async fn bash_workspace_cd_outside_rejected() {
        let dir = tempfile::tempdir().unwrap();
        let ws = dir.path().canonicalize().unwrap();
        let tool = BashTool::with_sandbox(ws.clone(), crate::workspace::EnvPolicy::Inherit);

        // Try to cd outside workspace
        tool.execute(
            &crate::ExecutionContext::default(),
            json!({"command": "cd /tmp"}),
        )
        .await
        .unwrap();

        // Verify cwd stayed inside workspace
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "pwd"}),
            )
            .await
            .unwrap();
        assert!(
            result.content.contains(&ws.display().to_string()),
            "cwd should stay in workspace after cd /tmp, got: {}",
            result.content
        );
    }

    #[tokio::test]
    async fn bash_workspace_cd_inside_allowed() {
        let dir = tempfile::tempdir().unwrap();
        let ws = dir.path().canonicalize().unwrap();
        let sub = ws.join("subdir");
        std::fs::create_dir(&sub).unwrap();

        let tool = BashTool::with_sandbox(ws, crate::workspace::EnvPolicy::Inherit);

        tool.execute(
            &crate::ExecutionContext::default(),
            json!({"command": "cd subdir"}),
        )
        .await
        .unwrap();

        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "pwd"}),
            )
            .await
            .unwrap();
        assert!(
            result.content.contains("subdir"),
            "cwd should be in subdir, got: {}",
            result.content
        );
    }

    #[tokio::test]
    async fn bash_trailing_ampersand_no_syntax_error() {
        let tool = BashTool::new();
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "sleep 0.01 &"}),
            )
            .await
            .unwrap();
        assert!(
            !result.is_error,
            "trailing & should not cause syntax error: {}",
            result.content
        );
        assert!(result.content.contains("exit code: 0"));
    }

    #[tokio::test]
    async fn bash_background_with_foreground() {
        let tool = BashTool::new();
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "echo before & echo after"}),
            )
            .await
            .unwrap();
        assert!(
            !result.is_error,
            "background & foreground should work: {}",
            result.content
        );
        assert!(result.content.contains("after"));
    }

    #[tokio::test]
    async fn bash_with_sandbox_env_filtering() {
        use crate::workspace::EnvPolicy;

        // SAFETY: set_var is unsafe in multi-threaded contexts but acceptable in tests
        unsafe { std::env::set_var("__HEARTBIT_TEST_SECRET", "super_secret_123") };
        let dir = tempfile::tempdir().unwrap();
        let tool = BashTool::with_sandbox(
            dir.path().canonicalize().unwrap(),
            EnvPolicy::Allowlist(vec!["PATH".into(), "HOME".into()]),
        );
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "echo $__HEARTBIT_TEST_SECRET"}),
            )
            .await
            .unwrap();
        // Should NOT contain the secret since it's not in the allowlist
        assert!(
            !result.content.contains("super_secret_123"),
            "env var should be filtered: {}",
            result.content
        );
        unsafe { std::env::remove_var("__HEARTBIT_TEST_SECRET") };
    }

    #[tokio::test]
    async fn bash_inherit_env_passes_vars() {
        // SAFETY: set_var is unsafe in multi-threaded contexts but acceptable in tests
        unsafe { std::env::set_var("__HEARTBIT_TEST_VAR", "visible_value") };
        let tool = BashTool::new();
        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "echo $__HEARTBIT_TEST_VAR"}),
            )
            .await
            .unwrap();
        assert!(result.content.contains("visible_value"));
        unsafe { std::env::remove_var("__HEARTBIT_TEST_VAR") };
    }

    #[tokio::test]
    async fn bash_with_path_policy_rejects_outside_cwd() {
        use crate::sandbox::CorePathPolicy;
        use std::sync::Arc;

        let inside = tempfile::tempdir().unwrap();
        let outside = tempfile::tempdir().unwrap();
        let policy = Arc::new(
            CorePathPolicy::builder()
                .allow_dir(inside.path())
                .build()
                .unwrap(),
        );

        // Set cwd to outside the policy
        let tool = BashTool::with_sandbox(
            outside.path().canonicalize().unwrap(),
            crate::workspace::EnvPolicy::Inherit,
        )
        .with_path_policy(policy);

        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "echo hello"}),
            )
            .await
            .unwrap();
        assert!(
            result.is_error,
            "expected path policy error, got: {:?}",
            result.content
        );
        assert!(
            result.content.contains("path policy"),
            "expected path policy error, got: {:?}",
            result.content
        );
    }

    #[tokio::test]
    async fn bash_with_path_policy_allows_inside_cwd() {
        use crate::sandbox::CorePathPolicy;
        use std::sync::Arc;

        let inside = tempfile::tempdir().unwrap();
        let policy = Arc::new(
            CorePathPolicy::builder()
                .allow_dir(inside.path())
                .build()
                .unwrap(),
        );

        // Set cwd to inside the policy
        let tool = BashTool::with_sandbox(
            inside.path().canonicalize().unwrap(),
            crate::workspace::EnvPolicy::Inherit,
        )
        .with_path_policy(policy);

        let result = tool
            .execute(
                &crate::ExecutionContext::default(),
                json!({"command": "echo hello"}),
            )
            .await
            .unwrap();
        assert!(
            !result.is_error,
            "expected success, got: {:?}",
            result.content
        );
        assert!(result.content.contains("hello"));
    }
}