diff --git a/src/domain/permissions/checker.rs b/src/domain/permissions/checker.rs index 7d3ce1d..fd2bfca 100644 --- a/src/domain/permissions/checker.rs +++ b/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>, - permissions: Mutex>, + 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) -> Self { - Self { - created: Mutex::new(Vec::new()), - permissions: Mutex::new(permissions), - } + fn store(&self) -> Arc { + Arc::new(SqlitePermissionStore::new(self.pool.clone())) } } - impl PermissionStore for TestStore { - fn create_permission(&self, permission: Permission) -> Result { - 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, 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) { + 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(), diff --git a/src/domain/tools/patch_files.rs b/src/domain/tools/patch_files.rs index 6d85664..34541e3 100644 --- a/src/domain/tools/patch_files.rs +++ b/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 { + 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, } } diff --git a/src/domain/tools/shell_exec.rs b/src/domain/tools/shell_exec.rs index a1ec056..3f51e6b 100644 --- a/src/domain/tools/shell_exec.rs +++ b/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>, @@ -18,6 +24,8 @@ pub struct ShellExecInput { pub command: String, #[serde(rename = "working_dir", default)] pub working_dir: Option, + #[serde(rename = "timeout_ms", default)] + pub timeout_ms: Option, } #[derive(Debug, Clone)] @@ -25,9 +33,13 @@ struct ShellExecInputParsed { raw: String, command: String, working_dir: Option, + timeout_ms: Option, 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 { + 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(mut reader: R, budget: usize) -> Vec { + 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()); } -} \ No newline at end of file +}