Skip to main content

yarp_cli/
rewrite.rs

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