Merge pull request #70 from Sewer56/optimize-a-bit

Changed: Trim hot path overhead across bash and sandbox tools

Author: Sewer56

src/llm-coding-tools-bubblewrap/src/probe.rs

@@ -5,8 +5,9 @@
 use crate::path_util::normalize_path;
 use crate::{Availability, LinuxBwrapError, Preset};
 use parking_lot::RwLock;
+use std::collections::HashSet;
 use std::env;
-use std::ffi::OsString;
+use std::ffi::{OsStr, OsString};
 use std::path::{Path, PathBuf};
 use std::process::{Command, Output, Stdio};
 use std::sync::Arc;
@@ -91,10 +92,10 @@ fn profile_name(preset: Option<Preset>) -> &'static str {
     }
 }
 
-/// Searches `PATH` directories for a file named `name` and returns the first match.
-pub(crate) fn find_binary_on_path(name: &str) -> Option<Box<Path>> {
-    let path = env::var_os("PATH")?;
-    for dir in env::split_paths(&path) {
+#[inline]
+fn find_binary_on_path_in(name: &str, path: Option<&OsStr>) -> Option<Box<Path>> {
+    let path = path?;
+    for dir in env::split_paths(path) {
         if !dir.is_absolute() || dir.as_os_str().is_empty() {
             continue;
         }
@@ -115,29 +116,59 @@ pub(crate) fn first_shell_candidate_with<F, R>(mut classify: F) -> Option<(Box<P
 where
     F: FnMut(&Path) -> Option<R>,
 {
+    first_shell_candidate_with_in(env::var_os("PATH").as_deref(), &mut classify)
+}
+
+fn first_shell_candidate_with_in<F, R>(
+    env_path: Option<&OsStr>,
+    classify: &mut F,
+) -> Option<(Box<Path>, R)>
+where
+    F: FnMut(&Path) -> Option<R>,
+{
+    let mut seen = HashSet::with_capacity(10);
+
     for name in ["bash", "sh"] {
-        if let Some(path) = find_binary_on_path(name) {
-            let path = normalize_path(path.as_ref());
-            if let Some(r) = classify(path.as_ref()) {
-                return Some((path, r));
+        if let Some(shell_path) = find_binary_on_path_in(name, env_path) {
+            if let Some(result) = classify_shell_candidate(classify, &mut seen, shell_path) {
+                return Some(result);
             }
         }
     }
+
     for candidate in SHELL_CANDIDATES {
-        let path = PathBuf::from(candidate);
-        if path.is_file() {
-            let path = normalize_path(&path);
-            if let Some(r) = classify(path.as_ref()) {
-                return Some((path, r));
+        let candidate_path = PathBuf::from(candidate);
+        if candidate_path.is_file() {
+            if let Some(result) =
+                classify_shell_candidate(classify, &mut seen, candidate_path.into_boxed_path())
+            {
+                return Some(result);
             }
         }
     }
+
     None
 }
 
-/// Returns any available host shell (`bash` preferred, then `sh`).
-pub(crate) fn resolve_host_shell() -> Option<Box<Path>> {
-    first_shell_candidate_with(|_| Some(())).map(|(path, _)| path)
+#[inline]
+fn classify_shell_candidate<F, R>(
+    classify: &mut F,
+    seen: &mut HashSet<Box<Path>>,
+    path: Box<Path>,
+) -> Option<(Box<Path>, R)>
+where
+    F: FnMut(&Path) -> Option<R>,
+{
+    let path = normalize_path(path.as_ref());
+    if !seen.insert(path.clone()) {
+        return None;
+    }
+    classify(path.as_ref()).map(|result| (path, result))
+}
+
+#[inline]
+fn resolve_host_shell_in(path: Option<&OsStr>) -> Option<Box<Path>> {
+    first_shell_candidate_with_in(path, &mut |_| Some(())).map(|(path, _)| path)
 }
 
 fn probe_backend() -> LinuxBwrapBackend {
@@ -157,7 +188,7 @@ fn probe_backend() -> LinuxBwrapBackend {
         }
     }
 
-    let backend = probe_backend_uncached();
+    let backend = probe_backend_uncached(path.as_deref());
     *cache.write() = Some((path, backend.clone()));
     backend
 }
@@ -166,31 +197,14 @@ fn probe_backend() -> LinuxBwrapBackend {
 ///
 /// The probe binds the host root read-only and runs a tiny shell command. That
 /// checks both namespace support and shell visibility on FHS and Nix systems.
-fn probe_backend_uncached() -> LinuxBwrapBackend {
-    let Some(bwrap) = find_binary_on_path("bwrap") else {
+fn probe_backend_uncached(path: Option<&OsStr>) -> LinuxBwrapBackend {
+    let Some(bwrap) = find_binary_on_path_in("bwrap", path) else {
         return LinuxBwrapBackend::MissingBinary {
             reason: Box::from("`bwrap` was not found on PATH"),
         };
     };
 
-    let version = Command::new(bwrap.as_os_str())
-        .arg("--version")
-        .stdin(Stdio::null())
-        .stdout(Stdio::null())
-        .stderr(Stdio::piped())
-        .output();
-    let Ok(version) = version else {
-        return LinuxBwrapBackend::MissingBinary {
-            reason: format!("failed to execute {}", bwrap.display()).into_boxed_str(),
-        };
-    };
-    if !version.status.success() {
-        return LinuxBwrapBackend::Unusable {
-            reason: probe_failure_reason(&version, "`bwrap --version` failed"),
-        };
-    }
-
-    let Some(shell) = resolve_host_shell() else {
+    let Some(shell) = resolve_host_shell_in(path) else {
         return LinuxBwrapBackend::Unusable {
             reason: Box::from("no usable host shell (`bash` or `sh`) was found"),
         };
@@ -257,6 +271,11 @@ mod tests {
     use serial_test::serial;
     use tempfile::TempDir;
 
+    /// Searches `PATH` directories for a file named `name` and returns the first match.
+    fn find_binary_on_path(name: &str) -> Option<Box<Path>> {
+        find_binary_on_path_in(name, env::var_os("PATH").as_deref())
+    }
+
     // These tests swap PATH and exercise a process-wide availability cache, so
     // cases that probe `bwrap` run serially to avoid cross-test contamination.
 
@@ -298,14 +317,6 @@ mod tests {
         // Make `bwrap` look installed, then fail the "can it sandbox?" probe.
         let script = format!(
             r#"#!/bin/sh
-# Handle --version probe
-for arg in "$@"; do
-    if [ "$arg" = "--version" ]; then
-        echo "bubblewrap 0.8.0"
-        exit 0
-    fi
-done
-# Fail the "can it sandbox?" probe
 echo "{}" >&2
 exit 1
 "#,

src/llm-coding-tools-bubblewrap/src/profile/types.rs

@@ -468,6 +468,25 @@ impl Profile {
         )))
     }
 
+    /// Returns true only for exact path matches against prevalidated directories:
+    /// workspace(), synthetic_home(), cache_root() (when mount_cache_root()),
+    /// TmpBacking::BindHost host_dir, and entries in read_only_mounts() and read_write_mounts().
+    #[inline]
+    pub(crate) fn is_prevalidated_workdir(&self, workdir: &Path) -> bool {
+        workdir == self.workspace()
+            || workdir == self.synthetic_home()
+            || (self.mount_cache_root() && workdir == self.cache_root())
+            || matches!(self.tmp_backing(), TmpBacking::BindHost(host_dir) if workdir == host_dir.as_ref())
+            || self
+                .read_only_mounts()
+                .iter()
+                .any(|mount| workdir == mount.as_ref())
+            || self
+                .read_write_mounts()
+                .iter()
+                .any(|mount| workdir == mount.as_ref())
+    }
+
     fn sandbox_layout(&self) -> SandboxLayout<'_> {
         SandboxLayout {
             workspace: self.workspace(),
@@ -541,4 +560,50 @@ mod tests {
             Cow::Owned(mapped) => assert_eq!(mapped, Path::new("/workspace/subdir")),
         }
     }
+
+    #[test]
+    fn is_prevalidated_workdir_accepts_exact_profile_owned_roots() {
+        let profile = Profile {
+            preset: None,
+            workspace: Box::from(Path::new("/host/workspace")),
+            workspace_dest: Box::from(Path::new("/workspace")),
+            synthetic_home: Box::from(Path::new("/host/home")),
+            synthetic_home_dest: Box::from(Path::new("/home/sandbox")),
+            cache_root: Box::from(Path::new("/host/cache")),
+            tmp_backing: TmpBacking::BindHost(Box::from(Path::new("/host/tmp"))),
+            mount_cache_root: true,
+            compat_symlinks: Arc::new([]),
+            read_only_mounts: Arc::from([Box::from(Path::new("/host/ro"))]),
+            read_write_mounts: Arc::from([Box::from(Path::new("/host/rw"))]),
+            tmpfs_overlays: Arc::new([]),
+            file_overlays: Arc::new([]),
+            credential_file_mounts: Arc::new([]),
+            read_only_host_rootfs: false,
+            network_policy: NetworkPolicy::Disabled,
+            clear_env: false,
+            default_env: Arc::new([]),
+            extra_env: Arc::new([]),
+            availability: Availability::Unknown,
+            bwrap_program: Arc::from(Box::from(Path::new("/usr/bin/bwrap"))),
+            shell: Box::from(Path::new("/bin/sh")),
+            static_args: Arc::<[OsString]>::from([]),
+        };
+
+        assert!(profile.is_prevalidated_workdir(Path::new("/host/workspace")));
+        assert!(profile.is_prevalidated_workdir(Path::new("/host/home")));
+        assert!(profile.is_prevalidated_workdir(Path::new("/host/cache")));
+        assert!(profile.is_prevalidated_workdir(Path::new("/host/tmp")));
+        assert!(profile.is_prevalidated_workdir(Path::new("/host/ro")));
+        assert!(profile.is_prevalidated_workdir(Path::new("/host/rw")));
+    }
+
+    #[test]
+    fn is_prevalidated_workdir_rejects_nested_or_unmounted_paths() {
+        let profile = profile_with_workspace_dest("/workspace");
+
+        assert!(!profile.is_prevalidated_workdir(Path::new("/host/workspace/subdir")));
+        assert!(!profile.is_prevalidated_workdir(Path::new("/host/home/subdir")));
+        assert!(!profile.is_prevalidated_workdir(Path::new("/host/cache/subdir")));
+        assert!(!profile.is_prevalidated_workdir(Path::new("/outside")));
+    }
 }

src/llm-coding-tools-bubblewrap/src/wrap/command.rs

@@ -62,7 +62,11 @@ fn resolve_sandbox_cwd<'a>(
     profile: &'a Profile,
     workdir: Option<&'a Path>,
 ) -> Result<Cow<'a, Path>, LinuxBwrapError> {
-    validate_workdir(workdir)?;
+    if let Some(dir) = workdir {
+        if !profile.is_prevalidated_workdir(dir) {
+            validate_workdir(Some(dir))?;
+        }
+    }
     profile.map_workdir_to_sandbox(workdir)
 }
 

src/llm-coding-tools-core/src/tools/bash/blocking_impl.rs

@@ -1,8 +1,9 @@
 //! Blocking shell command execution.
 
 use super::{
-    timeout_error_with_kill_failure, timeout_message_with_buffered_output, validate_workdir,
-    BashExecutionMode, BashOutput, PIPE_BUFFER_CAPACITY,
+    string_from_utf8_or_lossy, timeout_error_with_kill_failure,
+    timeout_message_with_buffered_output, validate_workdir, BashExecutionMode, BashOutput,
+    PIPE_BUFFER_CAPACITY,
 };
 use crate::error::{ToolError, ToolResult};
 #[cfg(all(feature = "linux-bubblewrap", target_os = "linux"))]
@@ -14,6 +15,11 @@ use std::process::Stdio;
 use std::thread;
 use std::time::{Duration, Instant};
 
+/// Initial sleep between non-blocking wait polls.
+const INITIAL_POLL_INTERVAL_MS: u64 = 1;
+/// Upper bound for wait poll backoff.
+const MAX_POLL_INTERVAL_MS: u64 = 10;
+
 enum WaitOutcome {
     Exited(std::process::ExitStatus),
     TimedOut { kill_error: Option<std::io::Error> },
@@ -96,19 +102,24 @@ pub(in crate::tools::bash) fn run_wrapped_command(
     });
 
     let start = Instant::now();
+    let mut poll_interval_ms = INITIAL_POLL_INTERVAL_MS;
 
     // Poll for completion with timeout.
     let wait_outcome = loop {
         match child.try_wait() {
             Ok(Some(status)) => break WaitOutcome::Exited(status),
             Ok(None) => {
-                if start.elapsed() >= timeout {
+                let elapsed = start.elapsed();
+                if elapsed >= timeout {
                     // Kill entire process tree via Job Object (Windows) or process group (Unix)
                     break WaitOutcome::TimedOut {
                         kill_error: child.kill().err(),
                     };
                 }
-                thread::sleep(Duration::from_millis(10));
+
+                let remaining = timeout.saturating_sub(elapsed);
+                thread::sleep(remaining.min(Duration::from_millis(poll_interval_ms)));
+                poll_interval_ms = (poll_interval_ms << 1).min(MAX_POLL_INTERVAL_MS);
             }
             Err(e) => break WaitOutcome::WaitError(e),
         }
@@ -126,8 +137,8 @@ pub(in crate::tools::bash) fn run_wrapped_command(
     match wait_outcome {
         WaitOutcome::Exited(status) => Ok(BashOutput {
             exit_code: status.code(),
-            stdout: String::from_utf8_lossy(&stdout_data).into_owned(),
-            stderr: String::from_utf8_lossy(&stderr_data).into_owned(),
+            stdout: string_from_utf8_or_lossy(stdout_data),
+            stderr: string_from_utf8_or_lossy(stderr_data),
         }),
         WaitOutcome::TimedOut { kill_error } => Err(timeout_error_with_kill_failure(
             timeout_message_with_buffered_output(timeout, &stdout_data, &stderr_data),

src/llm-coding-tools-core/src/tools/bash/mod.rs

@@ -35,6 +35,7 @@ use crate::error::{ToolError, ToolResult};
 use crate::ToolOutput;
 use core::fmt::Write;
 use serde::Serialize;
+use std::borrow::Cow;
 use std::path::Path;
 use std::time::Duration;
 #[cfg(feature = "tokio")]
@@ -65,6 +66,17 @@ pub enum BashExecutionMode {
 /// 32KB covers typical command output without reallocations.
 const PIPE_BUFFER_CAPACITY: usize = 32 * 1024;
 
+#[inline]
+fn string_from_utf8_or_lossy(bytes: Vec<u8>) -> String {
+    match String::from_utf8(bytes) {
+        Ok(text) => text,
+        Err(error) => match String::from_utf8_lossy(&error.into_bytes()) {
+            Cow::Borrowed(text) => text.to_owned(),
+            Cow::Owned(text) => text,
+        },
+    }
+}
+
 #[inline]
 fn timeout_message_with_buffered_output(
     timeout: Duration,

src/llm-coding-tools-core/src/tools/bash/tokio_impl.rs

@@ -1,8 +1,9 @@
 //! Tokio-based async shell command execution.
 
 use super::{
-    timeout_error_with_kill_failure, timeout_message_with_buffered_output, validate_workdir,
-    BashExecutionMode, BashOutput, PIPE_BUFFER_CAPACITY,
+    string_from_utf8_or_lossy, timeout_error_with_kill_failure,
+    timeout_message_with_buffered_output, validate_workdir, BashExecutionMode, BashOutput,
+    PIPE_BUFFER_CAPACITY,
 };
 use crate::error::{ToolError, ToolResult};
 #[cfg(all(feature = "linux-bubblewrap", target_os = "linux"))]
@@ -55,7 +56,7 @@ where
 fn take_pipe_buffer(buffer: SharedPipeBuffer) -> Vec<u8> {
     match Arc::try_unwrap(buffer) {
         Ok(mutex) => mutex.into_inner(),
-        Err(shared) => shared.lock().clone(),
+        Err(shared) => core::mem::take(&mut *shared.lock()),
     }
 }
 
@@ -165,8 +166,8 @@ pub(in crate::tools::bash) async fn run_wrapped_command(
 
             Ok(BashOutput {
                 exit_code: status.code(),
-                stdout: String::from_utf8_lossy(&stdout_data).into_owned(),
-                stderr: String::from_utf8_lossy(&stderr_data).into_owned(),
+                stdout: string_from_utf8_or_lossy(stdout_data),
+                stderr: string_from_utf8_or_lossy(stderr_data),
             })
         }
         None => {