diff --git a/crates/agent/src/pattern_extraction.rs b/crates/agent/src/pattern_extraction.rs index 0e82bd3bafa909c710fb917ac481eb7fd7abd3de..69a7abae32d6df9c2755e53292ab1c1a1b5341de 100644 --- a/crates/agent/src/pattern_extraction.rs +++ b/crates/agent/src/pattern_extraction.rs @@ -7,28 +7,53 @@ fn normalize_separators(path_str: &str) -> String { path_str.replace('\\', "/") } -/// Extracts the command name from a shell command using the shell parser. +/// Returns true if the token looks like a command name or subcommand — i.e. it +/// contains only alphanumeric characters, hyphens, and underscores, and does not +/// start with a hyphen (which would make it a flag). +fn is_plain_command_token(token: &str) -> bool { + !token.starts_with('-') + && token + .chars() + .all(|c| c.is_alphanumeric() || c == '-' || c == '_') +} + +struct CommandPrefix { + command: String, + subcommand: Option, +} + +/// Extracts the command name and optional subcommand from a shell command using +/// the shell parser. /// -/// This parses the command properly to extract just the command name (first word), -/// handling shell syntax correctly. Returns `None` if parsing fails or if the -/// command name contains path separators (for security reasons). -fn extract_command_name(command: &str) -> Option { +/// This parses the command properly to extract the command name and optional +/// subcommand (e.g. "cargo" and "test" from "cargo test -p search"), handling shell +/// syntax correctly. Returns `None` if parsing fails or if the command name +/// contains path separators (for security reasons). +fn extract_command_prefix(command: &str) -> Option { let commands = extract_commands(command)?; let first_command = commands.first()?; - let first_token = first_command.split_whitespace().next()?; + let mut tokens = first_command.split_whitespace(); + let first_token = tokens.next()?; // Only allow alphanumeric commands with hyphens/underscores. // Reject paths like "./script.sh" or "/usr/bin/python" to prevent // users from accidentally allowing arbitrary script execution. - if first_token - .chars() - .all(|c| c.is_alphanumeric() || c == '-' || c == '_') - { - Some(first_token.to_string()) - } else { - None + if !is_plain_command_token(first_token) { + return None; } + + // Include the subcommand (second non-flag token) when present, to produce + // more specific patterns like "cargo test" instead of just "cargo". + let subcommand = tokens + .next() + .filter(|second_token| is_plain_command_token(second_token)) + .map(|second_token| second_token.to_string()); + + Some(CommandPrefix { + command: first_token.to_string(), + subcommand, + }) } /// Extracts a regex pattern from a terminal command based on the first token (command name). @@ -38,12 +63,26 @@ fn extract_command_name(command: &str) -> Option { /// rules for well-known command names (like `cargo`, `npm`, `git`), not for arbitrary /// scripts or absolute paths which could be manipulated by an attacker. pub fn extract_terminal_pattern(command: &str) -> Option { - let command_name = extract_command_name(command)?; - Some(format!("^{}\\b", regex::escape(&command_name))) + let prefix = extract_command_prefix(command)?; + let escaped_command = regex::escape(&prefix.command); + Some(match &prefix.subcommand { + Some(subcommand) => { + format!( + "^{}\\s+{}(\\s|$)", + escaped_command, + regex::escape(subcommand) + ) + } + None => format!("^{}\\b", escaped_command), + }) } pub fn extract_terminal_pattern_display(command: &str) -> Option { - extract_command_name(command) + let prefix = extract_command_prefix(command)?; + match prefix.subcommand { + Some(subcommand) => Some(format!("{} {}", prefix.command, subcommand)), + None => Some(prefix.command), + } } pub fn extract_path_pattern(path: &str) -> Option { @@ -125,20 +164,51 @@ mod tests { fn test_extract_terminal_pattern() { assert_eq!( extract_terminal_pattern("cargo build --release"), - Some("^cargo\\b".to_string()) + Some("^cargo\\s+build(\\s|$)".to_string()) + ); + assert_eq!( + extract_terminal_pattern("cargo test -p search"), + Some("^cargo\\s+test(\\s|$)".to_string()) ); assert_eq!( extract_terminal_pattern("npm install"), - Some("^npm\\b".to_string()) + Some("^npm\\s+install(\\s|$)".to_string()) ); assert_eq!( extract_terminal_pattern("git-lfs pull"), - Some("^git\\-lfs\\b".to_string()) + Some("^git\\-lfs\\s+pull(\\s|$)".to_string()) ); assert_eq!( extract_terminal_pattern("my_script arg"), - Some("^my_script\\b".to_string()) + Some("^my_script\\s+arg(\\s|$)".to_string()) + ); + + // Flags as second token: only the command name is used + assert_eq!( + extract_terminal_pattern("ls -la"), + Some("^ls\\b".to_string()) ); + assert_eq!( + extract_terminal_pattern("rm --force foo"), + Some("^rm\\b".to_string()) + ); + + // Single-word commands + assert_eq!(extract_terminal_pattern("ls"), Some("^ls\\b".to_string())); + + // Subcommand pattern does not match a hyphenated extension of the subcommand + // (e.g. approving "cargo build" should not approve "cargo build-foo") + assert_eq!( + extract_terminal_pattern("cargo build"), + Some("^cargo\\s+build(\\s|$)".to_string()) + ); + let pattern = regex::Regex::new(&extract_terminal_pattern("cargo build").unwrap()).unwrap(); + assert!(pattern.is_match("cargo build --release")); + assert!(pattern.is_match("cargo build")); + assert!(!pattern.is_match("cargo build-foo")); + assert!(!pattern.is_match("cargo builder")); + + // Path-like commands are rejected assert_eq!(extract_terminal_pattern("./script.sh arg"), None); assert_eq!(extract_terminal_pattern("/usr/bin/python arg"), None); } @@ -147,11 +217,23 @@ mod tests { fn test_extract_terminal_pattern_display() { assert_eq!( extract_terminal_pattern_display("cargo build --release"), - Some("cargo".to_string()) + Some("cargo build".to_string()) + ); + assert_eq!( + extract_terminal_pattern_display("cargo test -p search"), + Some("cargo test".to_string()) ); assert_eq!( extract_terminal_pattern_display("npm install"), - Some("npm".to_string()) + Some("npm install".to_string()) + ); + assert_eq!( + extract_terminal_pattern_display("ls -la"), + Some("ls".to_string()) + ); + assert_eq!( + extract_terminal_pattern_display("ls"), + Some("ls".to_string()) ); } diff --git a/crates/agent/src/tests/mod.rs b/crates/agent/src/tests/mod.rs index 232c626135b2128e62f15698d2387b7575202ca3..25bdcea4d806a173101d80822e5807dc9bcd239b 100644 --- a/crates/agent/src/tests/mod.rs +++ b/crates/agent/src/tests/mod.rs @@ -1053,7 +1053,47 @@ fn test_permission_options_terminal_with_pattern() { .map(|choice| choice.allow.name.as_ref()) .collect(); assert!(labels.contains(&"Always for terminal")); - assert!(labels.contains(&"Always for `cargo` commands")); + assert!(labels.contains(&"Always for `cargo build` commands")); + assert!(labels.contains(&"Only this time")); +} + +#[test] +fn test_permission_options_terminal_command_with_flag_second_token() { + let permission_options = + ToolPermissionContext::new(TerminalTool::NAME, vec!["ls -la".to_string()]) + .build_permission_options(); + + let PermissionOptions::Dropdown(choices) = permission_options else { + panic!("Expected dropdown permission options"); + }; + + assert_eq!(choices.len(), 3); + let labels: Vec<&str> = choices + .iter() + .map(|choice| choice.allow.name.as_ref()) + .collect(); + assert!(labels.contains(&"Always for terminal")); + assert!(labels.contains(&"Always for `ls` commands")); + assert!(labels.contains(&"Only this time")); +} + +#[test] +fn test_permission_options_terminal_single_word_command() { + let permission_options = + ToolPermissionContext::new(TerminalTool::NAME, vec!["whoami".to_string()]) + .build_permission_options(); + + let PermissionOptions::Dropdown(choices) = permission_options else { + panic!("Expected dropdown permission options"); + }; + + assert_eq!(choices.len(), 3); + let labels: Vec<&str> = choices + .iter() + .map(|choice| choice.allow.name.as_ref()) + .collect(); + assert!(labels.contains(&"Always for terminal")); + assert!(labels.contains(&"Always for `whoami` commands")); assert!(labels.contains(&"Only this time")); } diff --git a/crates/agent/src/tool_permissions.rs b/crates/agent/src/tool_permissions.rs index ef6e699d407eb9d7fbd53f9cdb8b8e46a2ed3b3e..7af196b295e8aa9fb3293c513818203c67aa2054 100644 --- a/crates/agent/src/tool_permissions.rs +++ b/crates/agent/src/tool_permissions.rs @@ -1174,24 +1174,38 @@ mod tests { #[test] fn always_allow_button_works_end_to_end() { // This test verifies that the "Always Allow" button behavior works correctly: - // 1. User runs a command like "cargo build" - // 2. They click "Always Allow for `cargo` commands" - // 3. The pattern extracted from that command should match future cargo commands + // 1. User runs a command like "cargo build --release" + // 2. They click "Always Allow for `cargo build` commands" + // 3. The pattern extracted should match future "cargo build" commands + // but NOT other cargo subcommands like "cargo test" let original_command = "cargo build --release"; let extracted_pattern = pattern(original_command); // The extracted pattern should allow the original command t(original_command).allow(&[extracted_pattern]).is_allow(); - // It should also allow other commands with the same base command - t("cargo test").allow(&[extracted_pattern]).is_allow(); - t("cargo fmt").allow(&[extracted_pattern]).is_allow(); + // It should allow other "cargo build" invocations with different flags + t("cargo build").allow(&[extracted_pattern]).is_allow(); + t("cargo build --features foo") + .allow(&[extracted_pattern]) + .is_allow(); + + // But NOT other cargo subcommands — the pattern is subcommand-specific + t("cargo test").allow(&[extracted_pattern]).is_confirm(); + t("cargo fmt").allow(&[extracted_pattern]).is_confirm(); + + // Hyphenated extensions of the subcommand should not match either + // (e.g. cargo plugins like "cargo build-foo") + t("cargo build-foo") + .allow(&[extracted_pattern]) + .is_confirm(); + t("cargo builder").allow(&[extracted_pattern]).is_confirm(); // But not commands with different base commands t("npm install").allow(&[extracted_pattern]).is_confirm(); - // And it should work with subcommand extraction (chained commands) - t("cargo build && cargo test") + // Chained commands: all must match the pattern + t("cargo build && cargo build --release") .allow(&[extracted_pattern]) .is_allow(); @@ -1201,6 +1215,32 @@ mod tests { .is_confirm(); } + #[test] + fn always_allow_button_works_without_subcommand() { + // When the second token is a flag (e.g. "ls -la"), the extracted pattern + // should only include the command name, not the flag. + let original_command = "ls -la"; + let extracted_pattern = pattern(original_command); + + // The extracted pattern should allow the original command + t(original_command).allow(&[extracted_pattern]).is_allow(); + + // It should allow other invocations of the same command + t("ls").allow(&[extracted_pattern]).is_allow(); + t("ls -R /tmp").allow(&[extracted_pattern]).is_allow(); + + // But not different commands + t("cat file.txt").allow(&[extracted_pattern]).is_confirm(); + + // Chained commands: all must match + t("ls -la && ls /tmp") + .allow(&[extracted_pattern]) + .is_allow(); + t("ls -la && cat file.txt") + .allow(&[extracted_pattern]) + .is_confirm(); + } + #[test] fn nested_command_substitution_all_checked() { t("echo $(cat $(whoami).txt)") diff --git a/crates/agent_ui/src/acp/thread_view.rs b/crates/agent_ui/src/acp/thread_view.rs index 4191fe20d35c519a151f8248d6492c5ef5205c4e..5c9b9e639e3e49c744fc21b8bb520b3d63f741ba 100644 --- a/crates/agent_ui/src/acp/thread_view.rs +++ b/crates/agent_ui/src/acp/thread_view.rs @@ -5007,7 +5007,7 @@ pub(crate) mod tests { "Missing 'Always for terminal' option" ); assert!( - labels.contains(&"Always for `cargo` commands"), + labels.contains(&"Always for `cargo build` commands"), "Missing pattern option" ); assert!(