yarp_cli/
lib.rs

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

yarp_cli/
lib.rs
