use std::ffi::OsString;
use std::fmt;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
pub const RELOAD_EXIT_CODE: u8 = 75;
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
pub const BEVY_VERSION: &str = "0.20.0";
pub mod env {
pub const HOME: &str = "RIG_HOME";
pub const SESSION: &str = "RIG_SESSION";
pub const LAUNCHER: &str = "RIG_LAUNCHER";
pub const NOTICE: &str = "RIG_NOTICE";
pub const BUILD_FAILURE: &str = "RIG_BUILD_FAILURE";
pub const RUN: &str = "RIG_RUN";
pub const START: &str = "RIG_START";
pub const AGENT_ONLY: &[&str] = &[SESSION, NOTICE, BUILD_FAILURE, RUN, START];
pub const DESCRIBE: &str = "RIG_DESCRIBE";
pub const SECRET_PREFIX: &str = "RIG_SECRET_";
pub const SOURCE: &str = "RIG_SOURCE";
}
pub fn launcher() -> Option<OsString> {
std::env::var_os(env::SESSION).filter(|session| !session.is_empty())?;
std::env::var_os(env::LAUNCHER).filter(|launcher| !launcher.is_empty())
}
pub fn secret_variable(name: &str) -> String {
format!(
"{}{}",
env::SECRET_PREFIX,
name.to_ascii_uppercase().replace('-', "_")
)
}
pub fn secret_variables() -> Vec<OsString> {
std::env::vars_os()
.map(|(name, _)| name)
.filter(|name| {
name.to_str()
.is_some_and(|name| name.starts_with(env::SECRET_PREFIX))
})
.collect()
}
pub fn first_errors<'a>(lines: impl IntoIterator<Item = &'a str>, count: usize) -> Vec<&'a str> {
lines
.into_iter()
.skip_while(|line| !line.starts_with("error"))
.take(count)
.collect()
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct SessionId(String);
impl SessionId {
pub fn generate() -> Self {
static MADE: AtomicU64 = AtomicU64::new(0);
let seconds = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|elapsed| elapsed.as_secs())
.unwrap_or_default();
let id = format!("{seconds}-{}", std::process::id());
match MADE.fetch_add(1, Ordering::Relaxed) {
0 => Self(id),
made => Self(format!("{id}-{made}")),
}
}
pub fn from_env() -> Result<Option<Self>, String> {
let id = std::env::var_os(env::SESSION).filter(|id| !id.is_empty());
id.map(|id| id.to_string_lossy().parse()).transpose()
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl FromStr for SessionId {
type Err = String;
fn from_str(id: &str) -> Result<Self, Self::Err> {
let valid = id.starts_with(|c: char| c.is_ascii_digit())
&& id.chars().all(|c| c.is_ascii_digit() || c == '-');
if valid {
Ok(Self(id.to_owned()))
} else {
Err(format!(
"`{id}` is not a session id: it must be digits and dashes"
))
}
}
}
impl fmt::Display for SessionId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum StartReason {
New,
Resumed,
Reloaded,
Restarted,
}
impl StartReason {
pub fn as_str(self) -> &'static str {
match self {
Self::New => "new",
Self::Resumed => "resumed",
Self::Reloaded => "reloaded",
Self::Restarted => "restarted",
}
}
pub fn from_env() -> Option<Self> {
std::env::var(env::START).ok()?.parse().ok()
}
}
impl FromStr for StartReason {
type Err = String;
fn from_str(text: &str) -> Result<Self, Self::Err> {
[Self::New, Self::Resumed, Self::Reloaded, Self::Restarted]
.into_iter()
.find(|reason| reason.as_str() == text)
.ok_or_else(|| format!("`{text}` is not a start reason"))
}
}
#[derive(Clone, Debug)]
pub struct Home {
root: PathBuf,
}
impl Home {
pub fn from_env() -> Self {
let root = std::env::var_os(env::HOME)
.filter(|root| !root.is_empty())
.map(PathBuf::from)
.or_else(|| std::env::home_dir().map(|home| home.join(".rig")))
.unwrap_or_else(|| PathBuf::from(".rig"));
Self::at(root)
}
pub fn at(root: impl Into<PathBuf>) -> Self {
let root = root.into();
let root = std::path::absolute(&root).unwrap_or(root);
Self { root }
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn plugins_toml(&self) -> PathBuf {
self.root.join("plugins.toml")
}
pub fn plugins(&self) -> PathBuf {
self.root.join("plugins")
}
pub fn settings(&self) -> PathBuf {
self.root.join("settings.toml")
}
pub fn defaults(&self) -> PathBuf {
self.root.join("defaults.json")
}
pub fn auth(&self) -> PathBuf {
self.root.join("auth")
}
pub fn project(&self) -> PathBuf {
self.root.join("project")
}
pub fn build_log(&self) -> PathBuf {
self.root.join("build.log")
}
pub fn target(&self) -> PathBuf {
self.root.join("target")
}
pub fn bin(&self) -> PathBuf {
self.root.join("bin")
}
pub fn good(&self) -> PathBuf {
self.bin().join("good")
}
pub fn staged(&self) -> PathBuf {
self.bin().join("staged")
}
pub fn described(&self, stamp: &str) -> PathBuf {
self.bin()
.join("described")
.join(stamp.trim().replace(' ', "-"))
}
pub fn rejected(&self) -> PathBuf {
self.bin().join("rejected")
}
pub fn build_lock(&self) -> PathBuf {
self.bin().join("lock")
}
pub fn staged_for(&self, session: &SessionId) -> PathBuf {
self.bin().join(format!("staged-{session}"))
}
pub fn trial_for(&self, session: &SessionId) -> PathBuf {
self.bin().join(format!("trial-{session}"))
}
pub fn owner_of(name: &str) -> Option<SessionId> {
name.strip_prefix("trial-")
.or_else(|| name.strip_prefix("staged-"))?
.parse()
.ok()
}
pub fn data(&self) -> PathBuf {
self.root.join("data")
}
pub fn projects(&self) -> PathBuf {
self.root.join("projects")
}
pub fn runs(&self) -> PathBuf {
self.root.join("run")
}
pub fn run(&self, launcher: u32) -> RunDir {
RunDir(self.runs().join(launcher.to_string()))
}
pub fn sessions(&self) -> PathBuf {
self.root.join("sessions")
}
pub fn session(&self, session: &SessionId) -> SessionDir {
SessionDir(self.sessions().join(session.as_str()))
}
pub fn new_session(&self) -> std::io::Result<(SessionId, SessionDir)> {
std::fs::create_dir_all(self.sessions())?;
let base = SessionId::generate();
let mut id = base.clone();
let mut taken = 0u64;
loop {
let dir = self.session(&id);
match std::fs::create_dir(dir.path()) {
Ok(()) => return Ok((id, dir)),
Err(failure) if failure.kind() == std::io::ErrorKind::AlreadyExists => {
taken += 1;
id = SessionId(format!("{base}-{taken}"));
}
Err(failure) => return Err(failure),
}
}
}
pub fn saved_sessions(&self) -> Vec<(SessionId, SessionDir)> {
let entries = std::fs::read_dir(self.sessions()).into_iter().flatten();
let ids = entries.filter_map(|entry| entry.ok()?.file_name().to_str()?.parse().ok());
ids.map(|id| {
let dir = self.session(&id);
(id, dir)
})
.filter(|(_, dir)| dir.is_saved())
.collect()
}
pub fn find(&self, name_or_id: &str) -> Result<SessionId, String> {
let wanted = name_or_id.trim();
if let Ok(id) = wanted.parse::<SessionId>()
&& self.session(&id).is_saved()
{
return Ok(id);
}
let mut named: Vec<SessionId> = self
.saved_sessions()
.into_iter()
.filter(|(_, dir)| dir.name().as_deref() == Some(wanted))
.map(|(id, _)| id)
.collect();
named.sort_by(|a, b| a.as_str().cmp(b.as_str()));
match named.as_slice() {
[id] => Ok(id.clone()),
[] => Err(format!("`{wanted}` names no saved session")),
several => Err(format!(
"`{wanted}` names {} saved sessions: {}; resume one by its id",
several.len(),
several
.iter()
.map(SessionId::as_str)
.collect::<Vec<_>>()
.join(", ")
)),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RunDir(PathBuf);
impl RunDir {
pub fn from_env() -> Option<Self> {
let path = std::env::var_os(env::RUN).filter(|path| !path.is_empty())?;
Some(Self(PathBuf::from(path)))
}
pub fn path(&self) -> &Path {
&self.0
}
pub fn ready(&self) -> PathBuf {
self.0.join("ready")
}
pub fn log(&self) -> PathBuf {
self.0.join("agent.log")
}
pub fn refusal(&self) -> PathBuf {
self.0.join("refusal")
}
pub fn lock(&self) -> PathBuf {
self.0.join("lock")
}
}
#[derive(Clone, Debug)]
pub struct SessionDir(PathBuf);
impl SessionDir {
pub fn at(path: impl Into<PathBuf>) -> Self {
Self(path.into())
}
pub fn path(&self) -> &Path {
&self.0
}
pub fn id(&self) -> Option<SessionId> {
self.0.file_name()?.to_str()?.parse().ok()
}
pub fn is_saved(&self) -> bool {
self.saved().exists() && !self.is_deleting()
}
pub fn is_deleting(&self) -> bool {
self.deleting().exists()
}
pub fn mark_deleting(&self) -> std::io::Result<()> {
std::fs::write(self.deleting(), b"")
}
fn deleting(&self) -> PathBuf {
self.0.join("deleting")
}
pub fn mark_saved(&self) -> std::io::Result<()> {
std::fs::write(self.saved(), b"")
}
fn saved(&self) -> PathBuf {
self.0.join("saved")
}
pub fn started(&self) -> PathBuf {
self.0.join("started")
}
pub fn log(&self) -> PathBuf {
self.0.join("agent.log")
}
pub fn launcher_lock(&self) -> PathBuf {
self.0.join("launcher.lock")
}
pub fn hold(&self) -> std::io::Result<Option<std::fs::File>> {
std::fs::create_dir_all(&self.0)?;
let file = lock_file(&self.launcher_lock(), true)?;
match file.try_lock() {
Ok(()) => Ok(Some(file)),
Err(std::fs::TryLockError::WouldBlock) => Ok(None),
Err(std::fs::TryLockError::Error(failure)) => Err(failure),
}
}
pub fn runs_elsewhere(&self) -> bool {
!is_free(&self.launcher_lock())
}
pub fn directory(&self) -> PathBuf {
self.0.join("directory")
}
pub fn name(&self) -> Option<String> {
let text = std::fs::read_to_string(self.name_file()).ok()?;
let name = text.trim();
(!name.is_empty()).then(|| name.to_owned())
}
pub fn set_name(&self, name: Option<&str>) -> std::io::Result<()> {
let file = self.name_file();
match name.map(str::trim).filter(|name| !name.is_empty()) {
Some(name) => {
let part = self.0.join("name.part");
std::fs::write(&part, name)?;
std::fs::rename(part, file)
}
None => match std::fs::remove_file(file) {
Err(failure) if failure.kind() == std::io::ErrorKind::NotFound => Ok(()),
removed => removed,
},
}
}
fn name_file(&self) -> PathBuf {
self.0.join("name")
}
pub fn switch(&self) -> PathBuf {
self.0.join("switch")
}
pub fn working_directory(&self) -> Option<PathBuf> {
let text = std::fs::read_to_string(self.directory()).ok()?;
let directory = PathBuf::from(text.trim_end_matches('\n'));
directory.is_dir().then_some(directory)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct Invocation {
pub print: Option<String>,
pub model: Option<String>,
pub flags: Vec<(String, Option<String>)>,
pub command: Option<String>,
pub arguments: Vec<String>,
}
pub const INVOCATION_USAGE: &str =
" -p, --print [prompt…] Answer one prompt and exit: the answer goes to stdout.
Text piped in on stdin follows the prompt. With
--json, the run as JSON lines; --json alone reads
requests on stdin (`rig help-agents`).
-m, --model <model> The catalog model (vendor/model) to use.
--<option> [<value>] An option of one of the agent's plugins, also written
--<option>=<value>.
";
impl Invocation {
pub fn parse<S: AsRef<str>>(args: &[S]) -> Result<Self, String> {
match read_args(args, |flag, _| plugin_flag(flag))? {
(invocation, None) => Ok(invocation),
(_, Some(_)) => Err("`--resume` is the launcher's".to_owned()),
}
}
pub fn from_env() -> Result<Self, String> {
let args: Vec<String> = std::env::args().skip(1).collect();
Self::parse(&args)
}
pub fn is_headless(&self) -> bool {
self.print.is_some()
}
pub fn to_args(&self) -> Vec<OsString> {
let mut args: Vec<OsString> = Vec::new();
if let Some(model) = &self.model {
args.extend(["--model".into(), model.into()]);
}
for (name, value) in &self.flags {
args.push(match value {
Some(value) => format!("--{name}={value}").into(),
None => format!("--{name}").into(),
});
}
if let Some(prompt) = &self.print {
args.push("--print".into());
if !prompt.is_empty() {
args.extend(["--".into(), prompt.into()]);
}
} else if let Some(command) = &self.command {
args.push(command.into());
args.extend(self.arguments.iter().map(OsString::from));
}
args
}
}
fn read_args<'a, S: AsRef<str>>(
args: &'a [S],
mut option: impl FnMut(
&'a str,
&mut dyn Iterator<Item = &'a str>,
) -> Result<(String, Option<String>), String>,
) -> Result<(Invocation, Option<String>), String> {
let mut print = false;
let mut model = None;
let mut resume = None;
let mut flags = Vec::new();
let mut words: Vec<&str> = Vec::new();
let mut args = args.iter().map(AsRef::as_ref);
while let Some(arg) = args.next() {
match arg {
"-p" | "--print" => print = true,
"-m" | "--model" => {
let name = args.next().ok_or("--model needs a vendor/model")?;
model = Some(name.to_owned());
}
"--" => words.extend(args.by_ref()),
flag if flag.starts_with("--model=") => {
model = flag.strip_prefix("--model=").map(str::to_owned);
}
"-r" | "--resume" => {
let session = args.next().ok_or("--resume needs a session id or name")?;
resume = Some(session.to_owned());
}
flag if flag.starts_with("--resume=") => {
resume = flag.strip_prefix("--resume=").map(str::to_owned);
}
flag if flag.starts_with("--") => flags.push(option(flag, &mut args)?),
flag if flag.starts_with('-') && flag.len() > 1 => {
return Err(format!("unknown option `{flag}`"));
}
word if !print => {
words.push(word);
words.extend(args.by_ref());
}
word => words.push(word),
}
}
let mut words = words.into_iter().map(str::to_owned);
let (print, command) = if print {
(Some(words.by_ref().collect::<Vec<_>>().join(" ")), None)
} else {
(None, words.next())
};
let invocation = Invocation {
print,
model,
flags,
command,
arguments: words.collect(),
};
Ok((invocation, resume))
}
fn plugin_flag(flag: &str) -> Result<(String, Option<String>), String> {
let option = flag.strip_prefix("--").unwrap_or(flag);
let (name, value) = match option.split_once('=') {
Some((name, value)) => (name, Some(value.to_owned())),
None => (option, None),
};
if !is_name(name) {
return Err(format!("unknown option `{flag}`"));
}
Ok((name.to_owned(), value))
}
pub fn lock_file(path: &Path, create: bool) -> std::io::Result<std::fs::File> {
std::fs::File::options()
.create(create)
.write(true)
.truncate(false)
.open(path)
}
pub fn is_free(path: &Path) -> bool {
match lock_file(path, false) {
Ok(lock) => lock.try_lock().is_ok(),
Err(failure) => failure.kind() == std::io::ErrorKind::NotFound,
}
}
pub fn is_name(name: &str) -> bool {
name.starts_with(|c: char| c.is_ascii_lowercase())
&& name
.chars()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
}
pub const RESERVED_COMMANDS: &[&str] = &["build", "plugin", "help", "help-agents"];
pub const ESSENTIAL_PLUGINS: &[(&str, &str)] = &[
(
"rig_tui::TuiPlugin",
"the terminal view, leaving only -p and piped prompts",
),
(
"rig_basics::reload::plugins::PluginsPlugin",
"/plugins itself",
),
(
"rig_basics::reload::ReloadPlugin",
"/reload, so a change applies only when rig starts again",
),
];
pub fn essential(type_path: &str) -> Option<&'static str> {
ESSENTIAL_PLUGINS
.iter()
.find(|(essential, _)| *essential == type_path)
.map(|(_, what)| *what)
}
const DESCRIBED_HEADER: &str = "rig-describe 1";
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct AgentDescription {
pub commands: Vec<DescribedCommand>,
pub flags: Vec<DescribedFlag>,
pub secrets: Vec<DescribedSecret>,
pub conflicts: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DescribedCommand {
pub name: String,
pub usage: String,
pub terminal: bool,
pub sessionless: bool,
pub help: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DescribedFlag {
pub name: String,
pub value: Option<String>,
pub help: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DescribedSecret {
pub name: String,
pub help: String,
}
impl AgentDescription {
pub fn parse(text: &str) -> Self {
let mut lines = text.lines();
let mut described = Self::default();
if lines.next() != Some(DESCRIBED_HEADER) {
return described;
}
let given = |field: &str| (field != "-").then(|| field.to_owned());
for line in lines {
match line.split('\t').collect::<Vec<_>>().as_slice() {
["command", name, usage, run, session, help, ..] => {
described.commands.push(DescribedCommand {
name: (*name).to_owned(),
usage: given(usage).unwrap_or_default(),
terminal: *run == "terminal",
sessionless: *session == "sessionless",
help: (*help).to_owned(),
});
}
["flag", name, value, help, ..] => described.flags.push(DescribedFlag {
name: (*name).to_owned(),
value: given(value),
help: (*help).to_owned(),
}),
["secret", name, help, ..] => described.secrets.push(DescribedSecret {
name: (*name).to_owned(),
help: (*help).to_owned(),
}),
["conflict", conflict, ..] => described.conflicts.push((*conflict).to_owned()),
_ => {}
}
}
described
}
pub fn write(&self) -> String {
let clean = |field: &str| field.replace(['\t', '\n', '\r'], " ");
let or_dash = |field: &str| match field {
"" => "-".to_owned(),
field => clean(field),
};
let mut lines = vec![DESCRIBED_HEADER.to_owned()];
for command in &self.commands {
lines.push(format!(
"command\t{}\t{}\t{}\t{}\t{}",
clean(&command.name),
or_dash(&command.usage),
if command.terminal {
"terminal"
} else {
"headless"
},
if command.sessionless {
"sessionless"
} else {
"session"
},
clean(&command.help),
));
}
for flag in &self.flags {
lines.push(format!(
"flag\t{}\t{}\t{}",
clean(&flag.name),
or_dash(flag.value.as_deref().unwrap_or_default()),
clean(&flag.help),
));
}
for secret in &self.secrets {
lines.push(format!(
"secret\t{}\t{}",
clean(&secret.name),
clean(&secret.help)
));
}
for conflict in &self.conflicts {
lines.push(format!("conflict\t{}", clean(conflict)));
}
lines.push(String::new());
lines.join("\n")
}
pub fn command(&self, name: &str) -> Option<&DescribedCommand> {
self.commands.iter().find(|command| command.name == name)
}
pub fn invocation<S: AsRef<str>>(&self, args: &[S]) -> Result<Invocation, String> {
let invocation = read_args(args, |flag, rest| {
let (name, value) = plugin_flag(flag).map_err(|_| self.unknown_option(flag))?;
let Some(known) = self.flags.iter().find(|known| known.name == name) else {
return Err(self.unknown_option(flag));
};
match (&known.value, value) {
(Some(_), Some(value)) => Ok((name, Some(value))),
(Some(what), None) => match rest.next() {
Some(value) => Ok((name, Some(value.to_owned()))),
None => Err(format!("--{name} needs a value: --{name} {what}")),
},
(None, Some(_)) => Err(format!("--{name} takes no value")),
(None, None) => Ok((name, None)),
}
})?
.0;
if let Some(command) = &invocation.command {
if RESERVED_COMMANDS.contains(&command.as_str()) {
return Err(format!("`{command}` is the launcher's; put it first"));
}
if self.command(command).is_none() {
return Err(format!(
"unknown command `{command}`; commands: {}",
self.command_names().join(", ")
));
}
}
Ok(invocation)
}
pub fn resume<S: AsRef<str>>(&self, args: &[S]) -> Option<String> {
let read = read_args(args, |flag, rest| {
let (name, value) = plugin_flag(flag)?;
let spaced = value.is_none()
&& self
.flags
.iter()
.any(|known| known.name == name && known.value.is_some());
Ok((
name,
if spaced {
rest.next().map(str::to_owned)
} else {
value
},
))
});
read.ok().and_then(|(_, resume)| resume)
}
pub fn usage(&self) -> String {
let mut text = String::from("\nPlugin commands (of the last build):\n");
if self.commands.is_empty() {
text.push_str(" (none)\n");
}
for command in &self.commands {
let usage = format!("rig {} {}", command.name, command.usage);
text.push_str(&help_entry(usage.trim_end(), &command.help));
}
text.push_str("\nPlugin options (of the last build):\n");
if self.flags.is_empty() {
text.push_str(" (none)\n");
}
for flag in &self.flags {
let option = match &flag.value {
Some(value) => format!("--{} {value}", flag.name),
None => format!("--{}", flag.name),
};
text.push_str(&help_entry(&option, &flag.help));
}
if !self.secrets.is_empty() {
text.push_str("\nPlugin secrets (of the last build), asked for when not set:\n");
}
for secret in &self.secrets {
let variable = secret_variable(&secret.name);
text.push_str(&help_entry(&variable, &secret.help));
}
text
}
pub fn command_names(&self) -> Vec<&str> {
RESERVED_COMMANDS
.iter()
.copied()
.chain(self.commands.iter().map(|command| command.name.as_str()))
.collect()
}
fn unknown_option(&self, flag: &str) -> String {
let options: Vec<String> = ["-p", "-m"]
.into_iter()
.map(str::to_owned)
.chain(self.flags.iter().map(|flag| match &flag.value {
Some(value) => format!("--{} {value}", flag.name),
None => format!("--{}", flag.name),
}))
.collect();
format!("unknown option {flag}; options: {}", options.join(", "))
}
}
fn help_entry(name: &str, help: &str) -> String {
const COLUMN: usize = 25;
let name = format!(" {name}");
if name.chars().count() + 2 <= COLUMN {
format!("{name:<COLUMN$}{help}\n")
} else {
format!("{name}\n{:COLUMN$}{help}\n", "")
}
}
#[cfg(test)]
mod tests;