diff --git a/src/discover/README.md b/src/discover/README.md index bf1cc46dc1..77ad36b684 100644 --- a/src/discover/README.md +++ b/src/discover/README.md @@ -21,7 +21,7 @@ When a hook sends `cargo fmt --all && cargo test 2>&1 | tail -20`: → [Arg("cargo"), Arg("test"), Redirect("2>&1"), Operator("&&"), Arg("git"), Arg("status")] ``` -**Compound splitting** — The rewrite engine walks the tokens, splitting on `Operator` (`&&`, `||`, `;`, and the `case` terminators `;;`, `;&`, `;;&`) and typed `Pipe` tokens (`|`, `|&`). For normal pipelines, intermediate stages stay raw. A final stage whose rule's `pipeline_safety` allows it (ordinary `grep` and `rg` invocations) is rewritten; search pattern-file forms (`-f`/`--file`) defer because they can consume pipeline stdin as configuration. The producer stage is rewritten when its rule's `pipeline_safety` allows it and every downstream stage is a display-only consumer (`cat`, `head`, non-following `tail`) with no file-target redirect (#3171). Stderr pipelines (`|&`) and pipelines containing opaque shell groups remain raw. +**Compound splitting** — The rewrite engine walks the tokens, splitting on `Operator` (`&&`, `||`, `;`, and the `case` terminators `;;`, `;&`, `;;&`) and typed `Pipe` tokens (`|`, `|&`). For normal pipelines, intermediate stages stay raw. A final stage whose rule's `pipeline_safety` allows it (ordinary `grep` and `rg` invocations) is rewritten; search pattern-file forms (`-f`/`--file`) defer because they can consume pipeline stdin as configuration. The producer stage is rewritten when its rule's `pipeline_safety` allows it and every downstream stage is a display-only consumer (`cat`, `head`, non-following `tail`) with no file-target redirect (#3171). A producer whose rtk filter moves lines relative to the raw output (`rtk pytest` prints its summary first) stays raw in front of a positional consumer (`head`, `tail`), so the consumer still selects the lines the agent asked for (#4445). Stderr pipelines (`|&`) and pipelines containing opaque shell groups remain raw. **Quoted shell wrappers** — An exact `sh -c`, `bash -c`, `zsh -c` or `fish -c` segment whose script is one balanced quoted argument has that script rewritten in place, leaving the shell path, quote delimiters, spacing, suffix arguments and redirects byte-identical (`shell_wrapper.rs`). Recursion is capped at one level, so a wrapper inside a wrapper stays raw. A script carrying a newline, heredoc, `$((`, or any construct the lexer cannot attest defers unchanged. The wrapper names the shell, so the script is attested under *that* shell's grammar (`ShellDialect`): inside `fish -c`, `(cmd)` is substitution and `and`/`or`/`end` are control flow, and either one defers — while the same shapes in a host command string or a POSIX wrapper stay ordinary bash and keep rewriting. diff --git a/src/discover/registry.rs b/src/discover/registry.rs index e4d2d5381a..91a678556c 100644 --- a/src/discover/registry.rs +++ b/src/discover/registry.rs @@ -1083,6 +1083,8 @@ struct PipelineAnalysis { next_clause_offset: Option, final_stage_start: Option, all_consumers_safe: bool, + /// A consumer selects lines by position (`head`, `tail`) rather than showing them all. + positional_consumer: bool, } fn analyze_pipeline( @@ -1105,6 +1107,7 @@ fn analyze_pipeline( let mut final_stage_start = None; let mut has_supported_structure = true; let mut consumers_all_safe = true; + let mut positional_consumer = false; for (i, token) in tokens.iter().enumerate() { if token.offset >= end_offset { @@ -1126,10 +1129,11 @@ fn analyze_pipeline( if cmd[stage_start..token.offset].trim().is_empty() || kind == PipeKind::StdoutAndStderr { has_supported_structure = false; } - if token.offset > first_pipe_offset - && !is_safe_pipe_consumer(cmd[stage_start..token.offset].trim()) - { - consumers_all_safe = false; + if token.offset > first_pipe_offset { + match safe_pipe_consumer(cmd[stage_start..token.offset].trim()) { + Some(consumer) => positional_consumer |= consumer.positional, + None => consumers_all_safe = false, + } } stage_start = token.offset + token.value.len(); @@ -1138,8 +1142,11 @@ fn analyze_pipeline( if cmd[stage_start..end_offset].trim().is_empty() { has_supported_structure = false; - } else if !is_safe_pipe_consumer(cmd[stage_start..end_offset].trim()) { - consumers_all_safe = false; + } else { + match safe_pipe_consumer(cmd[stage_start..end_offset].trim()) { + Some(consumer) => positional_consumer |= consumer.positional, + None => consumers_all_safe = false, + } } PipelineAnalysis { @@ -1151,6 +1158,7 @@ fn analyze_pipeline( None }, all_consumers_safe: has_supported_structure && consumers_all_safe, + positional_consumer, } } @@ -1222,7 +1230,9 @@ fn rewrite_pipeline_producer( cmd, segment_start, first_pipe_offset, - RewriteContext::PipelineProducer, + RewriteContext::PipelineProducer { + positional_consumer: analysis.positional_consumer, + }, excluded, transparent_prefixes, shell_wrapper_depth, @@ -1462,6 +1472,8 @@ const PROCESS_WRAPPERS: &[ProcessWrapper] = &[ struct SafePipeConsumer { name: &'static str, + /// Selects lines by position, so it needs the producer's lines in their raw order. + positional: bool, unsafe_flags: &'static [&'static str], unsafe_flag_chars: &'static [char], } @@ -1469,17 +1481,20 @@ struct SafePipeConsumer { const SAFE_PIPE_CONSUMERS: &[SafePipeConsumer] = &[ SafePipeConsumer { name: "cat", + positional: false, unsafe_flags: &[], unsafe_flag_chars: &[], }, SafePipeConsumer { name: "head", + positional: true, unsafe_flags: &[], unsafe_flag_chars: &[], }, // #3171: only non-following tail is display-only SafePipeConsumer { name: "tail", + positional: true, unsafe_flags: &["--follow"], unsafe_flag_chars: &['f', 'F'], }, @@ -1500,18 +1515,21 @@ fn arg_matches_unsafe_flag(consumer: &SafePipeConsumer, arg: &str) -> bool { }) } -fn is_safe_pipe_consumer(stage: &str) -> bool { +fn safe_pipe_consumer(stage: &str) -> Option<&'static SafePipeConsumer> { let words = shell_split(stage); let mut words = words.iter(); - let Some(head) = words.next() else { - return false; - }; - let Some(consumer) = SAFE_PIPE_CONSUMERS.iter().find(|c| c.name == head.as_str()) else { - return false; - }; - !words.any(|arg| arg_matches_unsafe_flag(consumer, arg)) + let head = words.next()?; + let consumer = SAFE_PIPE_CONSUMERS + .iter() + .find(|c| c.name == head.as_str())?; + (!words.any(|arg| arg_matches_unsafe_flag(consumer, arg))).then_some(consumer) } +/// #4445: rtk filters that move lines relative to the raw output (`rtk pytest` prints its +/// summary first and tee pointers last, where pytest ends on the summary), so a positional +/// consumer would select different lines than the command the agent wrote. +const POSITION_SHIFTING_PRODUCERS: &[&str] = &["rtk pytest"]; + /// Every built-in transparent wrapper, paired with whether it may fall through. /// Derived from the two lists above so they cannot drift apart. fn builtin_transparent_prefixes() -> impl Iterator { @@ -1530,7 +1548,7 @@ const MAX_SHELL_WRAPPER_DEPTH: usize = 1; enum RewriteContext { Normal, PipelineFinal, - PipelineProducer, + PipelineProducer { positional_consumer: bool }, } /// Checks whether grep or rg reads patterns from a file. @@ -1924,10 +1942,13 @@ fn rewrite_segment_inner( return None; } // #3171 - if context == RewriteContext::PipelineProducer + if let RewriteContext::PipelineProducer { + positional_consumer, + } = context && (!rule.pipeline_safety.producer_safe() || !pipeline_command_is_safe(rule.rtk_cmd, cmd_part) - || !producer_output_is_line_faithful(rule.rtk_cmd, cmd_part)) + || !producer_output_is_line_faithful(rule.rtk_cmd, cmd_part) + || (positional_consumer && POSITION_SHIFTING_PRODUCERS.contains(&rule.rtk_cmd))) { return None; } @@ -4259,8 +4280,8 @@ mod tests { #[test] fn test_rewrite_pipe_producer_batch_rules_rewritten() { assert_eq!( - rewrite_command_no_prefixes("pytest | tail -20", &[]), - Some("rtk pytest | tail -20".into()) + rewrite_command_no_prefixes("pytest | cat", &[]), + Some("rtk pytest | cat".into()) ); assert_eq!( rewrite_command_no_prefixes("terraform plan | head -40", &[]), @@ -4268,6 +4289,34 @@ mod tests { ); } + /// #4445: `rtk pytest` prints its summary first and the tee pointers last, so a + /// positional consumer would select different lines than it does on raw pytest. + #[test] + fn test_rewrite_pipe_reordering_producer_before_positional_consumer_stays_raw() { + for cmd in [ + "python3 -m pytest t.py -q 2>&1 | tail -3", + "pytest -q | tail -n 5", + "pytest | tail", + "pytest -x | head -5", + "uv run pytest -q | tail -3", + "pytest -q | cat | tail -3", + ] { + assert_eq!(rewrite_command_no_prefixes(cmd, &[]), None, "{cmd}"); + } + assert_eq!( + rewrite_command_no_prefixes("pytest -q", &[]), + Some("rtk pytest -q".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("pytest -q 2>&1 | cat", &[]), + Some("rtk pytest -q 2>&1 | cat".into()) + ); + assert_eq!( + rewrite_command_no_prefixes("pytest -q | tail -3 && pytest -q", &[]), + Some("pytest -q | tail -3 && rtk pytest -q".into()) + ); + } + #[test] fn test_rewrite_pipe_final_grep_beats_producer_path() { assert_eq!(