rig_harness/host/cli.rs
1//! What the command line chooses beyond the launcher's own options: the
2//! options plugins take ([`Flag`]), the front of the run ([`Front`]) and
3//! the commands that choose one ([`Front::on_command`]), and the secrets
4//! plugins ask for ([`Secret`]).
5//!
6//! A plugin declares its options in `build` with `app.add_flag(..)` and
7//! reads them from [`Invoked`]: `rig --json` reaches the agent as the flag
8//! `json`, `rig --serve 127.0.0.1:80` as `serve` with its value. The agent
9//! refuses an option no plugin of its build declared. Each declaration is
10//! an entity (`front:json`, `flag:serve`, `secret:crypt`), so it is marked
11//! [`ProvidedBy`](crate::ProvidedBy) its plugin, and another plugin can
12//! turn it off with `Disabled`. The launcher reads the command line
13//! against what the build declared, which the build describes when asked
14//! ([`describe`](crate::describe)).
15//!
16//! A front shows the agents to someone and takes what they type, such as the
17//! terminal view or `--print`. One front takes each run: the one whose
18//! command or flag was given, else the first offered (in `plugins.toml`
19//! order) for the kind of run. Each sets itself up in `finish` only if
20//! `app.runs_front(name)`. The plugin guide's
21//! [fronts section](crate::plugin_guide#fronts-options-and-commands) covers
22//! sessionless fronts and fronts with a store of their own.
23//!
24//! ```
25//! use rig_harness::prelude::*;
26//!
27//! #[derive(Default)]
28//! struct JsonFront;
29//!
30//! impl Plugin for JsonFront {
31//! fn build(&self, app: &mut App) {
32//! app.add_front(Front::new("json").on_flag("json", "Stream the session as JSON lines"))
33//! // `rig stats [--all]`: not `on_terminal`, so it runs headless
34//! // and its output and exit code pass through the launcher.
35//! .add_front(Front::new("stats").sessionless().on_command(
36//! "stats",
37//! "[--all]",
38//! "Count the saved sessions",
39//! ));
40//! }
41//!
42//! fn finish(&self, app: &mut App) {
43//! if app.runs_front("json") {
44//! // Own stdout and stdin from here.
45//! }
46//! if app.runs_front("stats") {
47//! let words = app.world().get_resource::<Invoked>().and_then(Invoked::command);
48//! let all = words.is_some_and(|(_, words)| words.iter().any(|word| word == "--all"));
49//! // Add a `Startup` system that prints, then triggers `Exit`.
50//! # drop(all);
51//! }
52//! }
53//! }
54//! ```
55
56use std::fmt::Write as _;
57use std::sync::Arc;
58
59use bevy_app::prelude::*;
60use bevy_ecs::prelude::*;
61use rig_harness_protocol::{RESERVED_COMMANDS, SessionDir, StartReason};
62use rig_memory::{SessionError, Store};
63
64use super::launcher::Invoked;
65use super::session::{AppStart, StoreOpener};
66use crate::describe::Describing;
67use crate::secrets::Secret;
68
69/// An option a plugin takes on the command line: `--name`, or with a
70/// value, `--name VALUE` (also written `--name=VALUE`).
71#[derive(Component, Clone, Debug, PartialEq, Eq)]
72pub struct Flag {
73 /// Its name: lowercase letters, digits and dashes.
74 pub name: &'static str,
75 /// What its value is, such as `ADDR`, when it takes one.
76 pub value: Option<&'static str>,
77 /// What it does, for `rig help`.
78 pub help: &'static str,
79}
80
81impl Flag {
82 /// `--name`, which takes no value.
83 pub fn new(name: &'static str, help: &'static str) -> Self {
84 Self {
85 name,
86 value: None,
87 help,
88 }
89 }
90
91 /// The same option, taking a value described as `value`.
92 pub fn value(mut self, value: &'static str) -> Self {
93 self.value = Some(value);
94 self
95 }
96
97 /// How a refusal lists it.
98 fn usage(&self) -> String {
99 let option = match self.value {
100 Some(value) => format!("--{} {value}", self.name),
101 None => format!("--{}", self.name),
102 };
103 format!(" {option:<23}{}", self.help)
104 }
105}
106
107/// A top-level command, `rig <name> <usage>`, that gives its front the run.
108#[derive(Clone, Debug, PartialEq, Eq)]
109pub struct Command {
110 /// Its name: lowercase letters, digits and dashes, and none of the
111 /// launcher's own (`build`, `plugin`, `help`).
112 pub name: &'static str,
113 /// What its arguments are, such as `<tool> [path]`; may be empty.
114 pub usage: &'static str,
115 /// What it does, for `rig help`.
116 pub help: &'static str,
117}
118
119/// A front a plugin offers, and which runs it can take.
120#[derive(Component, Clone, Debug)]
121pub struct Front {
122 /// Its name, which [`AppCliExt::runs_front`] asks about.
123 pub name: &'static str,
124 /// Whether it takes a run on an interactive terminal. A command's
125 /// front that does not runs headless: the launcher passes its output
126 /// and exit code through and never restarts it.
127 pub terminal: bool,
128 /// Whether it takes a run without one (`--print`, stdin piped in).
129 /// This and [`terminal`](Self::terminal) only pick the front of a run
130 /// that names no command or option: a front with either takes only
131 /// the runs that name it.
132 pub headless: bool,
133 /// The option that gives it the run, whatever the terminal.
134 pub flag: Option<Flag>,
135 /// The command that gives it the run, whatever the terminal.
136 pub command: Option<Command>,
137 /// Whether its runs keep no session: a run it takes removes its
138 /// session's directory when it ends, and a resumed session's run is
139 /// refused.
140 pub sessionless: bool,
141 /// Where a run it takes keeps its session, in place of
142 /// [`StoreOpener`] ([`Self::session_store`]).
143 pub session_store: Option<StoreOpener>,
144}
145
146impl Front {
147 /// A front that takes no run until told which.
148 pub fn new(name: &'static str) -> Self {
149 Self {
150 name,
151 terminal: false,
152 headless: false,
153 flag: None,
154 command: None,
155 sessionless: false,
156 session_store: None,
157 }
158 }
159
160 /// It takes a run on an interactive terminal.
161 pub fn on_terminal(mut self) -> Self {
162 self.terminal = true;
163 self
164 }
165
166 /// It takes a run without an interactive terminal.
167 pub fn headless(mut self) -> Self {
168 self.headless = true;
169 self
170 }
171
172 /// Its runs keep no session.
173 pub fn sessionless(mut self) -> Self {
174 self.sessionless = true;
175 self
176 }
177
178 /// A run it takes keeps its session in the store `open` opens for the
179 /// session's directory, such as one of a chat thread's, instead of the
180 /// one [`StoreOpener`] opens. Only the running session's store comes
181 /// from it: imports and listings still open theirs with
182 /// [`StoreOpener`]. The front is chosen, and the store opened, once
183 /// every plugin is built.
184 pub fn session_store(
185 mut self,
186 open: impl Fn(&SessionDir) -> Result<Arc<dyn Store>, SessionError> + Send + Sync + 'static,
187 ) -> Self {
188 self.session_store = Some(StoreOpener::new(open));
189 self
190 }
191
192 /// It takes any run with `--flag`, an option it declares with `help`.
193 pub fn on_flag(mut self, flag: &'static str, help: &'static str) -> Self {
194 self.flag = Some(Flag::new(flag, help));
195 self
196 }
197
198 /// It takes any run of `rig <name> …`, a command whose arguments are
199 /// `usage` and which does what `help` says; the words after the name
200 /// are read with [`Invoked::command`]. Two plugins with one command,
201 /// or a command of the launcher's own name, fail the build.
202 pub fn on_command(
203 mut self,
204 name: &'static str,
205 usage: &'static str,
206 help: &'static str,
207 ) -> Self {
208 self.command = Some(Command { name, usage, help });
209 self
210 }
211}
212
213/// The front that takes this run, once chosen, and the declarations of
214/// options, fronts and secrets in the order they were made.
215#[derive(Resource, Default, Debug)]
216pub struct Cli {
217 declared: Vec<Entity>,
218 /// The front of this run, once chosen; why none could be, if so.
219 chosen: Option<Result<Option<Front>, String>>,
220}
221
222impl Cli {
223 /// The front of this run, once chosen.
224 pub fn front(&self) -> Option<&'static str> {
225 self.chosen().map(|front| front.name)
226 }
227
228 /// Whether the front of this run keeps no session.
229 pub fn sessionless(&self) -> bool {
230 self.chosen().is_some_and(|front| front.sessionless)
231 }
232
233 /// The front of this run, once chosen.
234 pub(crate) fn chosen(&self) -> Option<&Front> {
235 self.chosen.as_ref()?.as_ref().ok()?.as_ref()
236 }
237}
238
239/// Every declaration of a `T` not turned off, in the order made, with its
240/// entity.
241pub(crate) fn declared<T: Component + Clone>(world: &mut World) -> Vec<(Entity, T)> {
242 let order = world
243 .get_resource::<Cli>()
244 .map(|cli| cli.declared.clone())
245 .unwrap_or_default();
246 let mut query = world.query::<&T>();
247 order
248 .into_iter()
249 .filter_map(|entity| Some((entity, query.get(world, entity).ok()?.clone())))
250 .collect()
251}
252
253/// The declared options, the fronts' included, in order.
254pub(crate) fn flags(world: &mut World) -> Vec<Flag> {
255 let fronts = declared::<Front>(world);
256 declared::<Flag>(world)
257 .into_iter()
258 .map(|(_, flag)| flag)
259 .chain(fronts.into_iter().filter_map(|(_, front)| front.flag))
260 .collect()
261}
262
263/// Why the command line does not fit this build: an option no plugin
264/// declared, a value where none goes or none where one does, an unknown
265/// command, two fronts given the run, or a front that keeps no session
266/// given a resumed one.
267fn refusal_of(world: &mut World, invoked: &Invoked) -> Option<String> {
268 let flags = flags(world);
269 for (name, value) in &invoked.args.flags {
270 let Some(flag) = flags.iter().find(|flag| flag.name == name) else {
271 let known: String = flags.iter().fold(String::new(), |mut text, flag| {
272 let _ = writeln!(text, "{}", flag.usage());
273 text
274 });
275 return Some(if known.is_empty() {
276 format!("unknown option `--{name}`: no plugin of this build takes options")
277 } else {
278 format!("unknown option `--{name}`; this build's plugins take:\n{known}")
279 });
280 };
281 match (flag.value, value) {
282 (Some(value_name), None) => {
283 return Some(format!("`--{name}` needs a value: --{name} {value_name}"));
284 }
285 (None, Some(_)) => return Some(format!("`--{name}` takes no value")),
286 _ => {}
287 }
288 }
289 let resumed = world
290 .get_resource::<AppStart>()
291 .is_some_and(|start| start.0 == StartReason::Resumed);
292 match world
293 .get_resource::<Cli>()
294 .and_then(|cli| cli.chosen.as_ref())
295 {
296 Some(Err(why)) => Some(why.clone()),
297 Some(Ok(Some(front))) if front.sessionless && resumed => Some(format!(
298 "the front `{}` keeps no session, so it cannot resume one",
299 front.name
300 )),
301 _ => None,
302 }
303}
304
305/// Chooses the front of the run, once: none while the build describes
306/// itself.
307pub(crate) fn choose(world: &mut World) -> Option<&'static str> {
308 if world.get_resource_or_init::<Cli>().chosen.is_none() {
309 let chosen = if world.contains_resource::<Describing>() {
310 Ok(None)
311 } else {
312 let invoked = world.get_resource::<Invoked>().cloned().unwrap_or_default();
313 let fronts: Vec<Front> = declared::<Front>(world)
314 .into_iter()
315 .map(|(_, front)| front)
316 .collect();
317 front_of(&fronts, &invoked)
318 };
319 world.resource_mut::<Cli>().chosen = Some(chosen);
320 }
321 world.resource::<Cli>().front()
322}
323
324/// The front of `fronts` that takes the run `invoked` describes.
325fn front_of(fronts: &[Front], invoked: &Invoked) -> Result<Option<Front>, String> {
326 let command = invoked.command().map(|(name, _)| name);
327 let given: Vec<&Front> = fronts
328 .iter()
329 .filter(|front| {
330 let flagged = front
331 .flag
332 .as_ref()
333 .is_some_and(|flag| invoked.flag(flag.name));
334 let commanded = front
335 .command
336 .as_ref()
337 .is_some_and(|given| Some(given.name) == command);
338 flagged || commanded
339 })
340 .collect();
341 let interactive = invoked.interactive();
342 match given.as_slice() {
343 [front] => Ok(Some((*front).clone())),
344 [] => match command {
345 Some(command) => Err(format!(
346 "unknown command `{command}`: no plugin of this build adds it"
347 )),
348 None => Ok(fronts
349 .iter()
350 .filter(|front| front.command.is_none() && front.flag.is_none())
351 .find(|front| {
352 if interactive {
353 front.terminal
354 } else {
355 front.headless
356 }
357 })
358 .cloned()),
359 },
360 several => {
361 let names: Vec<String> = several
362 .iter()
363 .filter_map(|front| match (&front.command, &front.flag) {
364 (Some(given), _) if Some(given.name) == command => {
365 Some(format!("the command `{}`", given.name))
366 }
367 (_, Some(flag)) => Some(format!("--{}", flag.name)),
368 _ => None,
369 })
370 .collect();
371 Err(format!(
372 "{} each choose the front; give one",
373 names.join(" and ")
374 ))
375 }
376 }
377}
378
379/// Declaring options, fronts and secrets while the app is built.
380pub trait AppCliExt {
381 /// Declares an option of this plugin.
382 fn add_flag(&mut self, flag: Flag) -> &mut Self;
383
384 /// Offers a front, and declares its option and command.
385 fn add_front(&mut self, front: Front) -> &mut Self;
386
387 /// Asks for a secret, which the launcher gets before the agent starts
388 /// (see [`Secrets`](crate::secrets::Secrets)).
389 fn add_secret(&mut self, secret: Secret) -> &mut Self;
390
391 /// Whether the front `name` takes this run. Ask in `finish` or
392 /// `cleanup`, once every plugin offered its fronts: the first answer
393 /// chooses the run's front.
394 fn runs_front(&mut self, name: &str) -> bool;
395}
396
397/// Spawns the declaration `what` as an entity named `name`.
398fn declare(app: &mut App, name: String, what: impl Bundle) {
399 let world = app.world_mut();
400 let entity = world.spawn((Name::new(name), what)).id();
401 world.get_resource_or_init::<Cli>().declared.push(entity);
402}
403
404impl AppCliExt for App {
405 fn add_flag(&mut self, flag: Flag) -> &mut Self {
406 declare(self, format!("flag:{}", flag.name), flag);
407 self
408 }
409
410 fn add_front(&mut self, front: Front) -> &mut Self {
411 declare(self, format!("front:{}", front.name), front);
412 self
413 }
414
415 fn add_secret(&mut self, secret: Secret) -> &mut Self {
416 declare(self, format!("secret:{}", secret.name), secret);
417 self
418 }
419
420 fn runs_front(&mut self, name: &str) -> bool {
421 choose(self.world_mut()) == Some(name)
422 }
423}
424
425impl Invoked {
426 /// Whether the command line has the plugin option `--name`.
427 pub fn flag(&self, name: &str) -> bool {
428 self.args.flags.iter().any(|(flag, _)| flag == name)
429 }
430
431 /// The value of the plugin option `--name VALUE`, if given.
432 pub fn flag_value(&self, name: &str) -> Option<&str> {
433 self.args
434 .flags
435 .iter()
436 .find(|(flag, _)| flag == name)
437 .and_then(|(_, value)| value.as_deref())
438 }
439
440 /// The command `rig <command> …` runs and the words after it, if the
441 /// run is a command's.
442 pub fn command(&self) -> Option<(&str, &[String])> {
443 let command = self.args.command.as_deref()?;
444 Some((command, &self.args.arguments))
445 }
446}
447
448/// Why `command`, a front's, cannot be a command: a launcher's word, or a
449/// name the launcher cannot read.
450pub(crate) fn command_refusal(command: &Command) -> Option<String> {
451 if RESERVED_COMMANDS.contains(&command.name) {
452 return Some(format!(
453 "the command `{}` is the launcher's own",
454 command.name
455 ));
456 }
457 (!rig_harness_protocol::is_name(command.name)).then(|| {
458 format!(
459 "the command `{}` is not lowercase letters, digits and dashes",
460 command.name
461 )
462 })
463}
464
465/// Chooses the front of the run, if no front asked yet, and why the run
466/// cannot go on: a command line the build does not take, or an
467/// interactive terminal no front takes. Nothing while the build describes
468/// itself.
469pub(crate) fn refusal(app: &mut App) -> Option<String> {
470 let world = app.world_mut();
471 if world.contains_resource::<Describing>() {
472 return None;
473 }
474 let invoked = world.get_resource::<Invoked>().cloned().unwrap_or_default();
475 let front = choose(world);
476 if let Some(why) = refusal_of(world, &invoked) {
477 return Some(why);
478 }
479 (invoked.interactive() && front.is_none()).then(|| {
480 "this build has no front for an interactive terminal: enable rig-tui \
481 (`rig plugin add rig_tui::TuiPlugin --crate rig-tui`), or run with -p or with stdin \
482 piped in"
483 .to_owned()
484 })
485}
486
487#[cfg(test)]
488mod tests;