Skip to main content

yarp_cli/
rewrite.rs

1use crate::rules::{PackRequest, Registry, Selection, digest_hex};
2use crate::shell::{self, Connector, ShellItem, SimpleCommand};
3use serde::Serialize;
4use yarp_rule_pack::{OutputPolicy, Rule, Transform};
5
6/// Metadata passed to an archived shell wrapper.
7#[derive(Clone, Copy, Debug)]
8pub struct ArchiveCommandRef<'a> {
9    pub agent: &'a str,
10    pub account: &'a str,
11    pub session_id: &'a str,
12    pub call_id: &'a str,
13}
14
15/// Return a wrapper command only for a simple shell command selected by a built-in rule.
16#[must_use]
17pub fn rewrite(command: &str) -> Option<String> {
18    rewrite_with_archive(command, None)
19}
20
21/// Return an archived wrapper command for a command selected by a built-in rule.
22#[must_use]
23pub fn rewrite_with_archive(
24    command: &str,
25    archive: Option<ArchiveCommandRef<'_>>,
26) -> Option<String> {
27    rewrite_with_options(command, archive, &[]).ok().flatten()
28}
29
30#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
31pub struct ShellPlan {
32    pub version: u32,
33    pub execution: ExecutionPlan,
34    pub result: ResultPolicy,
35}
36
37#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
38#[serde(tag = "kind", rename_all = "snake_case")]
39pub enum ExecutionPlan {
40    Original,
41    Rewrite { command: String },
42}
43
44#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
45#[serde(tag = "kind", rename_all = "snake_case")]
46pub enum ResultPolicy {
47    Ordinary,
48    Recovery,
49}
50
51/// Plan shell execution and result handling using one conservative parse.
52///
53/// # Errors
54///
55/// Returns an error when configured rule packs cannot be loaded or selected safely.
56pub fn plan_with_options(
57    command: &str,
58    archive: Option<ArchiveCommandRef<'_>>,
59    packs: &[PackRequest],
60) -> Result<ShellPlan, String> {
61    let command = command.trim();
62    let Some(words) = shell::parse_simple_words(command).ok() else {
63        return Ok(original_plan(ResultPolicy::Ordinary));
64    };
65    if is_recovery_command(&words) {
66        return Ok(original_plan(ResultPolicy::Recovery));
67    }
68    let mut registry = Registry::load(packs)?;
69    let Selection::Reduce(selected) = registry.select(&words)? else {
70        return Ok(original_plan(ResultPolicy::Ordinary));
71    };
72    let wrapper = selected_wrapper(command, archive, &registry, &selected)?;
73    Ok(ShellPlan {
74        version: 1,
75        execution: ExecutionPlan::Rewrite { command: wrapper },
76        result: ResultPolicy::Ordinary,
77    })
78}
79
80/// Return a wrapper command using built-in and explicitly supplied compiled rule packs.
81///
82/// # Errors
83///
84/// Returns an error when a configured pack cannot be loaded or a selected pack path cannot be
85/// represented in a shell wrapper.
86pub fn rewrite_with_options(
87    command: &str,
88    archive: Option<ArchiveCommandRef<'_>>,
89    packs: &[PackRequest],
90) -> Result<Option<String>, String> {
91    plan_with_options(command, archive, packs).map(|plan| match plan.execution {
92        ExecutionPlan::Original => None,
93        ExecutionPlan::Rewrite { command } => Some(command),
94    })
95}
96
97const fn original_plan(result: ResultPolicy) -> ShellPlan {
98    ShellPlan {
99        version: 1,
100        execution: ExecutionPlan::Original,
101        result,
102    }
103}
104
105fn is_recovery_command(words: &[String]) -> bool {
106    matches!(
107        words,
108        [program, subcommand, arguments @ ..]
109            if program == "yarp"
110                && matches!(subcommand.as_str(), "search" | "read")
111                && !arguments.iter().any(|argument| matches!(argument.as_str(), "--help" | "-h"))
112    )
113}
114
115fn selected_wrapper(
116    command: &str,
117    archive: Option<ArchiveCommandRef<'_>>,
118    registry: &Registry,
119    selected: &crate::rules::SelectedRule,
120) -> Result<String, String> {
121    let mut wrapper = format!(
122        "yarp run --selected-pack {} --selected-rule {} --selected-digest {}",
123        shell_quote(&selected.pack_id),
124        shell_quote(&selected.rule.id),
125        shell_quote(&digest_hex(&selected.source_digest)),
126    );
127    for reference in registry.references() {
128        let path = reference
129            .path
130            .to_str()
131            .ok_or_else(|| "rule pack path is not valid UTF-8".to_owned())?;
132        wrapper.push_str(" --rule-pack ");
133        wrapper.push_str(&shell_quote(path));
134        wrapper.push_str(" --rule-pack-digest ");
135        wrapper.push_str(&shell_quote(&digest_hex(&reference.source_digest)));
136        wrapper.push_str(" --rule-pack-compiled-digest ");
137        wrapper.push_str(&shell_quote(&digest_hex(&reference.compiled_digest)));
138    }
139    if let Some(reference) = archive {
140        wrapper.push_str(" --archive-agent ");
141        wrapper.push_str(&shell_quote(reference.agent));
142        wrapper.push_str(" --archive-account ");
143        wrapper.push_str(&shell_quote(reference.account));
144        wrapper.push_str(" --archive-session ");
145        wrapper.push_str(&shell_quote(reference.session_id));
146        wrapper.push_str(" --archive-call ");
147        wrapper.push_str(&shell_quote(reference.call_id));
148    }
149    wrapper.push_str(" -- ");
150    wrapper.push_str(command);
151    Ok(wrapper)
152}
153
154fn shell_quote(value: &str) -> String {
155    format!("'{}'", value.replace('\'', "'\\''"))
156}
157
158/// Parse one conservative shell command and select it using built-in rules.
159///
160/// # Errors
161///
162/// Returns an error when the shell source is unsupported or rule selection fails.
163pub fn select_builtin_command(command: &str) -> Result<(Vec<String>, Selection), String> {
164    let words = shell::parse_simple_words(command)?;
165    let selection = Registry::builtins_only().select(&words)?;
166    Ok((words, selection))
167}
168
169/// Select a command from an already parsed argument list using built-in rules.
170///
171/// # Errors
172///
173/// Returns an error only when the embedded registry cannot read a selected record.
174pub fn select_builtin_argv(arguments: &[String]) -> Result<Selection, String> {
175    Registry::builtins_only().select(arguments)
176}
177
178#[derive(Clone, Copy, Debug, Eq, PartialEq)]
179pub enum StatusConfidence {
180    Complete,
181    FinalStageOnly,
182    Conditional,
183}
184
185const TRANSFORM_DIAGNOSTIC_PREFIXES: [(&str, &[u8]); 6] = [
186    ("cat", b"cat:"),
187    ("tee", b"tee:"),
188    ("head", b"head:"),
189    ("tail", b"tail:"),
190    ("sort", b"sort:"),
191    ("uniq", b"uniq:"),
192];
193
194#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
195pub struct TransformDiagnostics(u8);
196
197impl TransformDiagnostics {
198    fn insert(&mut self, program: &str) {
199        if let Some(index) = TRANSFORM_DIAGNOSTIC_PREFIXES
200            .iter()
201            .position(|(candidate, _)| *candidate == program)
202        {
203            self.0 |= 1_u8 << index;
204        }
205    }
206
207    fn merge(&mut self, other: Self) {
208        self.0 |= other.0;
209    }
210
211    #[must_use]
212    pub fn matches_line(self, line: &[u8]) -> bool {
213        TRANSFORM_DIAGNOSTIC_PREFIXES
214            .iter()
215            .enumerate()
216            .any(|(index, (_, prefix))| self.0 & (1_u8 << index) != 0 && line.starts_with(prefix))
217    }
218}
219
220#[derive(Clone, Debug, Eq, PartialEq)]
221pub struct ResultPlan {
222    pub rule: Rule,
223    pub status_confidence: StatusConfidence,
224    pub transform_diagnostics: TransformDiagnostics,
225    pub fail_open_setup_diagnostics: bool,
226}
227
228/// Select one compatible typed plan for conservative post-result reduction.
229///
230/// # Errors
231///
232/// Returns an error when shell syntax is ambiguous, guarded, unsupported, or mixes reducer
233/// families. Callers fail open on every error.
234pub fn select_result_plan(command: &str) -> Result<ResultPlan, String> {
235    let program = shell::parse(command)?;
236    let mut selected: Option<Rule> = None;
237    let mut confidence = StatusConfidence::Complete;
238    let mut transform_diagnostics = TransformDiagnostics::default();
239    let mut fail_open_setup_diagnostics = false;
240    let mut pipefail = Some(false);
241    let mut previous_had_output = false;
242
243    for (index, item) in program.items.iter().enumerate() {
244        let connector = index
245            .checked_sub(1)
246            .and_then(|connector| program.connectors.get(connector))
247            .copied();
248        if let ShellItem::Simple(command) = item
249            && let Some(value) = pipefail_setting(&command.words)
250        {
251            pipefail = if matches!(connector, None | Some(Connector::Sequence)) {
252                Some(value)
253            } else {
254                None
255            };
256        }
257        let (candidate, setup) = match item {
258            ShellItem::Simple(command) => select_simple(command)?,
259            ShellItem::Pipeline(stages) => {
260                if pipefail != Some(true) {
261                    confidence = merge_confidence(confidence, StatusConfidence::FinalStageOnly);
262                }
263                let (rule, diagnostics) = select_pipeline(stages)?;
264                transform_diagnostics.merge(diagnostics);
265                (Some(rule), false)
266            }
267        };
268        let has_output = candidate.is_some();
269        let may_have_setup_diagnostics =
270            setup && matches!(item, ShellItem::Simple(command) if setup_may_emit(&command.words));
271        fail_open_setup_diagnostics |= may_have_setup_diagnostics;
272        let may_have_visible_output = has_output || may_have_setup_diagnostics;
273        match connector {
274            Some(Connector::Sequence) if previous_had_output => {
275                confidence = merge_confidence(confidence, StatusConfidence::FinalStageOnly);
276            }
277            Some(Connector::Or) if previous_had_output || may_have_visible_output => {
278                confidence = merge_confidence(confidence, StatusConfidence::Conditional);
279            }
280            None if index > 0 => {
281                return Err("compound command is missing a connector".to_owned());
282            }
283            Some(Connector::And | Connector::Sequence | Connector::Or) | None => {}
284        }
285        if let Some(candidate) = candidate {
286            merge_rule(&mut selected, candidate)?;
287        }
288        previous_had_output = previous_had_output || may_have_visible_output;
289    }
290
291    let rule =
292        selected.ok_or_else(|| "compound command has no supported output command".to_owned())?;
293    Ok(ResultPlan {
294        rule,
295        status_confidence: confidence,
296        transform_diagnostics,
297        fail_open_setup_diagnostics,
298    })
299}
300
301/// Select one compatible typed rule for callers that do not need status confidence.
302///
303/// # Errors
304///
305/// Returns the same errors as [`select_result_plan`].
306pub fn select_result_rule(command: &str) -> Result<Rule, String> {
307    select_result_plan(command).map(|plan| plan.rule)
308}
309
310fn select_simple(command: &SimpleCommand) -> Result<(Option<Rule>, bool), String> {
311    if is_setup_command(&command.words) || command.words.iter().all(|word| is_assignment(word)) {
312        return Ok((None, true));
313    }
314    let Some(words) = normalize_result_words(command.words.clone())? else {
315        return Ok((None, true));
316    };
317    let mut registry = Registry::builtins_only();
318    match registry
319        .select(&words)
320        .map_err(|error| format!("could not classify compound command: {error}"))?
321    {
322        Selection::Reduce(candidate) => Ok((Some((*candidate.rule).clone()), false)),
323        Selection::Transform(_) => Err("line transform has no pipeline input".to_owned()),
324        Selection::Passthrough(_) | Selection::Ambiguous(_) | Selection::Unsupported => {
325            Err("compound command contains an unsupported or guarded command".to_owned())
326        }
327    }
328}
329
330fn select_pipeline(stages: &[SimpleCommand]) -> Result<(Rule, TransformDiagnostics), String> {
331    let mut selected = None;
332    let mut diagnostics = TransformDiagnostics::default();
333    for stage in stages {
334        if is_setup_command(&stage.words) || stage.words.iter().all(|word| is_assignment(word)) {
335            return Err("pipeline contains a setup command".to_owned());
336        }
337        let words = normalize_result_words(stage.words.clone())?
338            .ok_or_else(|| "pipeline stage has no output command".to_owned())?;
339        let mut registry = Registry::builtins_only();
340        match registry
341            .select(&words)
342            .map_err(|error| format!("could not classify pipeline stage: {error}"))?
343        {
344            Selection::Reduce(candidate) => merge_rule(&mut selected, (*candidate.rule).clone())?,
345            Selection::Transform(candidate) => {
346                if candidate.rule.transform != Some(Transform::LinePreserving) {
347                    return Err("pipeline transform is unsupported".to_owned());
348                }
349                validate_line_preserving(&words)?;
350                if selected.is_none() {
351                    return Err("pipeline transform has no typed input".to_owned());
352                }
353                diagnostics.insert(&words[0]);
354            }
355            Selection::Passthrough(_) | Selection::Ambiguous(_) | Selection::Unsupported => {
356                return Err("pipeline contains an unsupported or guarded command".to_owned());
357            }
358        }
359    }
360    selected
361        .map(|rule| (rule, diagnostics))
362        .ok_or_else(|| "pipeline has no supported output command".to_owned())
363}
364
365fn validate_line_preserving(words: &[String]) -> Result<(), String> {
366    let Some((program, arguments)) = words.split_first() else {
367        return Err("pipeline transform has no program".to_owned());
368    };
369    match program.as_str() {
370        "cat" => validate_stdin_only(arguments),
371        "tee" => validate_tee(arguments),
372        "head" | "tail" => validate_line_selector(program, arguments),
373        "sort" => validate_sort(arguments),
374        "uniq" => validate_uniq(arguments),
375        _ => Err("pipeline transform program is unsupported".to_owned()),
376    }
377}
378
379fn validate_stdin_only(arguments: &[String]) -> Result<(), String> {
380    let mut operands = false;
381    for argument in arguments {
382        if argument == "--" && !operands {
383            operands = true;
384        } else if argument != "-" {
385            return Err("pipeline transform must read only standard input".to_owned());
386        }
387    }
388    Ok(())
389}
390
391fn validate_tee(arguments: &[String]) -> Result<(), String> {
392    let mut operands = false;
393    for argument in arguments {
394        if operands {
395            continue;
396        }
397        match argument.as_str() {
398            "--" => operands = true,
399            "-a" | "-i" | "-p" | "--append" | "--ignore-interrupts" | "--output-error" => {}
400            value if value.starts_with("--output-error=") => {}
401            value if !value.starts_with('-') || value == "-" => operands = true,
402            _ => return Err("tee transform uses an unsupported option".to_owned()),
403        }
404    }
405    Ok(())
406}
407
408fn validate_line_selector(program: &str, arguments: &[String]) -> Result<(), String> {
409    let mut index = 0;
410    let mut operands = false;
411    while index < arguments.len() {
412        let argument = &arguments[index];
413        if operands {
414            if argument != "-" {
415                return Err(format!("{program} transform must read only standard input"));
416            }
417            index += 1;
418            continue;
419        }
420        if argument == "--" {
421            operands = true;
422            index += 1;
423            continue;
424        }
425        if argument == "-" {
426            operands = true;
427            index += 1;
428            continue;
429        }
430        if matches!(argument.as_str(), "-q" | "--quiet" | "--silent") {
431            index += 1;
432            continue;
433        }
434        if argument == "-n" || argument == "--lines" {
435            let count = arguments
436                .get(index + 1)
437                .ok_or_else(|| format!("{program} line option is missing its value"))?;
438            validate_line_count(count, program)?;
439            index += 2;
440            continue;
441        }
442        if let Some(count) = argument.strip_prefix("--lines=") {
443            validate_line_count(count, program)?;
444            index += 1;
445            continue;
446        }
447        if let Some(count) = argument.strip_prefix("-n")
448            && !count.is_empty()
449        {
450            validate_line_count(count, program)?;
451            index += 1;
452            continue;
453        }
454        if argument.strip_prefix('-').is_some_and(|count| {
455            !count.is_empty() && count.bytes().all(|byte| byte.is_ascii_digit())
456        }) {
457            index += 1;
458            continue;
459        }
460        return Err(format!(
461            "{program} transform uses an unsupported option or operand"
462        ));
463    }
464    Ok(())
465}
466
467fn validate_line_count(value: &str, program: &str) -> Result<(), String> {
468    let digits = value.strip_prefix(['+', '-']).unwrap_or(value).as_bytes();
469    if digits.is_empty() || !digits.iter().all(u8::is_ascii_digit) {
470        return Err(format!("{program} line count is not a literal integer"));
471    }
472    Ok(())
473}
474
475fn validate_sort(arguments: &[String]) -> Result<(), String> {
476    let mut index = 0;
477    let mut operands = false;
478    while index < arguments.len() {
479        let argument = &arguments[index];
480        if operands || argument == "-" || !argument.starts_with('-') {
481            if argument != "-" {
482                return Err("sort transform must read only standard input".to_owned());
483            }
484            operands = true;
485            index += 1;
486            continue;
487        }
488        if argument == "--" {
489            operands = true;
490            index += 1;
491            continue;
492        }
493        if sort_option_takes_value(argument) {
494            if sort_option_has_attached_value(argument) {
495                index += 1;
496            } else {
497                index = index
498                    .checked_add(2)
499                    .filter(|next| *next <= arguments.len())
500                    .ok_or_else(|| "sort option is missing its value".to_owned())?;
501            }
502            continue;
503        }
504        if is_safe_sort_flag(argument) {
505            index += 1;
506            continue;
507        }
508        return Err("sort transform uses an unsupported option".to_owned());
509    }
510    Ok(())
511}
512
513fn sort_option_takes_value(value: &str) -> bool {
514    [
515        "-k",
516        "-t",
517        "-S",
518        "-T",
519        "--key",
520        "--field-separator",
521        "--buffer-size",
522        "--temporary-directory",
523        "--batch-size",
524        "--random-source",
525    ]
526    .iter()
527    .any(|option| value == *option || value.starts_with(&format!("{option}=")))
528        || ["-k", "-t", "-S", "-T"]
529            .iter()
530            .any(|option| value.starts_with(option) && value.len() > option.len())
531}
532
533fn sort_option_has_attached_value(value: &str) -> bool {
534    value.contains('=')
535        || ["-k", "-t", "-S", "-T"]
536            .iter()
537            .any(|option| value.starts_with(option) && value.len() > option.len())
538}
539
540fn is_safe_sort_flag(value: &str) -> bool {
541    if value.starts_with("--") {
542        return matches!(
543            value,
544            "--stable"
545                | "--unique"
546                | "--reverse"
547                | "--numeric-sort"
548                | "--general-numeric-sort"
549                | "--human-numeric-sort"
550                | "--version-sort"
551                | "--month-sort"
552                | "--random-sort"
553                | "--ignore-case"
554                | "--ignore-leading-blanks"
555                | "--dictionary-order"
556                | "--ignore-nonprinting"
557        );
558    }
559    value.len() > 1
560        && value[1..].bytes().all(|byte| {
561            matches!(
562                byte,
563                b'b' | b'd'
564                    | b'f'
565                    | b'g'
566                    | b'h'
567                    | b'i'
568                    | b'M'
569                    | b'n'
570                    | b'R'
571                    | b'r'
572                    | b's'
573                    | b'u'
574                    | b'V'
575            )
576        })
577}
578
579fn validate_uniq(arguments: &[String]) -> Result<(), String> {
580    let mut index = 0;
581    let mut operands = false;
582    let mut operand_count = 0_usize;
583    while index < arguments.len() {
584        let argument = &arguments[index];
585        if operands || argument == "-" || !argument.starts_with('-') {
586            if argument != "-" || operand_count > 0 {
587                return Err("uniq transform must read only standard input".to_owned());
588            }
589            operands = true;
590            operand_count += 1;
591            index += 1;
592            continue;
593        }
594        if argument == "--" {
595            operands = true;
596            index += 1;
597            continue;
598        }
599        if uniq_option_takes_value(argument) {
600            if uniq_option_has_attached_value(argument) {
601                index += 1;
602            } else {
603                index = index
604                    .checked_add(2)
605                    .filter(|next| *next <= arguments.len())
606                    .ok_or_else(|| "uniq option is missing its value".to_owned())?;
607            }
608            continue;
609        }
610        if is_safe_uniq_flag(argument) {
611            index += 1;
612            continue;
613        }
614        return Err("uniq transform uses an unsupported option".to_owned());
615    }
616    Ok(())
617}
618
619fn uniq_option_takes_value(value: &str) -> bool {
620    [
621        "-f",
622        "-s",
623        "-w",
624        "--skip-fields",
625        "--skip-chars",
626        "--check-chars",
627    ]
628    .iter()
629    .any(|option| value == *option || value.starts_with(&format!("{option}=")))
630        || ["-f", "-s", "-w"]
631            .iter()
632            .any(|option| value.starts_with(option) && value.len() > option.len())
633}
634
635fn uniq_option_has_attached_value(value: &str) -> bool {
636    value.contains('=')
637        || ["-f", "-s", "-w"]
638            .iter()
639            .any(|option| value.starts_with(option) && value.len() > option.len())
640}
641
642fn is_safe_uniq_flag(value: &str) -> bool {
643    if value.starts_with("--") {
644        return matches!(
645            value,
646            "--all-repeated" | "--repeated" | "--unique" | "--ignore-case"
647        );
648    }
649    value.len() > 1
650        && value[1..]
651            .bytes()
652            .all(|byte| matches!(byte, b'D' | b'd' | b'u' | b'i'))
653}
654
655fn merge_rule(selected: &mut Option<Rule>, candidate: Rule) -> Result<(), String> {
656    let Some(existing) = selected else {
657        *selected = Some(candidate);
658        return Ok(());
659    };
660    if existing.reducer != candidate.reducer {
661        return Err("compound command mixes reducer families".to_owned());
662    }
663    let existing_success = existing
664        .success
665        .ok_or_else(|| "selected rule has no success policy".to_owned())?;
666    let candidate_success = candidate
667        .success
668        .ok_or_else(|| "candidate rule has no success policy".to_owned())?;
669    let existing_failure = existing
670        .failure
671        .ok_or_else(|| "selected rule has no failure policy".to_owned())?;
672    let candidate_failure = candidate
673        .failure
674        .ok_or_else(|| "candidate rule has no failure policy".to_owned())?;
675    existing.success = Some(merge_policy(existing_success, candidate_success));
676    existing.failure = Some(merge_policy(existing_failure, candidate_failure));
677    Ok(())
678}
679
680const fn merge_confidence(left: StatusConfidence, right: StatusConfidence) -> StatusConfidence {
681    match (left, right) {
682        (StatusConfidence::Conditional, _) | (_, StatusConfidence::Conditional) => {
683            StatusConfidence::Conditional
684        }
685        (StatusConfidence::FinalStageOnly, _) | (_, StatusConfidence::FinalStageOnly) => {
686            StatusConfidence::FinalStageOnly
687        }
688        (StatusConfidence::Complete, StatusConfidence::Complete) => StatusConfidence::Complete,
689    }
690}
691
692fn pipefail_setting(words: &[String]) -> Option<bool> {
693    match words {
694        [program, option, value] if program == "set" && value == "pipefail" && option == "-o" => {
695            Some(true)
696        }
697        [program, option, value] if program == "set" && value == "pipefail" && option == "+o" => {
698            Some(false)
699        }
700        _ => None,
701    }
702}
703
704fn merge_policy(left: OutputPolicy, right: OutputPolicy) -> OutputPolicy {
705    OutputPolicy {
706        max_line_bytes: left.max_line_bytes.max(right.max_line_bytes),
707        max_output_bytes: left.max_output_bytes.max(right.max_output_bytes),
708        min_savings_bytes: left.min_savings_bytes.max(right.min_savings_bytes),
709        min_savings_basis_points: left
710            .min_savings_basis_points
711            .max(right.min_savings_basis_points),
712    }
713}
714
715fn normalize_result_words(mut words: Vec<String>) -> Result<Option<Vec<String>>, String> {
716    if is_setup_command(&words) {
717        return Ok(None);
718    }
719    while words.first().is_some_and(|word| is_assignment(word)) {
720        words.remove(0);
721    }
722    let Some(program) = words.first().map(String::as_str) else {
723        return Ok(None);
724    };
725    match program {
726        "env" => {
727            words.remove(0);
728            if words.first().is_some_and(|word| word == "--") {
729                words.remove(0);
730            }
731            while words.first().is_some_and(|word| is_assignment(word)) {
732                words.remove(0);
733            }
734        }
735        "command" | "exec" => {
736            words.remove(0);
737            if words.first().is_some_and(|word| word == "--") {
738                words.remove(0);
739            }
740        }
741        "time" if words.get(1).is_some_and(|word| !word.starts_with('-')) => {
742            words.remove(0);
743        }
744        "timeout"
745            if words.len() >= 3 && words.get(1).is_some_and(|word| !word.starts_with('-')) =>
746        {
747            words.drain(0..2);
748        }
749        _ => {}
750    }
751    if words.is_empty() {
752        return Ok(None);
753    }
754    if words.first().is_some_and(|word| {
755        matches!(
756            word.as_str(),
757            "env" | "command" | "exec" | "time" | "timeout"
758        )
759    }) {
760        return Err("compound command uses an unsupported command wrapper".to_owned());
761    }
762    Ok(Some(words))
763}
764
765fn is_setup_command(words: &[String]) -> bool {
766    match words {
767        [program] if matches!(program.as_str(), ":" | "true" | "false") => true,
768        [program, path] if program == "cd" && !path.starts_with('-') => true,
769        [program, flag] if program == "set" && is_quiet_set_flag(flag) => true,
770        [program, flag, value]
771            if program == "set" && matches!(flag.as_str(), "-o" | "+o") && value == "pipefail" =>
772        {
773            true
774        }
775        [program, values @ ..] if program == "export" && !values.is_empty() => values
776            .iter()
777            .all(|value| is_assignment(value) || is_identifier(value)),
778        [program, value] if program == "umask" && !value.starts_with('-') => true,
779        _ => false,
780    }
781}
782
783/// Return whether output contains a diagnostic from an accepted fallible setup command.
784#[must_use]
785pub fn contains_setup_diagnostic(input: &[u8]) -> bool {
786    [
787        b"cd:".as_slice(),
788        b"umask:".as_slice(),
789        b"export:".as_slice(),
790        b"readonly variable".as_slice(),
791        b"not a valid identifier".as_slice(),
792    ]
793    .iter()
794    .any(|pattern| {
795        input.windows(pattern.len()).any(|window| {
796            window
797                .iter()
798                .zip(*pattern)
799                .all(|(left, right)| left.eq_ignore_ascii_case(right))
800        })
801    })
802}
803
804fn setup_may_emit(words: &[String]) -> bool {
805    words.iter().all(|word| is_assignment(word))
806        || words
807            .first()
808            .is_some_and(|program| matches!(program.as_str(), "cd" | "export" | "umask"))
809}
810
811fn is_quiet_set_flag(value: &str) -> bool {
812    matches!(
813        value,
814        "-e" | "+e"
815            | "-u"
816            | "+u"
817            | "-f"
818            | "+f"
819            | "-C"
820            | "+C"
821            | "-m"
822            | "+m"
823            | "-b"
824            | "+b"
825            | "-n"
826            | "+n"
827            | "-E"
828            | "+E"
829            | "-T"
830            | "+T"
831            | "-P"
832            | "+P"
833            | "-h"
834            | "+h"
835            | "-k"
836            | "+k"
837            | "-p"
838            | "+p"
839    )
840}
841
842fn is_identifier(value: &str) -> bool {
843    let mut characters = value.chars();
844    characters
845        .next()
846        .is_some_and(|character| character == '_' || character.is_ascii_alphabetic())
847        && characters.all(|character| character == '_' || character.is_ascii_alphanumeric())
848}
849
850fn is_assignment(word: &str) -> bool {
851    word.split_once('=')
852        .is_some_and(|(name, _)| is_identifier(name))
853}
854
855#[cfg(test)]
856mod tests {
857    use super::*;
858
859    #[test]
860    fn rewrites_selected_commands_without_changing_the_original_text() {
861        let expected = format!(
862            "yarp run --selected-pack 'yarp-builtins' --selected-rule 'git/diff' --selected-digest '{}' -- git diff -- 'file name'",
863            digest_hex(&crate::rules::BUILTIN_SOURCE_DIGEST)
864        );
865        assert_eq!(
866            rewrite("  git diff -- 'file name'  ").as_deref(),
867            Some(expected.as_str())
868        );
869        assert!(
870            rewrite("npm run lint -- --fix")
871                .is_some_and(|value| value.ends_with("-- npm run lint -- --fix"))
872        );
873        assert!(rewrite("pytest -q").is_some());
874        assert!(rewrite("rg -n needle .").is_some());
875        assert!(rewrite("egrep -n needle file").is_some());
876        assert!(rewrite("fgrep needle file").is_some());
877        assert!(rewrite("podman ps").is_some());
878        assert!(rewrite("git reflog --date=iso").is_some());
879        assert!(rewrite("git remote show origin").is_some());
880        assert!(rewrite("gh repo view osolmaz/yarp").is_some());
881        assert!(rewrite("pnpm -s build").is_some());
882        assert!(rewrite("git -C /repo diff --check").is_some());
883        assert!(rewrite("pnpm -C frontend test").is_some());
884        assert!(rewrite("pnpm --filter web -r lint").is_some());
885        assert!(rewrite("ls -la").is_some());
886        assert!(rewrite("find src -name '*.rs'").is_some());
887        assert!(rewrite("ps aux").is_some());
888    }
889
890    #[test]
891    fn adds_archive_identifiers_with_shell_quoting() {
892        let rewritten = rewrite_with_archive(
893            "git status --short",
894            Some(ArchiveCommandRef {
895                agent: "pi",
896                account: "o'nur",
897                session_id: "session-1",
898                call_id: "call-1",
899            }),
900        )
901        .expect("rewrite");
902        assert_eq!(
903            rewritten,
904            format!(
905                "yarp run --selected-pack 'yarp-builtins' --selected-rule 'git/status' --selected-digest '{}' --archive-agent 'pi' --archive-account 'o'\\''nur' --archive-session 'session-1' --archive-call 'call-1' -- git status --short",
906                digest_hex(&crate::rules::BUILTIN_SOURCE_DIGEST)
907            )
908        );
909    }
910
911    #[test]
912    fn classifies_only_direct_archive_queries_as_recovery() {
913        for command in [
914            "yarp search yr_0123456789abcdef0123456789abcdef error",
915            "yarp read yr_0123456789abcdef0123456789abcdef stdout 1:20",
916            "yarp search yr_0123456789abcdef0123456789abcdef 'error|warning' -C 2",
917        ] {
918            let plan = plan_with_options(command, None, &[]).expect("recovery plan");
919            assert_eq!(plan.execution, ExecutionPlan::Original);
920            assert_eq!(plan.result, ResultPolicy::Recovery);
921        }
922        for command in [
923            "command yarp search ref error",
924            "env DEBUG=1 yarp search ref error",
925            "/usr/bin/yarp read ref 1:20",
926            "yarp search ref error | head",
927            "yarp read ref 1:20 > output",
928            "yarp search ref error; printf extra",
929            "yarp search $(printf ref) error",
930            "yarp search ref error &",
931            "yarp search --help",
932        ] {
933            let plan = plan_with_options(command, None, &[]).expect("ordinary plan");
934            assert_eq!(
935                plan.result,
936                ResultPolicy::Ordinary,
937                "classified {command:?}"
938            );
939        }
940    }
941
942    #[test]
943    fn leaves_unknown_and_guarded_commands_unchanged() {
944        for command in [
945            "",
946            "cat .env",
947            "curl https://example.com",
948            "git push",
949            "npm install",
950            "git show HEAD:file",
951            "rg --json needle .",
952            "rg --files",
953            "rg -l needle .",
954            "rg -0l needle .",
955            "grep -l needle file",
956            "grep -lZ needle file",
957            "egrep -l needle file",
958            "fgrep --count needle file",
959            "git status -sbz",
960            "git grep -l needle",
961            "git grep -lz needle",
962            "git log --oneline -z",
963            "ls --zero",
964            "find . -print0",
965            "find . -exec echo {} +",
966            "fd -0 needle .",
967            "fd -x echo {}",
968            "tree -J .",
969            "du --null .",
970            "df --output=source,target",
971            "lsof -F pcfn",
972            "free --seconds 1",
973            "ps aux -o pid=",
974            "cmake -S . -B build --trace",
975            "git worktree list --porcelain",
976            "git tag --list -z",
977            "git reflog delete main@{0}",
978            "git reflog --format=%H",
979            "gh workflow view ci.yml --yaml",
980            "docker ps --format json",
981            "docker images --quiet",
982            "podman ps --format json",
983            "podman logs --follow container",
984            "journalctl --output=json",
985            "systemctl status -o json sshd",
986            "ninja -t compdb",
987            "tsc --showConfig",
988            "eslint --format json .",
989            "npx eslint --format=json .",
990            "pyright --outputjson",
991            "ruff check --output-format json .",
992            "uv run ruff --show-settings",
993            "mypy --output=json .",
994            "pytest --json-report",
995            "python3 -m pytest --collect-only",
996            "go vet -json ./...",
997            "go test -json ./...",
998            "cargo test --message-format=json",
999            "cargo audit --json",
1000            "pnpm exec oxlint --format json",
1001            "gh pr view 1 --json=number,title",
1002            "git diff --stat",
1003            "kubectl get pods -ojson",
1004            "npm test -- --reporter=json",
1005            "codex review --json --base main",
1006            "yarp run -- git status",
1007        ] {
1008            assert_eq!(rewrite(command), None, "accepted {command:?}");
1009        }
1010    }
1011
1012    #[test]
1013    fn rejects_shell_control_syntax_and_expansion() {
1014        for command in [
1015            "git status | cat",
1016            "git status > result",
1017            "git status && echo done",
1018            "git status; echo done",
1019            "git status $(touch bad)",
1020            "git status `touch bad`",
1021            "git status # comment",
1022            "git status\nwhoami",
1023            "git status $HOME",
1024        ] {
1025            assert_eq!(rewrite(command), None, "accepted {command:?}");
1026        }
1027    }
1028
1029    #[test]
1030    fn classifies_only_compatible_compound_results() {
1031        let rule = select_result_rule("cd repo && cargo check; cargo build")
1032            .expect("compatible build result");
1033        assert_eq!(rule.action, yarp_rule_pack::Action::Reduce);
1034        assert!(matches!(
1035            rule.reducer,
1036            Some(yarp_rule_pack::Reducer::BuildSummary)
1037        ));
1038        assert!(select_result_rule("cargo test && cargo build").is_err());
1039        assert!(matches!(
1040            select_result_rule("GOWORK=off go test ./... 2>&1 || true")
1041                .expect("assignment and stream merge"),
1042            Rule {
1043                reducer: Some(yarp_rule_pack::Reducer::TestSummary),
1044                ..
1045            }
1046        ));
1047        assert!(matches!(
1048            select_result_rule("env CI=1 pnpm build && :").expect("environment wrapper"),
1049            Rule {
1050                reducer: Some(yarp_rule_pack::Reducer::BuildSummary),
1051                ..
1052            }
1053        ));
1054        assert!(matches!(
1055            select_result_rule("cargo test | cat").expect("line-preserving pipeline"),
1056            Rule {
1057                reducer: Some(yarp_rule_pack::Reducer::TestSummary),
1058                ..
1059            }
1060        ));
1061        assert!(select_result_rule("cargo test > result.log").is_err());
1062        assert!(select_result_rule("cargo test && echo done").is_err());
1063        assert!(select_result_rule("set && cargo test").is_err());
1064        assert!(select_result_rule("yarp search ref error").is_err());
1065    }
1066
1067    #[test]
1068    fn plans_safe_composite_results_and_rejects_uncertain_stages() {
1069        let search = select_result_plan("rg TODO . | sort | head -50").expect("search pipeline");
1070        assert!(matches!(
1071            search.rule.reducer,
1072            Some(yarp_rule_pack::Reducer::SearchSummary)
1073        ));
1074        assert_eq!(search.status_confidence, StatusConfidence::FinalStageOnly);
1075
1076        let list = select_result_plan("find . -type f | sort | uniq").expect("list pipeline");
1077        assert!(matches!(
1078            list.rule.reducer,
1079            Some(yarp_rule_pack::Reducer::ListSummary)
1080        ));
1081        assert_eq!(list.status_confidence, StatusConfidence::FinalStageOnly);
1082
1083        let test = select_result_plan("cargo test | tee test.log").expect("test pipeline");
1084        assert!(matches!(
1085            test.rule.reducer,
1086            Some(yarp_rule_pack::Reducer::TestSummary)
1087        ));
1088        assert_eq!(test.status_confidence, StatusConfidence::FinalStageOnly);
1089
1090        let pipefail = select_result_plan("set -o pipefail && cargo test | head -100")
1091            .expect("pipefail pipeline");
1092        assert_eq!(pipefail.status_confidence, StatusConfidence::Complete);
1093        let skipped_pipefail = select_result_plan("false && set -o pipefail; cargo test | head -1")
1094            .expect("conditionally skipped pipefail");
1095        assert_eq!(
1096            skipped_pipefail.status_confidence,
1097            StatusConfidence::FinalStageOnly
1098        );
1099        let conditional_disable =
1100            select_result_plan("set -o pipefail; false && set +o pipefail; cargo test | head -1")
1101                .expect("conditional pipefail disable");
1102        assert_eq!(
1103            conditional_disable.status_confidence,
1104            StatusConfidence::FinalStageOnly
1105        );
1106
1107        let sequence = select_result_plan("cargo test; cargo test").expect("test sequence");
1108        assert_eq!(sequence.status_confidence, StatusConfidence::FinalStageOnly);
1109        for command in [
1110            "cd missing; cargo test",
1111            "export CI=1; cargo test",
1112            "VALUE=1; cargo test",
1113            "umask invalid; cargo test",
1114        ] {
1115            let plan = select_result_plan(command).expect("fallible setup sequence");
1116            assert_eq!(
1117                plan.status_confidence,
1118                StatusConfidence::FinalStageOnly,
1119                "{command:?}"
1120            );
1121        }
1122        let conjunction = select_result_plan("cargo test && cargo test").expect("test conjunction");
1123        assert_eq!(conjunction.status_confidence, StatusConfidence::Complete);
1124        let multiline = select_result_plan("cargo test\ncargo test").expect("multiline tests");
1125        assert_eq!(
1126            multiline.status_confidence,
1127            StatusConfidence::FinalStageOnly
1128        );
1129        let masked = select_result_plan("cargo test; true").expect("masked test status");
1130        assert_eq!(masked.status_confidence, StatusConfidence::FinalStageOnly);
1131        assert!(select_result_plan("export CI=1; umask 077; cargo test").is_ok());
1132
1133        for command in [
1134            "cargo test | head source.txt",
1135            "find . -type f | sort existing.txt",
1136            "find . -type f | sort --compress-program=cat",
1137            "find . -type f | sort --compress-program cat",
1138            "ls -D",
1139            "git tag --list --format=%(refname)",
1140            "git stash list --pretty=format:%H",
1141            "rg --json TODO . | jq .",
1142            "find . -print0 | xargs -0 echo",
1143            "cat source.rs | sed 's/x/y/'",
1144            "set -o; cargo test",
1145            "set -x; cargo test",
1146            "set -v; cargo test",
1147            "export -p; cargo test",
1148            "umask -S; cargo test",
1149            "echo \"$VALUE\" | rg x",
1150            "cargo test | unknown-filter",
1151            "cargo test |& head",
1152            "cat source.rs | head",
1153        ] {
1154            assert!(select_result_plan(command).is_err(), "planned {command:?}");
1155        }
1156    }
1157
1158    #[test]
1159    fn handles_quotes_and_escapes_conservatively() {
1160        assert!(rewrite("git diff file\\ name").is_some());
1161        assert!(rewrite("'git' \"status\"").is_some());
1162        assert!(rewrite("git diff '$HOME'").is_some());
1163        assert_eq!(rewrite("git diff \"$HOME\""), None);
1164        assert_eq!(rewrite("git diff 'unterminated"), None);
1165        assert_eq!(rewrite("git diff trailing\\"), None);
1166    }
1167
1168    #[test]
1169    fn validates_parsed_child_arguments() {
1170        assert!(matches!(
1171            select_builtin_argv(&strings(&["cargo", "test", "--workspace"])),
1172            Ok(Selection::Reduce(_))
1173        ));
1174        assert!(matches!(
1175            select_builtin_argv(&strings(&["go", "test", "./..."])),
1176            Ok(Selection::Reduce(_))
1177        ));
1178        assert!(matches!(
1179            select_builtin_argv(&strings(&["git", "show", "HEAD:file"])),
1180            Ok(Selection::Passthrough(_))
1181        ));
1182        assert_eq!(
1183            select_builtin_argv(&strings(&["git", "push"])),
1184            Ok(Selection::Unsupported)
1185        );
1186    }
1187
1188    #[test]
1189    fn classifies_reviewed_commands_and_guards_streaming_or_structured_output() {
1190        for arguments in [
1191            &["hf", "jobs", "logs", "job-id"][..],
1192            &["hf", "spaces", "logs", "owner/space", "--build"][..],
1193        ] {
1194            let Ok(Selection::Reduce(selected)) = select_builtin_argv(&strings(arguments)) else {
1195                panic!("did not reduce {arguments:?}");
1196            };
1197            assert!(matches!(
1198                selected.rule.reducer,
1199                Some(yarp_rule_pack::Reducer::LogSummary)
1200            ));
1201        }
1202        for arguments in [
1203            &["herdr", "pane", "list"][..],
1204            &["herdr", "tab", "list", "--workspace", "w1"][..],
1205            &["herdr", "workspace", "list"][..],
1206        ] {
1207            let Ok(Selection::Reduce(selected)) = select_builtin_argv(&strings(arguments)) else {
1208                panic!("did not reduce {arguments:?}");
1209            };
1210            assert!(matches!(
1211                selected.rule.reducer,
1212                Some(yarp_rule_pack::Reducer::ListSummary)
1213            ));
1214        }
1215        let arguments = ["pnpm", "vitest", "run", "src/example.test.ts"];
1216        let Ok(Selection::Reduce(selected)) = select_builtin_argv(&strings(&arguments)) else {
1217            panic!("did not reduce {arguments:?}");
1218        };
1219        assert!(matches!(
1220            selected.rule.reducer,
1221            Some(yarp_rule_pack::Reducer::TestSummary)
1222        ));
1223        for arguments in [
1224            &["hf", "jobs", "logs", "job-id", "--follow"][..],
1225            &["hf", "jobs", "logs", "job-id", "-f"][..],
1226            &["hf", "jobs", "logs", "job-id", "--json"][..],
1227            &["hf", "jobs", "logs", "job-id", "--format=json"][..],
1228            &["hf", "jobs", "logs", "job-id", "--format", "quiet"][..],
1229            &["pnpm", "vitest", "run", "--watch"][..],
1230        ] {
1231            assert!(matches!(
1232                select_builtin_argv(&strings(arguments)),
1233                Ok(Selection::Passthrough(_))
1234            ));
1235        }
1236        for arguments in [&["pnpm", "vitest"][..], &["pnpm", "vitest", "watch"][..]] {
1237            assert!(!matches!(
1238                select_builtin_argv(&strings(arguments)),
1239                Ok(Selection::Reduce(_))
1240            ));
1241        }
1242    }
1243
1244    fn strings(values: &[&str]) -> Vec<String> {
1245        values.iter().map(ToString::to_string).collect()
1246    }
1247}