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

yarp_cli/
lib.rs
