From c536cb54f030e0d1c22223f24092601eee4fbc85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=86=A0=E8=BE=B0?= Date: Wed, 24 Jun 2026 16:40:03 +0800 Subject: [PATCH] fix(rewrite): handle process wrapper prefixes --- src/discover/registry.rs | 222 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 221 insertions(+), 1 deletion(-) diff --git a/src/discover/registry.rs b/src/discover/registry.rs index fc18b6be0..832ff938f 100644 --- a/src/discover/registry.rs +++ b/src/discover/registry.rs @@ -3,7 +3,7 @@ use lazy_static::lazy_static; use regex::{Regex, RegexSet}; -use super::lexer::{split_on_operators, tokenize, TokenKind}; +use super::lexer::{split_on_operators, tokenize, ParsedToken, TokenKind}; use super::rules::{IGNORED_EXACT, IGNORED_PREFIXES, RULES}; /// Result of classifying a command. @@ -770,6 +770,11 @@ fn rewrite_segment_inner( } } + if let Some((prefix, rest)) = strip_process_wrapper_prefix(trimmed) { + return rewrite_segment_inner(rest, excluded, transparent_prefixes, depth + 1) + .map(|rewritten| format!("{} {}", prefix, rewritten)); + } + // User-configured wrapper prefixes (e.g. `docker exec mycontainer`). Same // strip-recurse-reprepend contract as the builtin list above. for prefix in transparent_prefixes { @@ -860,6 +865,175 @@ fn rewrite_segment_inner( None } +fn strip_process_wrapper_prefix(cmd: &str) -> Option<(&str, &str)> { + let tokens = tokenize(cmd); + let args: Vec<&ParsedToken> = tokens + .iter() + .filter(|token| token.kind == TokenKind::Arg) + .collect(); + let first = args.first()?; + let inner_idx = match command_basename(&first.value) { + "timeout" => timeout_inner_command_index(&args)?, + "time" => time_inner_command_index(&args)?, + "nice" => nice_inner_command_index(&args)?, + "nohup" => nohup_inner_command_index(&args)?, + "stdbuf" => stdbuf_inner_command_index(&args)?, + _ => return None, + }; + + let inner = args.get(inner_idx)?; + let prefix = cmd[..inner.offset].trim_end(); + let rest = cmd[inner.offset..].trim_start(); + if rest.is_empty() { + None + } else { + Some((prefix, rest)) + } +} + +fn command_basename(command: &str) -> &str { + command.rsplit('/').next().unwrap_or(command) +} + +fn timeout_inner_command_index(args: &[&ParsedToken]) -> Option { + let mut idx = 1; + while idx < args.len() { + let arg = args[idx].value.as_str(); + if arg == "--" { + idx += 1; + break; + } + if !arg.starts_with('-') || arg == "-" { + break; + } + if timeout_option_takes_separate_value(arg) { + idx += 1; + if idx >= args.len() { + return None; + } + } + idx += 1; + } + + // `timeout` requires a duration before the command. + idx += 1; + (idx < args.len()).then_some(idx) +} + +fn timeout_option_takes_separate_value(arg: &str) -> bool { + matches!(arg, "-s" | "-k") || matches!(arg, "--signal" | "--kill-after") +} + +fn time_inner_command_index(args: &[&ParsedToken]) -> Option { + let mut idx = 1; + while idx < args.len() { + let arg = args[idx].value.as_str(); + if arg == "--" { + idx += 1; + break; + } + if !arg.starts_with('-') || arg == "-" { + break; + } + if time_option_takes_separate_value(arg) { + idx += 1; + if idx >= args.len() { + return None; + } + } + idx += 1; + } + + (idx < args.len()).then_some(idx) +} + +fn time_option_takes_separate_value(arg: &str) -> bool { + matches!(arg, "-f" | "-o") || matches!(arg, "--format" | "--output") +} + +fn nice_inner_command_index(args: &[&ParsedToken]) -> Option { + let mut idx = 1; + while idx < args.len() { + let arg = args[idx].value.as_str(); + if arg == "--" { + idx += 1; + break; + } + if arg == "-n" || arg == "--adjustment" { + idx += 2; + if idx > args.len() { + return None; + } + continue; + } + if arg.starts_with("--adjustment=") || is_nice_adjustment(arg) { + idx += 1; + continue; + } + if arg.starts_with('-') { + return None; + } + break; + } + + (idx < args.len()).then_some(idx) +} + +fn is_nice_adjustment(arg: &str) -> bool { + let Some(rest) = arg.strip_prefix('-').or_else(|| arg.strip_prefix('+')) else { + return false; + }; + !rest.is_empty() && rest.chars().all(|ch| ch.is_ascii_digit()) +} + +fn nohup_inner_command_index(args: &[&ParsedToken]) -> Option { + let idx = if args.get(1).is_some_and(|arg| arg.value == "--") { + 2 + } else { + 1 + }; + (idx < args.len()).then_some(idx) +} + +fn stdbuf_inner_command_index(args: &[&ParsedToken]) -> Option { + let mut idx = 1; + while idx < args.len() { + let arg = args[idx].value.as_str(); + if arg == "--" { + idx += 1; + break; + } + if !arg.starts_with('-') || arg == "-" { + break; + } + if stdbuf_option_takes_separate_value(arg) { + idx += 1; + if idx >= args.len() { + return None; + } + } else if !stdbuf_option_is_inline_or_flag(arg) { + return None; + } + idx += 1; + } + + (idx < args.len()).then_some(idx) +} + +fn stdbuf_option_takes_separate_value(arg: &str) -> bool { + matches!(arg, "-i" | "-o" | "-e") || matches!(arg, "--input" | "--output" | "--error") +} + +fn stdbuf_option_is_inline_or_flag(arg: &str) -> bool { + let short_mode = ["-i", "-o", "-e"] + .iter() + .any(|prefix| arg.starts_with(prefix) && arg.len() > prefix.len()); + let long_mode = ["--input=", "--output=", "--error="] + .iter() + .any(|prefix| arg.starts_with(prefix)); + short_mode || long_mode +} + /// Strip a command prefix with word-boundary check. /// Returns the remainder of the command after the prefix, or `None` if no match. fn strip_word_prefix<'a>(cmd: &'a str, prefix: &str) -> Option<&'a str> { @@ -3907,6 +4081,52 @@ mod tests { ); } + #[test] + fn test_process_wrapper_prefixes_rewrite_inner_command() { + assert_eq!( + rewrite_command_no_prefixes("timeout 300 cargo test", &[]), + Some("timeout 300 rtk cargo test".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("time cargo build", &[]), + Some("time rtk cargo build".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("nice -n 10 cargo test", &[]), + Some("nice -n 10 rtk cargo test".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("nohup cargo build", &[]), + Some("nohup rtk cargo build".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("stdbuf -oL cargo test", &[]), + Some("stdbuf -oL rtk cargo test".into()) + ); + } + + #[test] + fn test_process_wrapper_prefixes_handle_options_and_compounds() { + assert_eq!( + rewrite_command_no_prefixes("timeout -k 5s 300 cargo test", &[]), + Some("timeout -k 5s 300 rtk cargo test".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("CI=1 timeout 300 cargo test", &[]), + Some("CI=1 timeout 300 rtk cargo test".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("timeout 300 cargo test && time git status", &[]), + Some("timeout 300 rtk cargo test && time rtk git status".into()) + ); + } + + #[test] + fn test_process_wrapper_prefix_without_inner_command_returns_none() { + assert_eq!(rewrite_command_no_prefixes("timeout 300", &[]), None); + assert_eq!(rewrite_command_no_prefixes("time", &[]), None); + } + #[test] fn test_transparent_prefix_multiple_configured() { let prefixes = vec!["shadowenv exec --".to_string(), "direnv exec .".to_string()];