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#[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}