Skip to main content

yarp_cli/
lib.rs

1#![forbid(unsafe_code)]
2
3pub mod agent_skill;
4pub mod archive;
5pub mod archive_query;
6pub mod config;
7mod config_cli;
8pub mod reducers;
9mod result_reducer;
10pub mod rewrite;
11pub mod rule_pack;
12pub mod rules;
13mod rules_cli;
14pub mod runner;
15pub mod shell;
16
17use rewrite::ArchiveCommandRef;
18use std::io;
19use std::io::Write as _;
20use time::OffsetDateTime;
21use time::format_description::well_known::Rfc3339;
22
23const HELP: &str = "YARP prunes developer command output and archives Pi tool calls.\n\nUsage:\n  yarp plan --json [--rule-pack <path>]... <shell-command>\n  yarp rewrite [--rule-pack <path>]... <shell-command>\n  yarp run [--rule-pack <path>]... -- <command> [arguments...]\n  yarp config <path|init|show|get|set|unset|check>\n  yarp rules check <source-pack>\n  yarp rules compile <source-pack> --output <compiled-pack>\n  yarp rules verify <compiled-pack>\n  yarp rules list [--rule-pack <path>]... [--json]\n  yarp rules explain [--rule-pack <path>]... [--json] -- <command> [arguments...]\n  yarp search REF PATTERN [options]\n  yarp read REF [SOURCE] START:END\n  yarp read REF SOURCE --bytes START:END\n  yarp archive stats\n  yarp archive verify\n  yarp archive prune --before <UTC timestamp>\n  yarp archive ingest\n  yarp --skill list\n  yarp --skill show yarp\n  yarp --skill export yarp\n  yarp --help\n  yarp --version\n";
24
25/// Run the command-line interface and return the process exit code.
26#[must_use]
27pub fn run_cli(arguments: &[String]) -> i32 {
28    match arguments {
29        [] => {
30            print!("{HELP}");
31            0
32        }
33        [argument] if argument == "--help" || argument == "-h" => {
34            print!("{HELP}");
35            0
36        }
37        [argument] if argument == "--version" || argument == "-V" => {
38            println!("yarp {}", env!("CARGO_PKG_VERSION"));
39            0
40        }
41        [command, format, rest @ ..] if command == "plan" && format == "--json" => {
42            match parse_rewrite(rest) {
43                Ok((packs, reference, shell_command)) => plan(shell_command, reference, &packs),
44                Err(error) => usage_error(&error),
45            }
46        }
47        [command, rest @ ..] if command == "rewrite" => match parse_rewrite(rest) {
48            Ok((packs, reference, shell_command)) => rewrite(shell_command, reference, &packs),
49            Err(error) => usage_error(&error),
50        },
51        [command, rest @ ..] if command == "run" => match parse_run(rest) {
52            Ok((key, child, packs, expected)) => {
53                match runner::run_with_rules(child, key.as_ref(), &packs, expected.as_ref()) {
54                    Ok(code) => code,
55                    Err(error) => usage_error(&error),
56                }
57            }
58            Err(error) => usage_error(&error),
59        },
60        [command, rest @ ..] if command == "config" => config_cli::run(rest),
61        [command, rest @ ..] if command == "rules" => rules_cli::run(rest),
62        [command, rest @ ..] if command == "search" => archive_search(rest),
63        [command, rest @ ..] if command == "read" => archive_read(rest),
64        [command] if command == "result-reduce" => {
65            match result_reducer::run(io::stdin().lock(), io::stdout().lock()) {
66                Ok(()) => 0,
67                Err(error) => {
68                    eprintln!("yarp: {error}");
69                    65
70                }
71            }
72        }
73        [command, subcommand] if command == "archive" && subcommand == "stats" => archive_stats(),
74        [command, subcommand] if command == "archive" && subcommand == "verify" => archive_verify(),
75        [command, subcommand] if command == "archive" && subcommand == "ingest" => {
76            match archive::run_ingest(io::stdin().lock(), io::stdout().lock()) {
77                Ok(()) => 0,
78                Err(error) => {
79                    eprintln!("yarp: {error}");
80                    74
81                }
82            }
83        }
84        [command, subcommand, rest @ ..] if command == "archive" && subcommand == "restore" => {
85            archive_restore(rest)
86        }
87        [command, subcommand, option, timestamp]
88            if command == "archive" && subcommand == "prune" && option == "--before" =>
89        {
90            archive_prune(timestamp)
91        }
92        _ => usage_error("invalid arguments"),
93    }
94}
95
96fn plan(
97    shell_command: &str,
98    reference: Option<ArchiveCommandRef<'_>>,
99    packs: &[rules::PackRequest],
100) -> i32 {
101    match rewrite::plan_with_options(shell_command, reference, packs)
102        .and_then(|plan| serde_json::to_string(&plan).map_err(|error| error.to_string()))
103    {
104        Ok(plan) => {
105            println!("{plan}");
106            0
107        }
108        Err(error) => {
109            eprintln!("yarp: {error}");
110            3
111        }
112    }
113}
114
115fn rewrite(
116    shell_command: &str,
117    reference: Option<ArchiveCommandRef<'_>>,
118    packs: &[rules::PackRequest],
119) -> i32 {
120    match rewrite::rewrite_with_options(shell_command, reference, packs) {
121        Ok(Some(rewritten)) => {
122            print!("{rewritten}");
123            0
124        }
125        Ok(None) => 3,
126        Err(error) => {
127            eprintln!("yarp: {error}");
128            3
129        }
130    }
131}
132
133fn archive_search(arguments: &[String]) -> i32 {
134    match archive_query::search(arguments) {
135        Ok(archive_query::SearchOutcome::Matches(output)) => {
136            if let Err(error) = io::stdout().write_all(&output) {
137                write_recovery_error(&format!("could not write search output: {error}"));
138                return 74;
139            }
140            0
141        }
142        Ok(archive_query::SearchOutcome::NoMatches(output)) => {
143            if let Err(error) = io::stdout().write_all(&output) {
144                write_recovery_error(&format!("could not write search output: {error}"));
145                return 74;
146            }
147            1
148        }
149        Err(error) => {
150            write_recovery_error(&error);
151            65
152        }
153    }
154}
155
156fn archive_read(arguments: &[String]) -> i32 {
157    match archive_query::read(arguments) {
158        Ok(output) => match io::stdout().write_all(&output) {
159            Ok(()) => 0,
160            Err(error) => {
161                write_recovery_error(&format!("could not write exact archive range: {error}"));
162                74
163            }
164        },
165        Err(error) => {
166            write_recovery_error(&error);
167            65
168        }
169    }
170}
171
172fn write_recovery_error(error: &str) {
173    const PREFIX_BYTES: usize = "yarp: ".len();
174    const NEWLINE_BYTES: usize = 1;
175    const ELLIPSIS: &str = "...";
176
177    let max_bytes = config::load().map_or(config::MIN_RECOVERY_CAP_BYTES, |resolved| {
178        resolved.output.recovery_cap_bytes
179    });
180    let body_budget = max_bytes.saturating_sub(PREFIX_BYTES + NEWLINE_BYTES);
181    let mut body = String::with_capacity(body_budget.min(error.len()));
182    let mut truncated = false;
183    for character in error.chars() {
184        let character = if character.is_control() {
185            ' '
186        } else {
187            character
188        };
189        if body.len().saturating_add(character.len_utf8()) > body_budget {
190            truncated = true;
191            break;
192        }
193        body.push(character);
194    }
195    if truncated && body_budget >= ELLIPSIS.len() {
196        while body.len().saturating_add(ELLIPSIS.len()) > body_budget {
197            body.pop();
198        }
199        body.push_str(ELLIPSIS);
200    }
201    eprintln!("yarp: {body}");
202}
203
204fn archive_stats() -> i32 {
205    match archive::Archive::open().and_then(|archive| archive.stats()) {
206        Ok(stats) => {
207            println!("sessions: {}", stats.sessions);
208            println!("calls: {}", stats.calls);
209            println!("incomplete_calls: {}", stats.incomplete_calls);
210            println!("logical_payload_bytes: {}", stats.logical_payload_bytes);
211            println!("stored_payload_bytes: {}", stats.stored_payload_bytes);
212            println!("database_bytes: {}", stats.database_bytes);
213            if let Some(oldest) = stats.oldest_call_ms {
214                println!("oldest_call_ms: {oldest}");
215            }
216            if let Some(newest) = stats.newest_call_ms {
217                println!("newest_call_ms: {newest}");
218            }
219            0
220        }
221        Err(error) => {
222            eprintln!("yarp: {error}");
223            74
224        }
225    }
226}
227
228fn archive_verify() -> i32 {
229    match archive::Archive::open_read_only().and_then(|archive| archive.verify()) {
230        Ok(report) => {
231            println!("incomplete_calls: {}", report.incomplete_calls);
232            if report.errors.is_empty() {
233                println!("archive: ok");
234                0
235            } else {
236                for error in report.errors {
237                    eprintln!("yarp: archive verification failed: {error}");
238                }
239                65
240            }
241        }
242        Err(error) => {
243            eprintln!("yarp: {error}");
244            65
245        }
246    }
247}
248
249fn archive_restore(arguments: &[String]) -> i32 {
250    let key = match parse_archive_key(arguments) {
251        Ok(value) => value,
252        Err(error) => return usage_error(&error),
253    };
254    match archive::Archive::open_read_only()
255        .and_then(|archive| archive.restore_streams(&key, io::stdout().lock(), io::stderr().lock()))
256    {
257        Ok(()) => 0,
258        Err(error) => {
259            eprintln!("yarp: {error}");
260            74
261        }
262    }
263}
264
265fn archive_prune(timestamp: &str) -> i32 {
266    let parsed = match OffsetDateTime::parse(timestamp, &Rfc3339) {
267        Ok(value) => value,
268        Err(error) => return usage_error(&format!("invalid UTC timestamp: {error}")),
269    };
270    if parsed.offset() != time::UtcOffset::UTC {
271        return usage_error("prune timestamp must use UTC with a Z suffix");
272    }
273    let Ok(milliseconds) = i64::try_from(parsed.unix_timestamp_nanos() / 1_000_000) else {
274        return usage_error("prune timestamp is outside the supported range");
275    };
276    match archive::Archive::open().and_then(|mut archive| archive.prune_before(milliseconds)) {
277        Ok(deleted) => {
278            println!("pruned_calls: {deleted}");
279            0
280        }
281        Err(error) => {
282            eprintln!("yarp: {error}");
283            74
284        }
285    }
286}
287
288fn parse_rewrite(
289    arguments: &[String],
290) -> Result<(Vec<rules::PackRequest>, Option<ArchiveCommandRef<'_>>, &str), String> {
291    let mut packs = rules::requests_from_config()?;
292    let mut project_root = None;
293    let mut project_pack_seen = false;
294    let mut agent = None;
295    let mut account = None;
296    let mut session = None;
297    let mut call = None;
298    let mut index = 0;
299    while index + 1 < arguments.len() {
300        let value = arguments
301            .get(index + 1)
302            .ok_or_else(|| format!("{} requires a value", arguments[index]))?;
303        match arguments[index].as_str() {
304            "--project-root" if project_root.is_none() => {
305                project_root = Some(std::path::PathBuf::from(value));
306            }
307            "--rule-pack" => {
308                let path = if let Some(root) = &project_root {
309                    project_pack_seen = true;
310                    rules::canonical_project_pack(root, std::path::Path::new(value))?
311                } else {
312                    value.into()
313                };
314                packs.push(rules::PackRequest {
315                    path,
316                    expected_digest: None,
317                    expected_compiled_digest: None,
318                });
319            }
320            "--archive-agent" if agent.is_none() => agent = Some(value.as_str()),
321            "--archive-account" if account.is_none() => account = Some(value.as_str()),
322            "--archive-session" if session.is_none() => session = Some(value.as_str()),
323            "--archive-call" if call.is_none() => call = Some(value.as_str()),
324            value if value.starts_with("--") => {
325                return Err(format!("unknown or duplicate rewrite option: {value}"));
326            }
327            _ => break,
328        }
329        index += 2;
330    }
331    let [shell_command] = &arguments[index..] else {
332        return Err("rewrite requires one shell command".to_owned());
333    };
334    if project_root.is_some() && !project_pack_seen {
335        return Err("--project-root requires a following --rule-pack".to_owned());
336    }
337    let reference = match (agent, account, session, call) {
338        (None, None, None, None) => None,
339        (Some(agent), Some(account), Some(session_id), Some(call_id)) => Some(ArchiveCommandRef {
340            agent,
341            account,
342            session_id,
343            call_id,
344        }),
345        _ => return Err("archived rewrite requires every archive identifier".to_owned()),
346    };
347    Ok((packs, reference, shell_command))
348}
349
350fn parse_archive_key(arguments: &[String]) -> Result<archive::ArchiveKey, String> {
351    let [
352        agent_flag,
353        agent,
354        account_flag,
355        account,
356        session_flag,
357        session,
358        call_flag,
359        call,
360    ] = arguments
361    else {
362        return Err("invalid archive key arguments".to_owned());
363    };
364    if agent_flag != "--archive-agent"
365        || account_flag != "--archive-account"
366        || session_flag != "--archive-session"
367        || call_flag != "--archive-call"
368    {
369        return Err("invalid archive key options".to_owned());
370    }
371    Ok(archive::ArchiveKey {
372        session: archive::SessionIdentity {
373            agent: agent.clone(),
374            account: account.clone(),
375            source_session_id: session.clone(),
376            started_at_ms: None,
377        },
378        source_call_id: call.clone(),
379    })
380}
381
382type ParsedRun<'a> = (
383    Option<archive::ArchiveKey>,
384    &'a [String],
385    Vec<rules::PackRequest>,
386    Option<runner::ExpectedSelection>,
387);
388
389fn parse_run(arguments: &[String]) -> Result<ParsedRun<'_>, String> {
390    let mut packs = Vec::new();
391    let mut selected_pack = None;
392    let mut selected_rule = None;
393    let mut selected_digest = None;
394    let mut agent = None;
395    let mut account = None;
396    let mut session = None;
397    let mut call = None;
398    let mut index = 0;
399    while index < arguments.len() && arguments[index] != "--" {
400        let option = arguments[index].as_str();
401        let value = arguments
402            .get(index + 1)
403            .ok_or_else(|| format!("{option} requires a value"))?;
404        match option {
405            "--selected-pack" if selected_pack.is_none() => selected_pack = Some(value.clone()),
406            "--selected-rule" if selected_rule.is_none() => selected_rule = Some(value.clone()),
407            "--selected-digest" if selected_digest.is_none() => {
408                selected_digest = Some(rules::parse_digest(value)?);
409            }
410            "--archive-agent" if agent.is_none() => agent = Some(value.clone()),
411            "--archive-account" if account.is_none() => account = Some(value.clone()),
412            "--archive-session" if session.is_none() => session = Some(value.clone()),
413            "--archive-call" if call.is_none() => call = Some(value.clone()),
414            "--rule-pack" => {
415                let (request, next_index) = parse_run_pack(arguments, index, value)?;
416                packs.push(request);
417                index = next_index;
418                continue;
419            }
420            value if value.starts_with("--") => {
421                return Err(format!("unknown or duplicate run option: {value}"));
422            }
423            _ => return Err(format!("invalid run option: {option}")),
424        }
425        index += 2;
426    }
427    let Some(separator) = arguments.get(index) else {
428        return Err("run requires -- and a child command".to_owned());
429    };
430    if separator != "--" || index + 1 >= arguments.len() {
431        return Err("run requires -- and a child command".to_owned());
432    }
433    let child = &arguments[index + 1..];
434    let expected = match (selected_pack, selected_rule, selected_digest) {
435        (None, None, None) => None,
436        (Some(pack_id), Some(rule_id), Some(source_digest)) => Some(runner::ExpectedSelection {
437            pack_id,
438            rule_id,
439            source_digest,
440        }),
441        _ => {
442            return Err(
443                "run requires selected pack, selected rule, and selected digest together"
444                    .to_owned(),
445            );
446        }
447    };
448    let archive_key = match (agent, account, session, call) {
449        (None, None, None, None) => None,
450        (Some(agent), Some(account), Some(session), Some(call)) => Some(archive::ArchiveKey {
451            session: archive::SessionIdentity {
452                agent,
453                account,
454                source_session_id: session,
455                started_at_ms: None,
456            },
457            source_call_id: call,
458        }),
459        _ => return Err("archived run requires every archive identifier".to_owned()),
460    };
461    if expected.is_none() {
462        let mut configured = rules::requests_from_config()?;
463        configured.extend(packs);
464        packs = configured;
465    }
466    Ok((archive_key, child, packs, expected))
467}
468
469fn parse_run_pack(
470    arguments: &[String],
471    mut index: usize,
472    path: &str,
473) -> Result<(rules::PackRequest, usize), String> {
474    let mut expected_digest = None;
475    let mut expected_compiled_digest = None;
476    index += 2;
477    if arguments
478        .get(index)
479        .is_some_and(|value| value == "--rule-pack-digest")
480    {
481        let digest = arguments
482            .get(index + 1)
483            .ok_or_else(|| "--rule-pack-digest requires a value".to_owned())?;
484        expected_digest = Some(rules::parse_digest(digest)?);
485        index += 2;
486    }
487    if arguments
488        .get(index)
489        .is_some_and(|value| value == "--rule-pack-compiled-digest")
490    {
491        let digest = arguments
492            .get(index + 1)
493            .ok_or_else(|| "--rule-pack-compiled-digest requires a value".to_owned())?;
494        expected_compiled_digest = Some(rules::parse_digest(digest)?);
495        index += 2;
496    }
497    if expected_digest.is_some() != expected_compiled_digest.is_some() {
498        return Err("rule pack source and compiled digests must be supplied together".to_owned());
499    }
500    Ok((
501        rules::PackRequest {
502            path: path.into(),
503            expected_digest,
504            expected_compiled_digest,
505        },
506        index,
507    ))
508}
509
510fn usage_error(error: &str) -> i32 {
511    eprintln!("yarp: {error}\n\n{HELP}");
512    64
513}
514
515#[cfg(test)]
516mod tests {
517    use super::*;
518
519    #[test]
520    fn unsupported_rewrite_fails_open_with_distinct_status() {
521        assert_eq!(run_cli(&["rewrite".into(), "cat .env".into()]), 3);
522    }
523
524    #[test]
525    fn invalid_arguments_return_usage_error() {
526        assert_eq!(run_cli(&["rewrite".into()]), 64);
527        assert_eq!(run_cli(&["unknown".into()]), 64);
528        assert_eq!(run_cli(&["run".into(), "--".into()]), 64);
529    }
530
531    #[test]
532    fn parses_archived_run_metadata() {
533        let arguments = vec![
534            "--archive-agent".into(),
535            "pi".into(),
536            "--archive-account".into(),
537            "onur".into(),
538            "--archive-session".into(),
539            "session".into(),
540            "--archive-call".into(),
541            "call".into(),
542            "--".into(),
543            "git".into(),
544            "status".into(),
545        ];
546        let (key, child, packs, expected) = parse_run(&arguments).expect("parse");
547        let key = key.expect("archive key");
548        assert_eq!(key.session.agent, "pi");
549        assert_eq!(key.source_call_id, "call");
550        assert_eq!(child, &["git".to_owned(), "status".to_owned()]);
551        assert!(packs.is_empty());
552        assert!(expected.is_none());
553    }
554}