fix(rewrite): keep pytest raw before a positional tail/head

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.
 

src/discover/registry.rs

@@ -1083,6 +1083,8 @@ struct PipelineAnalysis {
     next_clause_offset: Option<usize>,
     final_stage_start: Option<usize>,
     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,24 +1472,29 @@ 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],
 }
 
 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<Item = (&'static str, bool)> {
@@ -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,15 +4280,43 @@ 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", &[]),
             Some("rtk terraform plan | head -40".into())
         );
     }
 
+    /// #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!(