save

src/domain/permissions/checker.rs

@@ -180,6 +180,15 @@ impl PermissionChecker {
             return Ok(p.system_decision());
         }
 
+        if request.tool_name == "patch_files" {
+            if let Ok(Some(p)) = self
+                .store
+                .find_permission(&request.tool_name, project_id, "", "")
+            {
+                return Ok(p.system_decision());
+            }
+        }
+
         Ok(self.config.default_decision.clone())
     }
 
@@ -199,111 +208,59 @@ impl PermissionChecker {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::domain::permissions::store::SqlitePermissionStore;
     use crate::domain::permissions::types::Permission;
     use crate::domain::permissions::{PermissionRequest, PermissionScope};
     use crate::domain::session::{Request, SessionRequest};
     use crate::domain::tools::{Error, Tool, ToolResult};
+    use crate::infrastructure::db::DbPool;
+    use r2d2_sqlite::SqliteConnectionManager;
     use std::path::{Path, PathBuf};
     use std::sync::atomic::{AtomicUsize, Ordering};
-    use std::sync::Mutex;
+    use tempfile::TempDir;
 
-    struct TestStore {
-        created: Mutex<Vec<Permission>>,
-        permissions: Mutex<Vec<Permission>>,
+    struct TestDb {
+        _temp_dir: TempDir,
+        pool: DbPool,
     }
 
-    impl TestStore {
+    impl TestDb {
         fn new() -> Self {
+            let temp_dir = tempfile::tempdir().unwrap();
+            let db_path = temp_dir.path().join("permissions.db");
+            let manager = SqliteConnectionManager::file(&db_path);
+            let pool = r2d2::Pool::new(manager).unwrap();
+            let conn = pool.get().unwrap();
+
+            conn.execute(
+                "CREATE TABLE IF NOT EXISTS permissions (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    tool_name TEXT NOT NULL,
+                    command_pattern TEXT,
+                    resource_pattern TEXT,
+                    decision TEXT NOT NULL,
+                    scope TEXT NOT NULL,
+                    project_id INTEGER,
+                    created_at TEXT NOT NULL
+                )",
+                [],
+            )
+            .unwrap();
+
             Self {
-                created: Mutex::new(Vec::new()),
-                permissions: Mutex::new(Vec::new()),
+                _temp_dir: temp_dir,
+                pool,
             }
         }
 
-        fn with_permissions(permissions: Vec<Permission>) -> Self {
-            Self {
-                created: Mutex::new(Vec::new()),
-                permissions: Mutex::new(permissions),
-            }
+        fn store(&self) -> Arc<SqlitePermissionStore> {
+            Arc::new(SqlitePermissionStore::new(self.pool.clone()))
         }
     }
 
-    impl PermissionStore for TestStore {
-        fn create_permission(&self, permission: Permission) -> Result<Permission, StoreError> {
-            self.created.lock().unwrap().push(permission.clone());
-            Ok(permission)
-        }
-
-        fn find_permission(
-            &self,
-            tool: &str,
-            project_id: i32,
-            command_pattern: &str,
-            resource_pattern: &str,
-        ) -> Result<Option<Permission>, StoreError> {
-            let permissions = self.permissions.lock().unwrap();
-
-            // Find the most specific matching permission
-            // Priority: command+path > command > path > tool
-            let mut best_match: Option<(&Permission, i32)> = None;
-
-            for permission in permissions.iter() {
-                // Check if project_id matches
-                if let Some(perm_project_id) = permission.project_id {
-                    if perm_project_id != project_id {
-                        continue;
-                    }
-                }
-
-                // Check if tool matches
-                if permission.tool_name != tool {
-                    continue;
-                }
-
-                // Calculate specificity and check if it matches
-                let mut specificity = 0;
-                let mut matches = true;
-
-                // Check command pattern
-                match &permission.command_pattern {
-                    Some(pattern) => {
-                        if !command_pattern.is_empty() && pattern == command_pattern {
-                            specificity += 2;
-                        } else if !command_pattern.is_empty() {
-                            matches = false;
-                        }
-                    }
-                    None => {
-                        // None means it matches any command
-                    }
-                }
-
-                // Check resource pattern
-                match &permission.resource_pattern {
-                    Some(pattern) => {
-                        if !resource_pattern.is_empty() && pattern == resource_pattern {
-                            specificity += 1;
-                        } else if !resource_pattern.is_empty() {
-                            matches = false;
-                        }
-                    }
-                    None => {
-                        // None means it matches any resource
-                    }
-                }
-
-                if matches {
-                    if let Some((_, best_specificity)) = best_match {
-                        if specificity > best_specificity {
-                            best_match = Some((permission, specificity));
-                        }
-                    } else {
-                        best_match = Some((permission, specificity));
-                    }
-                }
-            }
-
-            Ok(best_match.map(|(p, _)| p.clone()))
+    fn seed_permissions(store: &SqlitePermissionStore, permissions: Vec<Permission>) {
+        for permission in permissions {
+            store.create_permission(permission).unwrap();
         }
     }
 
@@ -503,7 +460,8 @@ mod tests {
         let file_path = root.join("sample.txt");
         std::fs::write(&file_path, "data").unwrap();
 
-        let store = Arc::new(TestStore::new());
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let calls = Arc::new(AtomicUsize::new(0));
         let prompter = Arc::new(TestPrompter {
             calls: Arc::clone(&calls),
@@ -536,7 +494,8 @@ mod tests {
         let external_file = external_dir.path().join("external.txt");
         std::fs::write(&external_file, "data").unwrap();
 
-        let store = Arc::new(TestStore::new());
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let calls = Arc::new(AtomicUsize::new(0));
         let prompter = Arc::new(TestPrompter {
             calls: Arc::clone(&calls),
@@ -568,7 +527,8 @@ mod tests {
         let file_path = root.join("sample.txt");
         std::fs::write(&file_path, "data").unwrap();
 
-        let store = Arc::new(TestStore::new());
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let calls = Arc::new(AtomicUsize::new(0));
         let prompter = Arc::new(TestPrompter {
             calls: Arc::clone(&calls),
@@ -598,7 +558,8 @@ mod tests {
         let temp = tempfile::tempdir().unwrap();
         let root = temp.path().to_path_buf();
 
-        let store = Arc::new(TestStore::new());
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let calls = Arc::new(AtomicUsize::new(0));
         let prompter = Arc::new(TestPrompter {
             calls: Arc::clone(&calls),
@@ -630,7 +591,8 @@ mod tests {
         let root = temp.path().to_path_buf();
         let restricted_path = PathBuf::from("/etc");
 
-        let store = Arc::new(TestStore::new());
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let calls = Arc::new(AtomicUsize::new(0));
         let prompter = Arc::new(TestPrompter {
             calls: Arc::clone(&calls),
@@ -660,7 +622,8 @@ mod tests {
     fn resolve_permission_returns_ask_when_no_stored_permission() {
         let abs = std::fs::canonicalize(".").expect("failed to canonicalize path");
 
-        let store = Arc::new(TestStore::with_permissions(vec![]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -708,10 +671,12 @@ mod tests {
             PermissionScope::Project,
             Some(1),
         );
-        let store = Arc::new(TestStore::with_permissions(vec![
-            command_permission,
-            path_permission,
-        ]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(
+            store.as_ref(),
+            vec![command_permission, path_permission],
+        );
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -750,7 +715,9 @@ mod tests {
             PermissionScope::Project,
             Some(1),
         );
-        let store = Arc::new(TestStore::with_permissions(vec![path_permission]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![path_permission]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -781,11 +748,13 @@ mod tests {
             "command_tool".to_string(),
             None,
             None,
-            UserPermissionDecision::AllowOnce,
+            UserPermissionDecision::AlwaysAllow,
             PermissionScope::Project,
             Some(1),
         );
-        let store = Arc::new(TestStore::with_permissions(vec![tool_permission]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![tool_permission]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -806,13 +775,14 @@ mod tests {
 
         let decision = checker.resolve_permission(&request).unwrap();
 
-        // AllowOnce → system_decision = Ask
-        assert_eq!(decision, SystemPermissionDecision::Ask);
+        // AlwaysAllow → system_decision = Allow
+        assert_eq!(decision, SystemPermissionDecision::Allow);
     }
 
     #[test]
     fn resolve_permission_falls_back_to_default_rules() {
-        let store = Arc::new(TestStore::new());
+        let test_db = TestDb::new();
+        let store = test_db.store();
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig {
@@ -857,7 +827,9 @@ mod tests {
             Some(1),
         );
 
-        let store = Arc::new(TestStore::with_permissions(vec![perm]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![perm]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -899,7 +871,9 @@ mod tests {
             Some(1),
         );
 
-        let store = Arc::new(TestStore::with_permissions(vec![perm]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![perm]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -930,17 +904,19 @@ mod tests {
         let project_dir = tempfile::tempdir().unwrap();
         let project_root = project_dir.path().to_path_buf();
 
-        // Create an AlwaysAllow permission for write_tool (no resource pattern = matches any)
+        let target_file = project_root.join("internal.txt");
         let perm = Permission::new(
             "write_tool".to_string(),
             None,
-            None,
+            Some(target_file.to_string_lossy().to_string()),
             UserPermissionDecision::AlwaysAllow,
             PermissionScope::Project,
             Some(1),
         );
 
-        let store = Arc::new(TestStore::with_permissions(vec![perm]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![perm]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -953,7 +929,7 @@ mod tests {
         let request = PermissionRequest::new(
             "write_tool".to_string(),
             None,
-            vec![project_root.join("internal.txt")],
+            vec![target_file],
             PermissionScope::Project,
             Some(1),
             false,
@@ -972,17 +948,19 @@ mod tests {
         let internal_file = project_root.join("internal.txt");
         std::fs::write(&internal_file, "data").unwrap();
 
-        // Create an AlwaysAllow permission for read_only (no resource pattern = matches any)
+        // Create an AlwaysAllow permission for read_only matching a specific resource.
         let perm = Permission::new(
             "read_only".to_string(),
             None,
-            None,
+            Some(internal_file.to_string_lossy().to_string()),
             UserPermissionDecision::AlwaysAllow,
             PermissionScope::Project,
             Some(1),
         );
 
-        let store = Arc::new(TestStore::with_permissions(vec![perm]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![perm]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -1022,7 +1000,9 @@ mod tests {
             Some(1),
         );
 
-        let store = Arc::new(TestStore::with_permissions(vec![command_perm]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![command_perm]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),
@@ -1077,7 +1057,9 @@ mod tests {
             Some(1),
         );
 
-        let store = Arc::new(TestStore::with_permissions(vec![existing_perm]));
+        let test_db = TestDb::new();
+        let store = test_db.store();
+        seed_permissions(store.as_ref(), vec![existing_perm]);
         let checker = PermissionChecker::new_with_prompter(
             store,
             PermissionConfig::default(),

src/domain/tools/patch_files.rs

@@ -200,38 +200,11 @@ impl Tool for PatchFiles {
         // Compute file changes
         let mut file_changes = Vec::new();
         for path in &touched_paths {
-            let original = original_contents.get(path);
-            let new_content = new_vfs.get(path);
+            let original = original_contents.get(path).map(|s| s.as_str());
+            let new_content = new_vfs.get(path).map(|s| s.as_str());
 
-            let (added_lines, deleted_lines) =
-                if let (Some(orig), Some(new)) = (original, new_content) {
-                    if orig != new {
-                        compute_line_diff(orig, new)
-                    } else {
-                        (0, 0) // No change
-                    }
-                } else if original.is_some() && new_content.is_none() {
-                    // Deleted: count original lines as deleted
-                    (0, original.unwrap().lines().count() as u32)
-                } else if original.is_none() && new_content.is_some() {
-                    // Added: count new lines as added
-                    (new_content.unwrap().lines().count() as u32, 0)
-                } else {
-                    (0, 0) // No content change
-                };
-
-            if added_lines > 0 || deleted_lines > 0 {
-                let unified_diff = generate_unified_diff(
-                    path,
-                    original.map(|s| s.as_str()),
-                    new_content.map(|s| s.as_str()),
-                );
-                file_changes.push(FileChange {
-                    path: path.clone(),
-                    added_lines,
-                    deleted_lines,
-                    unified_diff,
-                });
+            if let Some(change) = compute_file_change(path, original, new_content) {
+                file_changes.push(change);
             }
         }
 
@@ -345,31 +318,36 @@ impl Tool for PatchFiles {
     }
 }
 
-fn compute_line_diff(old: &str, new: &str) -> (u32, u32) {
-    let diff = TextDiff::from_lines(old, new);
-    let mut added = 0;
-    let mut deleted = 0;
-    for change in diff.iter_all_changes() {
-        match change.tag() {
-            similar::ChangeTag::Insert => added += 1,
-            similar::ChangeTag::Delete => deleted += 1,
-            similar::ChangeTag::Equal => {}
-        }
-    }
-    (added, deleted)
-}
+const DIFF_LINE_LIMIT: usize = 2000;
+
+fn compute_file_change(path: &str, old: Option<&str>, new: Option<&str>) -> Option<FileChange> {
+    let old_lines = old.map(|content| content.lines().count()).unwrap_or(0);
+    let new_lines = new.map(|content| content.lines().count()).unwrap_or(0);
+    let max_lines = old_lines.max(new_lines);
+    let large_file = max_lines > DIFF_LINE_LIMIT;
 
-fn generate_unified_diff(path: &str, old: Option<&str>, new: Option<&str>) -> String {
     match (old, new) {
         (Some(old_content), Some(new_content)) => {
-            // Modified file
+            if old_content == new_content {
+                return None;
+            }
+            if large_file {
+                return Some(FileChange {
+                    path: path.to_string(),
+                    added_lines: new_lines.saturating_sub(old_lines) as u32,
+                    deleted_lines: old_lines.saturating_sub(new_lines) as u32,
+                    unified_diff: format!("Diff omitted for large file ({} lines).", max_lines),
+                });
+            }
+
             let diff = TextDiff::from_lines(old_content, new_content);
+            let mut added = 0;
+            let mut deleted = 0;
             let mut result = String::new();
 
             result.push_str(&format!("--- a/{}\n", path));
             result.push_str(&format!("+++ b/{}\n", path));
 
-            // Generate hunks
             let mut old_line = 1;
             let mut new_line = 1;
             let mut hunk_changes = Vec::new();
@@ -391,10 +369,12 @@ fn generate_unified_diff(path: &str, old: Option<&str>, new: Option<&str>) -> St
                     similar::ChangeTag::Delete => {
                         hunk_changes.push(format!("-{}", change));
                         old_line += 1;
+                        deleted += 1;
                     }
                     similar::ChangeTag::Insert => {
                         hunk_changes.push(format!("+{}", change));
                         new_line += 1;
+                        added += 1;
                     }
                 }
             }
@@ -411,42 +391,58 @@ fn generate_unified_diff(path: &str, old: Option<&str>, new: Option<&str>) -> St
                 }
             }
 
-            result
+            Some(FileChange {
+                path: path.to_string(),
+                added_lines: added,
+                deleted_lines: deleted,
+                unified_diff: result,
+            })
         }
         (None, Some(new_content)) => {
-            // New file
-            let mut result = String::new();
-            result.push_str(&format!("--- /dev/null\n"));
-            result.push_str(&format!("+++ b/{}\n", path));
+            let added_lines = new_lines as u32;
+            let unified_diff = if large_file {
+                format!("Diff omitted for large file ({} lines).", max_lines)
+            } else {
+                let mut result = String::new();
+                result.push_str("--- /dev/null\n");
+                result.push_str(&format!("+++ b/{}\n", path));
+                result.push_str(&format!("@@ -0,0 +1,{} @@\n", new_lines));
+                for line in new_content.lines() {
+                    result.push_str(&format!("+{}\n", line));
+                }
+                result
+            };
 
-            let line_count = new_content.lines().count();
-            result.push_str(&format!("@@ -0,0 +1,{} @@\n", line_count));
-
-            for line in new_content.lines() {
-                result.push_str(&format!("+{}\n", line));
-            }
-
-            result
+            Some(FileChange {
+                path: path.to_string(),
+                added_lines,
+                deleted_lines: 0,
+                unified_diff,
+            })
         }
         (Some(old_content), None) => {
-            // Deleted file
-            let mut result = String::new();
-            result.push_str(&format!("--- a/{}\n", path));
-            result.push_str(&format!("+++ /dev/null\n"));
+            let deleted_lines = old_lines as u32;
+            let unified_diff = if large_file {
+                format!("Diff omitted for large file ({} lines).", max_lines)
+            } else {
+                let mut result = String::new();
+                result.push_str(&format!("--- a/{}\n", path));
+                result.push_str("+++ /dev/null\n");
+                result.push_str(&format!("@@ -1,{} +0,0 @@\n", old_lines));
+                for line in old_content.lines() {
+                    result.push_str(&format!("-{}\n", line));
+                }
+                result
+            };
 
-            let line_count = old_content.lines().count();
-            result.push_str(&format!("@@ -1,{} +0,0 @@\n", line_count));
-
-            for line in old_content.lines() {
-                result.push_str(&format!("-{}\n", line));
-            }
-
-            result
-        }
-        (None, None) => {
-            // No change (shouldn't happen)
-            String::new()
+            Some(FileChange {
+                path: path.to_string(),
+                added_lines: 0,
+                deleted_lines,
+                unified_diff,
+            })
         }
+        (None, None) => None,
     }
 }
 

src/domain/tools/shell_exec.rs

@@ -2,9 +2,15 @@ use crate::domain::session::Request;
 use crate::domain::tools::{short_words, Error, Tool, ToolResult, TOOL_OUTPUT_BUDGET_CHARS};
 use serde::Deserialize;
 use serde_json::json;
+use std::io::Read;
 use std::path::PathBuf;
 use std::process::Command;
+use std::process::Output;
+use std::process::Stdio;
+use std::sync::mpsc;
 use std::sync::Mutex;
+use std::thread;
+use std::time::{Duration, Instant};
 
 pub struct ShellExec {
     input: Mutex<Option<ShellExecInputParsed>>,
@@ -18,6 +24,8 @@ pub struct ShellExecInput {
     pub command: String,
     #[serde(rename = "working_dir", default)]
     pub working_dir: Option<String>,
+    #[serde(rename = "timeout_ms", default)]
+    pub timeout_ms: Option<u64>,
 }
 
 #[derive(Debug, Clone)]
@@ -25,9 +33,13 @@ struct ShellExecInputParsed {
     raw: String,
     command: String,
     working_dir: Option<String>,
+    timeout_ms: Option<u64>,
     call_id: String,
 }
 
+const DEFAULT_TIMEOUT_MS: u64 = 15 * 60 * 1000;
+const TIMEOUT_POLL_MS: u64 = 50;
+
 fn is_allowed_read_only_command(command: &str) -> bool {
     let trimmed = command.trim();
     if trimmed.is_empty() {
@@ -50,9 +62,11 @@ fn is_allowed_read_only_command(command: &str) -> bool {
 
     match first {
         "rg" | "grep" | "glob" | "cat" | "head" | "tail" | "less" | "more" | "wc" | "cut"
-        | "sort" | "uniq" | "find" | "ls" | "tree" | "stat" | "file" | "awk" | "pwd"
-        | "which" | "type" => true,
-        "sed" => args.iter().any(|arg| *arg == "-n" || *arg == "--quiet" || *arg == "--silent"),
+        | "sort" | "uniq" | "find" | "ls" | "tree" | "stat" | "file" | "awk" | "pwd" | "which"
+        | "type" => true,
+        "sed" => args
+            .iter()
+            .any(|arg| *arg == "-n" || *arg == "--quiet" || *arg == "--silent"),
         _ => false,
     }
 }
@@ -76,6 +90,7 @@ impl Tool for ShellExec {
                     raw: trimmed.to_string(),
                     command: parsed.command,
                     working_dir: parsed.working_dir,
+                    timeout_ms: parsed.timeout_ms,
                     call_id,
                 });
                 None
@@ -122,19 +137,22 @@ impl Tool for ShellExec {
             None => request.project_root().to_path_buf(),
         };
 
+        let timeout_ms = input.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS);
+        let timeout_ms = if timeout_ms == 0 {
+            DEFAULT_TIMEOUT_MS
+        } else {
+            timeout_ms
+        };
+        let timeout = Duration::from_millis(timeout_ms);
+
         // Execute the command
-        let output = match Command::new("bash")
-            .arg("-c")
-            .arg(&input.command)
-            .current_dir(&work_dir)
-            .output()
-        {
+        let output = match run_command_with_timeout(&input.command, &work_dir, timeout) {
             Ok(o) => o,
-            Err(e) => {
+            Err(err) => {
                 return ToolResult::error(
                     self.name().to_string(),
                     input.raw.clone(),
-                    format!("Failed to execute command: {}", e),
+                    err,
                     input.call_id.clone(),
                 )
             }
@@ -186,6 +204,10 @@ impl Tool for ShellExec {
                 "working_dir": {
                     "type": "string",
                     "description": "optional; directory to run command in (default: project root)"
+                },
+                "timeout_ms": {
+                    "type": "number",
+                    "description": "optional; command timeout in milliseconds (default: 900000)"
                 }
             },
             "required": ["command"],
@@ -304,6 +326,108 @@ impl ShellExec {
     }
 }
 
+fn run_command_with_timeout(
+    command: &str,
+    work_dir: &std::path::Path,
+    timeout: Duration,
+) -> Result<Output, String> {
+    let mut child = Command::new("bash")
+        .arg("-c")
+        .arg(command)
+        .current_dir(work_dir)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()
+        .map_err(|e| format!("Failed to execute command: {}", e))?;
+
+    let stdout = child
+        .stdout
+        .take()
+        .ok_or_else(|| "Failed to capture stdout".to_string())?;
+    let stderr = child
+        .stderr
+        .take()
+        .ok_or_else(|| "Failed to capture stderr".to_string())?;
+
+    let (stdout_tx, stdout_rx) = mpsc::channel();
+    let (stderr_tx, stderr_rx) = mpsc::channel();
+
+    thread::spawn(move || {
+        let buf = read_stream_with_budget(stdout, TOOL_OUTPUT_BUDGET_CHARS);
+        let _ = stdout_tx.send(buf);
+    });
+
+    thread::spawn(move || {
+        let buf = read_stream_with_budget(stderr, TOOL_OUTPUT_BUDGET_CHARS);
+        let _ = stderr_tx.send(buf);
+    });
+
+    let start = Instant::now();
+    let mut timed_out = false;
+    let status = loop {
+        match child.try_wait() {
+            Ok(Some(status)) => break status,
+            Ok(None) => {
+                if start.elapsed() >= timeout {
+                    timed_out = true;
+                    let _ = child.kill();
+                    break child
+                        .wait()
+                        .map_err(|e| format!("Failed to stop command: {}", e))?;
+                }
+                thread::sleep(Duration::from_millis(TIMEOUT_POLL_MS));
+            }
+            Err(e) => return Err(format!("Failed while executing command: {}", e)),
+        }
+    };
+
+    let stdout_bytes = stdout_rx
+        .recv_timeout(Duration::from_secs(1))
+        .unwrap_or_default();
+    let stderr_bytes = stderr_rx
+        .recv_timeout(Duration::from_secs(1))
+        .unwrap_or_default();
+    let output = Output {
+        status,
+        stdout: stdout_bytes,
+        stderr: stderr_bytes,
+    };
+
+    if timed_out {
+        let stdout_preview = String::from_utf8_lossy(&output.stdout).to_string();
+        let stderr_preview = String::from_utf8_lossy(&output.stderr).to_string();
+        let mut message = format!("Command timed out after {}ms", timeout.as_millis());
+        if !stdout_preview.is_empty() {
+            message.push_str(&format!("\n[stdout]: {}", stdout_preview));
+        }
+        if !stderr_preview.is_empty() {
+            message.push_str(&format!("\n[stderr]: {}", stderr_preview));
+        }
+        return Err(message);
+    }
+
+    Ok(output)
+}
+
+fn read_stream_with_budget<R: Read>(mut reader: R, budget: usize) -> Vec<u8> {
+    let mut buf = Vec::new();
+    let mut temp = [0u8; 8192];
+    loop {
+        match reader.read(&mut temp) {
+            Ok(0) => break,
+            Ok(n) => {
+                let remaining = budget.saturating_sub(buf.len());
+                if remaining > 0 {
+                    let copy_len = remaining.min(n);
+                    buf.extend_from_slice(&temp[..copy_len]);
+                }
+            }
+            Err(_) => break,
+        }
+    }
+    buf
+}
+
 impl Default for ShellExec {
     fn default() -> Self {
         Self::new()
@@ -634,4 +758,4 @@ mod tests {
             .is_none());
         assert!(!tool.is_read_only());
     }
-}
+}