diff --git a/bench/socket_client.py b/bench/socket_client.py index df000b9b8..6c499e49f 100644 --- a/bench/socket_client.py +++ b/bench/socket_client.py @@ -76,7 +76,7 @@ def close(self): class SocketClient: - def __init__(self, binary, path, url, tools, extra=(), settings=None): + def __init__(self, binary, path, url, tools, extra=(), settings=None, env=None): # Keep AF_UNIX paths short even in deep remote snapshot directories. self.directory = tempfile.TemporaryDirectory(prefix='agent-bench-', dir='/tmp') self.socket_path = Path(self.directory.name) / 'daemon.sock' @@ -85,7 +85,7 @@ def __init__(self, binary, path, url, tools, extra=(), settings=None): self.process = subprocess.Popen( [str(binary), *serve_args(path, url, extra=extra), '--socket', str(self.socket_path)], stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, - start_new_session=True, env=clean_env()) + start_new_session=True, env=env or clean_env()) try: deadline = time.monotonic() + 5 while True: diff --git a/docs/CLI.md b/docs/CLI.md index 1321b8fdc..3bf4bce89 100644 --- a/docs/CLI.md +++ b/docs/CLI.md @@ -1,7 +1,7 @@ # CLI contract Agent uses flat verbs: `run`, `follow`, `fork`, `interrupt`, `wait`, `ls`, -`inspect`, `turns`, `usage`, `rm`, `prune`, `approvals`, `answer`, `models`, `start`, `stats`, +`inspect`, `turns`, `usage`, `rm`, `prune`, `approvals`, `answer`, `models`, `companion`, `start`, `stats`, `shutdown`, and `serve`. All named bots use the same commands. There is no parent/child command hierarchy. @@ -177,6 +177,40 @@ for both client commands and `serve`. Keep that registry, the implementation, and behavior tests aligned when adding an option. No CLI framework or runtime dependency is required. +## Companions + +A companion is a command the daemon keeps running for as long as it runs: +a client's helper that agents wait on, such as Sonny's lookout or the +approver, so it lives as long as the daemon whoever started either. The +daemon knows nothing of what one is. + +```sh +agent companion add NAME [--cwd DIR] [--env KEY=VALUE]... -- COMMAND [ARGS]... +agent companion rm NAME +agent companion ls +``` + +`add` keeps the companion in the store under NAME (1 to 64 letters, +digits, `.`, `_` or `-`, not starting with `.`) and starts it now; every +later daemon on that store starts it once it listens. The same spec again +changes nothing (`changed: false`); another spec under the name replaces +the old one, which is stopped first. `rm` stops and forgets one +(`removed: false` when there was none); `ls` lists them, as `stats` does +under `companions`, each with its spec, `running`, `pid`, `started_ms`, +`restarts` and `last_exit` (`code` and `signal`, or `spawn_failed` with +why, and `at_ms`). Each runs in its own process group with the daemon's +environment, its `--env` over that, and `AGENT_STORE`, `AGENT_SOCKET`, +`AGENT_BIN` and `AGENT_COMPANION` (its name), never a bot shell's marks; +its stdin is empty and its output goes to `STORE.companions/NAME.log`. One +that exits starts again after 1 s, the wait doubling up to 60 s while it +keeps exiting within a minute of starting. The daemon's shutdown stops +each with SIGTERM to its group, then SIGKILL after 5 s. A companion's own +sessions keep no daemon from `--idle-exit`, and a daemon on stdio, with no +socket to reach it on, runs none. The request is `companions` with `add` +(`name`, `command`, `args`, `env`, `cwd`), `remove` (a name) or neither; +a sandboxed bot's commands may not make it, since a companion runs as the +daemon does. + ## Output and status Default output is machine-readable JSON. `run` and `follow` stream one JSON diff --git a/docs/LOOKOUT.md b/docs/LOOKOUT.md index f99f3f38f..66e72eed1 100644 --- a/docs/LOOKOUT.md +++ b/docs/LOOKOUT.md @@ -20,19 +20,23 @@ and the rest, `sonny/src/served.rs`). ## One lookout a daemon -Sonny starts `sonny lookout --store STORE --socket SOCKET` detached when it -opens a project, and again after the daemon restarts. A lookout holds +`sonny lookout --store STORE --socket SOCKET` runs as the daemon's +companion `sonny-lookout` (CLI.md, Companions): Sonny adds it when it opens +a project, and the daemon keeps it running for as long as it runs, starts +it again if it exits, and starts it with every later daemon on that store, +whether or not a window is open. A lookout holds `STORE.lookout.lock` while it runs, so a second one prints `{"event":"already_watching"}` and exits, unless another build of Sonny started the one running, as an upgrade leaves it: that one is stopped, and the new one takes its place, so the tools the new Sonny gives its agents -are answered. `STORE.lookout.sonny` says which build runs. It watches for every bot with +are answered; the daemon's companion takes the place of one a window +started detached, of any build. `STORE.lookout.sonny` says which build runs. It watches for every bot with subscriptions on the daemon, whichever Sonny opened it, so a lead hears from its tasks even while no window shows that project. It exits -once the daemon has been unreachable for 30 seconds; the Sonny that starts -the next daemon starts the next lookout. Its own lines go to -`STORE.lookout.log`, and each subscriber it watches (`seated`) and each -note it sends (`woke`) to `STORE.lookout.jsonl`. +once the daemon has been unreachable for 30 seconds, and the daemon stops +it as it stops. Its own lines go to `STORE.companions/sonny-lookout.log`, +and each subscriber it watches (`seated`) and each note it sends (`woke`) +to `STORE.lookout.jsonl`. It keeps a cursor in `STORE.lookout.cursor`: the newest event it saw, or just before the oldest one whose news it has not told yet. A lookout that diff --git a/sonny/README.md b/sonny/README.md index 11aab072c..896aafdc2 100644 --- a/sonny/README.md +++ b/sonny/README.md @@ -50,7 +50,9 @@ folder under `~/.agent`. `~/.agent/env` (`KEY=VALUE` lines, readable only by you) holds the provider keys, and `AWS_PROFILE` for Bedrock, which Sonny applies over your shell's to -the daemon, the approver and the lookout. Settings go in +the daemon it starts; the approver and the lookout run as that daemon's +companions, with its environment, so they live as long as it does whether +or not a window is open. Settings go in `~/.agent/config.toml`, as Codex keeps them: `provider` (the daemon's `--provider` specs, after the shell's `AGENT_PROVIDER`), `model` (a new agent's when none is named) and the efforts below. When no provider is named @@ -176,7 +178,7 @@ notifies, `herdr agent wait` and `herdr agent prompt` reach it by name 0.9.2 lists Sonny but takes no command to restart it). Sonny reports from a pane of the Herdr server on its own machine: one reached by plain `ssh` from a Herdr pane is not listed, one in a remote machine Herdr itself connects to is. -The daemon, lookout and approver Sonny starts run in a session of their own, +The daemon Sonny starts, and the lookout and approver it keeps, run in a session of their own, so closing their pane or stopping Herdr leaves them and their turns running. In cmux, which reads OSC 7501 and OSC 777 itself, the workspace's sidebar row shows which task is working, waiting on an approval or done, and a @@ -188,8 +190,9 @@ then cmux offers it to run by hand); quitting Sonny clears it. cmux lets only programs started in its tabs reach it, by default; agents' shells run in the daemon, not in cmux, so set cmux's Settings > Automation to "Automation mode" (your own programs, no ancestry check) for agents to run -commands in tabs and read them. The lookout Sonny starts in cmux reaches it -either way; one started outside cmux needs Automation mode too. +commands in tabs and read them. The lookout, the daemon's companion, reaches it +either way when the daemon started from a cmux tab; otherwise it needs +Automation mode too. The place a project was laid out in is kept in `~/.agent/sonny/places`, so an agent's `run_in_pane` and `show_agent` calls go to that Herdr or cmux workspace or diff --git a/sonny/src/acp.rs b/sonny/src/acp.rs index 289ecd3b7..a7c4c0bac 100644 --- a/sonny/src/acp.rs +++ b/sonny/src/acp.rs @@ -1500,14 +1500,14 @@ async fn services( if let Some(project) = project { lookout::subscribe_lead(&agent.options.store, bot, project)?; } - crate::start_lookout(&agent.options)?; + crate::start_lookout(client).await?; if approvals::mode(&record) == Some("auto") { let model = format!( "{}/{}", record["provider"].as_str().unwrap_or(""), record["model"].as_str().unwrap_or("") ); - crate::serve_auto(client, &agent.options, &model).await?; + crate::serve_auto(client, &model).await?; } Ok(()) } @@ -1586,14 +1586,14 @@ async fn setup_lead( } } lookout::subscribe_lead(&agent.options.store, &lead.bot, &lead.project)?; - crate::start_lookout(&agent.options)?; + crate::start_lookout(client).await?; if gating.or(mode) == Some("auto") { let model = format!( "{}/{}", record["provider"].as_str().unwrap_or(""), record["model"].as_str().unwrap_or("") ); - crate::serve_auto(client, &agent.options, &model).await?; + crate::serve_auto(client, &model).await?; } Ok(()) } diff --git a/sonny/src/lookout/mod.rs b/sonny/src/lookout/mod.rs index 0dce16b91..aa227e35b 100644 --- a/sonny/src/lookout/mod.rs +++ b/sonny/src/lookout/mod.rs @@ -292,7 +292,7 @@ fn beside(store: &Path, suffix: &str) -> PathBuf { /// This Sonny, told apart from another build of it: its version and when /// its binary was written. -fn identity() -> String { +pub(crate) fn identity() -> String { let written = std::env::current_exe() .and_then(std::fs::metadata) .and_then(|meta| meta.modified()) @@ -308,7 +308,9 @@ fn identity() -> String { /// when the lookout holding it is this Sonny's own. fn replace(store: &Path) -> Result, String> { let theirs = std::fs::read_to_string(beside(store, ".lookout.sonny")).unwrap_or_default(); - if theirs == identity() { + // The daemon's own companion takes the place of one a window started + // before Sonny kept it as one, of this build too. + if theirs == identity() && std::env::var_os("AGENT_COMPANION").is_none() { return Ok(None); } let pid: Option = std::fs::read_to_string(beside(store, ".lookout.lock")) @@ -343,8 +345,14 @@ fn lock(store: &std::path::Path) -> Result, String> { .open(&path) .map_err(|e| format!("{}: {e}", std::path::Path::new(&path).display()))?; // SAFETY: flock on a descriptor `file` owns and keeps open; the lock is - // released when the process exits. - if unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0 { + // released when the process exits. Tried twice, since `running` holds + // it for an instant. + let take = || unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) } == 0; + let taken = take() || { + std::thread::sleep(Duration::from_millis(50)); + take() + }; + if taken { // Its process, which `/status` asks after. use std::io::Write; let _ = file.set_len(0); @@ -1010,14 +1018,24 @@ fn cursor_path(store: &Path) -> PathBuf { /// The most of the log `status` reads, from its end. const STATUS_BYTES: u64 = 1 << 20; +/// Whether a lookout watches `store`'s daemon: one holds its lock. The pid +/// the file names is no answer, since after a restart another process may +/// have it. The lock is tried shared and let go at once; a lookout +/// starting meanwhile tries again (`lock`). +pub fn running(store: &std::path::Path) -> bool { + use std::os::fd::AsRawFd; + let Ok(file) = std::fs::File::open(beside(store, ".lookout.lock")) else { + return false; + }; + // SAFETY: flock on a descriptor `file` owns; closing it lets go. + let taken = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_SH | libc::LOCK_NB) } == 0; + !taken && std::io::Error::last_os_error().raw_os_error() == Some(libc::EWOULDBLOCK) +} + /// One line on the lookout for a subscriber, as `/status` shows it: /// whether one runs and watches it, and today's notes. pub fn status(store: &std::path::Path, lead: &str) -> String { - let pid: Option = std::fs::read_to_string(beside(store, ".lookout.lock")) - .ok() - .and_then(|pid| pid.trim().parse().ok()); - // SAFETY: signal 0 sends nothing; it only asks whether the process lives. - if !pid.is_some_and(|pid| pid > 0 && unsafe { libc::kill(pid, 0) } == 0) { + if !running(store) { return "Lookout ยท not running".to_owned(); } let (mut seated, mut notes) = (false, 0); @@ -2339,6 +2357,24 @@ mod tests { assert!(lookout.updates.due()); } + #[test] + fn a_lookout_runs_while_one_holds_its_lock_whatever_pid_it_names() { + let dir = std::env::temp_dir().join(format!("lookout-running-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let store = dir.join("state.sqlite"); + assert!(!running(&store)); + // A lookout gone before a restart, its pid now another process's. + std::fs::write(dir.join("state.sqlite.lookout.lock"), "1").unwrap(); + assert!(!running(&store)); + let held = lock(&store).unwrap().unwrap(); + assert!(running(&store)); + // Asking leaves it held. + assert!(lock(&store).unwrap().is_none()); + drop(held); + assert!(!running(&store)); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn status_says_whether_it_watches_from_the_start() { let dir = std::env::temp_dir().join(format!("lookout-status-{}", std::process::id())); diff --git a/sonny/src/main.rs b/sonny/src/main.rs index 74c204f46..880a6aefa 100644 --- a/sonny/src/main.rs +++ b/sonny/src/main.rs @@ -92,7 +92,7 @@ use serde_json::{Value, json}; use std::{ collections::{BTreeMap, HashMap}, io::Write, - path::PathBuf, + path::{Path, PathBuf}, sync::Arc, time::{Duration, Instant}, }; @@ -414,74 +414,70 @@ fn apply_env_file(command: &mut std::process::Command) -> Result<(), String> { Ok(()) } -/// An `auto` gate needs an approver serving it: when none does, one is -/// started the way the CLI starts it, detached, logging beside the store, -/// judged by `model` unless a judge is set. -async fn serve_auto(client: &Client, options: &Options, model: &str) -> Result<(), String> { - let stats = client - .request("stats", json!({})) +/// The daemon's companion that answers an `auto` gate: `agent approver`, +/// judged by `model` unless a judge is set, which the daemon then keeps +/// running for as long as it runs, as it does the lookout. One kept already +/// keeps the judge it was given first, so leads on other models do not +/// start it again. +async fn serve_auto(client: &Client, model: &str) -> Result<(), String> { + let bin = agent_bin().ok_or("Agent's binary was not found to judge calls; set AGENT_BIN")?; + let args = daemon_args("approver"); + let kept = client + .request("companions", json!({})) .await .map_err(|e| e.to_string())?; - if stats["approvers"] - .as_array() - .is_some_and(|tags| tags.iter().any(|tag| tag == "auto")) + if (kept["companions"].as_array().into_iter().flatten()) + .any(|c| c["name"] == APPROVER && c["command"] == json!(bin) && c["args"] == json!(args)) { return Ok(()); } - let bin = agent_bin().ok_or("Agent's binary was not found to judge calls; set AGENT_BIN")?; - let mut log = options.store.clone().into_os_string(); - log.push(".approver.log"); - let log = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&log) - .map_err(|e| e.to_string())?; - let mut command = agent_command(&bin)?; - agent_client::own_session(&mut command); - command - .env("AGENT_MODEL", model) - .arg("approver") - .arg("--store") - .arg(&options.store) - .arg("--socket") - .arg(&options.socket) - .stdin(std::process::Stdio::null()) - .stdout(log.try_clone().map_err(|e| e.to_string())?) - .stderr(log) - .spawn() - .map_err(|e| format!("{}: {e}", bin.display()))?; - Ok(()) + keep(client, APPROVER, &bin, args, json!({"AGENT_MODEL": model})).await } /// The lookout, `sonny lookout`, which tells leads what happens in their -/// tasks and answers Sonny's tools: started detached, logging beside the -/// store. One already watching this daemon keeps it, and the new one -/// exits. -fn start_lookout(options: &Options) -> Result<(), String> { +/// tasks and answers Sonny's tools: the daemon's companion, so it runs for +/// as long as the daemon does, whether or not a window is open, and comes +/// back with the next daemon. This Sonny's own build replaces another's: +/// the build is part of what the daemon keeps, so a Sonny upgraded in +/// place has the daemon start the new one. +async fn start_lookout(client: &Client) -> Result<(), String> { let me = std::env::current_exe().map_err(|e| e.to_string())?; - let mut log = options.store.clone().into_os_string(); - log.push(".lookout.log"); - let log = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&log) - .map_err(|e| e.to_string())?; - let mut command = agent_command(&me)?; - agent_client::own_session(&mut command); - let mut child = command - .arg("lookout") - .arg("--store") - .arg(&options.store) - .arg("--socket") - .arg(&options.socket) - .stdin(std::process::Stdio::null()) - .stdout(log.try_clone().map_err(|e| e.to_string())?) - .stderr(log) - .spawn() - .map_err(|e| format!("{}: {e}", me.display()))?; - // Reaped when it exits, as it does at once when another one watches. - std::thread::spawn(move || child.wait()); - Ok(()) + keep( + client, + LOOKOUT, + &me, + daemon_args("lookout"), + json!({"SONNY_BUILD": lookout::identity()}), + ) + .await +} + +/// What Sonny names its companions in the daemon. +const LOOKOUT: &str = "sonny-lookout"; +const APPROVER: &str = "approver-auto"; + +/// A companion's arguments: its store and socket come from the daemon +/// that runs it (`AGENT_STORE`, `AGENT_SOCKET`), which may serve the store +/// on another socket later. +fn daemon_args(command: &str) -> Vec { + vec![command.to_owned()] +} + +/// Ask the daemon to keep `name` running `bin` with `args`; the same again +/// changes nothing. +async fn keep( + client: &Client, + name: &str, + bin: &Path, + args: Vec, + env: Value, +) -> Result<(), String> { + let add = json!({"name": name, "command": bin, "args": args, "env": env}); + client + .request("companions", json!({"add": add})) + .await + .map(drop) + .map_err(|e| format!("{name}: {e}")) } /// The daemon, started the way the CLI starts it when it is not running. @@ -890,7 +886,7 @@ async fn run() -> Result<(), String> { app.notice = Some(format!("the lookout: {error}")); } // It answers the tools Sonny gives every agent, a lead or not. - if let Err(error) = start_lookout(&app.options) { + if let Err(error) = start_lookout(&app.client).await { app.notice = Some(error); } // The lead told what its folder tells it today (skills and memory added @@ -950,7 +946,7 @@ async fn run() -> Result<(), String> { if fresh { app.gating = Some(approvals::MODES[0].0); } else if app.mode == Some("auto") - && let Err(error) = serve_auto(&app.client, &app.options, &app.model).await + && let Err(error) = serve_auto(&app.client, &app.model).await { app.notice = Some(error); } @@ -2957,8 +2953,7 @@ impl App { if let Some(mode) = gating { self.mode = Some(mode); if mode == "auto" - && let Err(error) = - serve_auto(&self.client, &self.options, &self.model).await + && let Err(error) = serve_auto(&self.client, &self.model).await { self.notice = Some(error); } @@ -3381,13 +3376,13 @@ impl App { } self.approvals_for(self.running.as_ref().map(|r| r.turn)) .await; - // A daemon that restarted took the approver and the lookout - // Sonny started with it. - if let Err(error) = start_lookout(&self.options) { + // A daemon that restarted brings back its companions; one + // started anew, as by another store's Sonny, gets them now. + if let Err(error) = start_lookout(&self.client).await { self.notice = Some(error); } if self.mode == Some("auto") - && let Err(error) = serve_auto(&self.client, &self.options, &self.model).await + && let Err(error) = serve_auto(&self.client, &self.model).await { self.notice = Some(error); } diff --git a/src/cli.rs b/src/cli.rs index 714681c80..c633558fe 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -152,6 +152,14 @@ const COMMANDS: &[Command] = &[ settings: "", startup: false, }, + Command { + name: "companion", + about: "Keep a command running for as long as the daemon runs, stop one, or list them", + usage: "companion (add NAME [--cwd DIR] [--env KEY=VALUE]... -- COMMAND [ARGS]... | rm NAME | ls)", + flags: "--cwd --env --pretty", + settings: "", + startup: true, + }, Command { name: "start", about: "Start the store's daemon if none runs; prints its ready line", @@ -379,6 +387,8 @@ fn print_flags(flags: &str) { "MODE", "A new bot's command sandbox: read-only, workspace-write or danger-full-access; default none", ), + "--cwd" => ("DIR", "The folder a companion runs in"), + "--env" => ("KEY=VALUE", "A variable a companion gets; repeat for more"), "--writable" => ( "PATH", "Another folder a workspace-write bot may change; repeat for more", @@ -537,7 +547,7 @@ pub fn prepare(args: Vec) -> Result>> { let mut positional = 0; while let Some(arg) = iter.next() { if arg == "--" { - if matches!(c.name, "run" | "wait") { + if matches!(c.name, "run" | "wait" | "companion") { out.push(arg); } for value in iter { @@ -570,7 +580,7 @@ pub fn prepare(args: Vec) -> Result>> { if flags.iter().any(|f| f == flag) && !matches!( flag, - "--provider" | "--image" | "--writable" | "--unsandboxed" + "--provider" | "--image" | "--writable" | "--unsandboxed" | "--env" ) { return fail_with("usage", format!("{flag} may only be specified once")); @@ -683,7 +693,7 @@ pub fn prepare(args: Vec) -> Result>> { if positional != 1 { return fail_with("usage", "login needs a provider: chatgpt"); } - } else if !matches!(c.name, "run" | "wait") && positional != 0 { + } else if !matches!(c.name, "run" | "wait" | "companion") && positional != 0 { return fail_with("usage", format!("{} takes no positional arguments", c.name)); } if c.name == "run" && has("--fork") { diff --git a/src/client.rs b/src/client.rs index afae8e1af..c25d99a53 100644 --- a/src/client.rs +++ b/src/client.rs @@ -133,6 +133,9 @@ struct Options { settings: serde_json::Map, settings_flags: Vec, positional: Vec, + /// A companion's folder and variables. + cwd: Option, + env: Vec, /// The `--store` and `--socket` flags, shell-quoted, that reach this /// daemon from any shell; empty when both are the defaults. target: String, @@ -202,6 +205,8 @@ fn parse(args: &[String]) -> Result { settings: serde_json::Map::new(), settings_flags: Vec::new(), positional: Vec::new(), + cwd: None, + env: Vec::new(), target: String::new(), }; let mut iter = args.iter(); @@ -240,6 +245,8 @@ fn parse(args: &[String]) -> Result { match flag { "--store" => store = Some(PathBuf::from(value)), "--socket" => socket = Some(PathBuf::from(value)), + "--cwd" => options.cwd = Some(std::path::absolute(&value)?), + "--env" => options.env.push(value), "--provider" => { options.providers.push(value); options.providers_explicit = true; @@ -1120,6 +1127,7 @@ pub fn main(args: Vec) -> Result { "approvals" => approvals(&options), "answer" => answer(&options), "approver" => approver(&options), + "companion" => companion(&options), "models" => models(&options), // The daemon's ready line, from the running one or one started now, // with the socket it answered on, so a caller that reached this @@ -1930,6 +1938,50 @@ fn approver(options: &Options) -> Result { }) } +/// `companion add`, `rm` or `ls`: what the daemon keeps running beside it. +fn companion(options: &Options) -> Result { + let usage = || { + fail_with( + "usage", + "companion needs add NAME -- COMMAND..., rm NAME or ls", + ) + }; + let request = match options + .positional + .iter() + .map(String::as_str) + .collect::>()[..] + { + ["ls"] => json!({}), + ["rm", name] => json!({"remove": name}), + ["add", name, command, ..] => { + let mut env = serde_json::Map::new(); + for pair in &options.env { + let Some((key, value)) = pair.split_once('=') else { + return fail_with("usage", format!("--env needs KEY=VALUE, not {pair}")); + }; + env.insert(key.to_owned(), json!(value)); + } + let cwd = (options.cwd.as_ref()) + .map(|cwd| { + cwd.to_str() + .ok_or(Error::with("usage", "--cwd needs a UTF-8 path")) + }) + .transpose()?; + json!({"add": {"name": name, "command": command, + "args": options.positional[3..], "env": env, "cwd": cwd}}) + } + _ => return usage(), + }; + if !request["add"].is_object() && (options.cwd.is_some() || !options.env.is_empty()) { + return fail_with("usage", "--cwd and --env go with companion add"); + } + let mut connection = ensure_daemon(options)?; + let answer = connection.request("companions", request)?; + print_json(&answer, options.pretty)?; + Ok(0) +} + fn answer(options: &Options) -> Result { // The same kind of guard as the creator identity: it stops a bot's own // shell from answering by accident, not a determined command. diff --git a/src/server/companions.rs b/src/server/companions.rs new file mode 100644 index 000000000..2eb8c7cee --- /dev/null +++ b/src/server/companions.rs @@ -0,0 +1,518 @@ +//! Companions: commands a client asks the daemon to keep running for as +//! long as it runs, such as a client's own helpers that agents wait on. +//! The daemon knows nothing of what they are. Each is kept in the store +//! by name, started once the daemon listens, started again with backoff +//! when it exits, and stopped when the daemon stops, so it lives as long +//! as the daemon whoever started it. +use agent_runtime::{Error, Result, fail_with}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::{ + collections::{BTreeMap, HashMap}, + path::{Path, PathBuf}, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; +use tokio::{sync::watch, task::JoinHandle}; + +/// The first wait before a companion that exited starts again, doubled +/// each time it exits again soon, up to `BACKOFF_MAX`. +const BACKOFF: Duration = Duration::from_secs(1); +const BACKOFF_MAX: Duration = Duration::from_secs(60); +/// A run this long counts as healthy: the next wait starts over. +const HEALTHY: Duration = Duration::from_secs(60); +/// How long a stopped companion has after SIGTERM before SIGKILL. +const GRACE: Duration = Duration::from_secs(5); +const NAME_BYTES: usize = 64; + +/// What a companion runs: a command and its arguments, as `exec` takes +/// them, with variables over the daemon's own and a folder to run in. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Spec { + pub command: String, + #[serde(default)] + pub args: Vec, + #[serde(default)] + pub env: BTreeMap, + pub cwd: Option, +} + +/// A companion's add request: its name and what it runs. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Added { + pub name: String, + pub command: String, + #[serde(default)] + pub args: Vec, + #[serde(default)] + pub env: BTreeMap, + pub cwd: Option, +} + +impl Added { + pub fn split(self) -> (String, Spec) { + let spec = Spec { + command: self.command, + args: self.args, + env: self.env, + cwd: self.cwd, + }; + (self.name, spec) + } +} + +impl Spec { + pub fn check(&self, name: &str) -> Result<()> { + let refuse = |why: String| fail_with("invalid_companion", format!("{name}: {why}")); + if name.is_empty() + || name.len() > NAME_BYTES + || name.starts_with('.') + || !name + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b"._-".contains(&b)) + { + return fail_with( + "invalid_companion", + format!( + "a companion's name is 1 to {NAME_BYTES} letters, digits, '.', '_' or '-', not starting with '.'" + ), + ); + } + if self.command.is_empty() { + return refuse("command is empty".into()); + } + if let Some(key) = self + .env + .keys() + .find(|key| key.is_empty() || key.contains(['=', '\0'])) + { + return refuse(format!("{key:?} is not a variable name")); + } + if self + .cwd + .as_deref() + .is_some_and(|cwd| !Path::new(cwd).is_absolute()) + { + return refuse("cwd must be an absolute path".into()); + } + let nul = |text: &String| text.contains('\0'); + if nul(&self.command) || self.args.iter().any(nul) || self.env.values().any(nul) { + return refuse("a NUL byte".into()); + } + Ok(()) + } +} + +/// How a companion is doing, as `companions` lists it. +#[derive(Default)] +struct Status { + pid: Option, + started_ms: Option, + /// Starts after the first. + restarts: u64, + last_exit: Option, +} + +struct Kept { + spec: Spec, + status: Arc>, + stop: watch::Sender, + task: JoinHandle<()>, +} + +/// The companions this daemon keeps, by name. +pub struct Companions { + /// The variables every companion gets over the daemon's own: where + /// the daemon is. + environment: Vec<(String, String)>, + /// Where each companion's output goes, `STORE.companions/NAME.log`. + logs: PathBuf, + /// Whether they run: only a daemon on a socket has one to reach it on. + runs: bool, + kept: HashMap, + /// Companions removed and still stopping: shutdown waits for them, and + /// one added again under the name starts once its last has stopped. + stopping: HashMap>, +} + +impl Companions { + pub fn new(store: &Path, environment: Vec<(String, String)>, runs: bool) -> Self { + let mut logs = store.as_os_str().to_owned(); + logs.push(".companions"); + Self { + environment, + logs: logs.into(), + runs, + kept: HashMap::new(), + stopping: HashMap::new(), + } + } + + /// Keep `spec` running under `name`: a new one starts, one whose spec + /// changed starts again once the old run has stopped, and the same + /// spec again changes nothing. + pub fn keep(&mut self, name: &str, spec: Spec) { + if self.kept.get(name).is_some_and(|kept| kept.spec == spec) { + return; + } + let before = match self.kept.remove(name) { + Some(kept) => { + let _ = kept.stop.send(true); + Some(kept.task) + } + None => self.stopping.remove(name), + }; + let status = Arc::new(Mutex::new(Status::default())); + let (stop, stopped) = watch::channel(false); + let run = Run { + name: name.to_owned(), + spec: spec.clone(), + environment: self.environment.clone(), + log: self.logs.join(format!("{name}.log")), + status: status.clone(), + }; + let runs = self.runs; + let task = tokio::spawn(async move { + if let Some(before) = before { + let _ = before.await; + } + if runs { + run.keep(stopped).await; + } + }); + self.kept.insert( + name.to_owned(), + Kept { + spec, + status, + stop, + task, + }, + ); + } + + /// Stop `name`'s companion; whether one was kept. + pub fn drop_one(&mut self, name: &str) -> bool { + self.stopping.retain(|_, task| !task.is_finished()); + let Some(kept) = self.kept.remove(name) else { + return false; + }; + let _ = kept.stop.send(true); + self.stopping.insert(name.to_owned(), kept.task); + true + } + + /// Stop every companion and wait for them to exit. + pub async fn stop(&mut self) { + let tasks: Vec<_> = self + .kept + .drain() + .map(|(_, kept)| { + let _ = kept.stop.send(true); + kept.task + }) + .chain(self.stopping.drain().map(|(_, task)| task)) + .collect(); + futures_util::future::join_all(tasks).await; + } + + /// Whether `pid` is a companion's process, or one in its group (what a + /// companion's script runs): its sessions keep no daemon from going idle. + pub fn runs_pid(&self, pid: u32) -> bool { + // SAFETY: getpgid only reads the process's group. + let group = unsafe { libc::getpgid(pid as i32) }; + self.kept.values().any(|kept| { + let leader = kept.status.lock().unwrap_or_else(|e| e.into_inner()).pid; + leader.is_some_and(|leader| leader == pid || leader as i32 == group) + }) + } + + /// Each companion, by name, with what it runs and how it is doing. + pub fn listed(&self) -> Value { + let mut names: Vec<_> = self.kept.keys().collect(); + names.sort(); + let listed: Vec = names + .into_iter() + .map(|name| self.shown(name).expect("kept")) + .collect(); + json!(listed) + } + + pub fn shown(&self, name: &str) -> Option { + let kept = self.kept.get(name)?; + let status = kept.status.lock().unwrap_or_else(|e| e.into_inner()); + let mut shown = json!({"name": name, "command": kept.spec.command, + "args": kept.spec.args, "env": kept.spec.env, "cwd": kept.spec.cwd, + "running": status.pid.is_some(), "pid": status.pid, + "started_ms": status.started_ms, "restarts": status.restarts, + "last_exit": status.last_exit}); + if !self.runs { + shown["running"] = json!(false); + } + Some(shown) + } +} + +/// One companion's runs, one after another. +struct Run { + name: String, + spec: Spec, + environment: Vec<(String, String)>, + log: PathBuf, + status: Arc>, +} + +impl Run { + fn status(&self) -> std::sync::MutexGuard<'_, Status> { + self.status.lock().unwrap_or_else(|e| e.into_inner()) + } + + async fn keep(self, mut stop: watch::Receiver) { + let mut backoff = BACKOFF; + let mut first = true; + while !*stop.borrow() { + if !first { + self.status().restarts += 1; + } + first = false; + let started = Instant::now(); + let exit = match self.spawn().await { + Ok(mut child) => { + { + let mut status = self.status(); + status.pid = child.id(); + status.started_ms = Some(epoch_ms()); + } + let exit = tokio::select! { + exit = child.wait() => exited(exit), + _ = stop.changed() => { + terminate(&mut child).await; + self.status().pid = None; + return; + } + }; + self.status().pid = None; + exit + } + Err(error) => json!({"error": "spawn_failed", "detail": error}), + }; + let mut exit = exit; + exit["at_ms"] = json!(epoch_ms()); + self.status().last_exit = Some(exit); + if started.elapsed() >= HEALTHY { + backoff = BACKOFF; + } + tokio::select! { + _ = tokio::time::sleep(backoff) => {} + _ = stop.changed() => return, + } + backoff = (backoff * 2).min(BACKOFF_MAX); + } + } + + async fn spawn(&self) -> std::result::Result { + // Off the runtime thread: a slow disk, or a FIFO at the log's + // path, holds only this companion. + tokio::fs::create_dir_all(self.log.parent().expect("in its folder")) + .await + .map_err(|e| e.to_string())?; + // Nonblocking, so a FIFO there with no reader fails at once rather + // than holding the companion, and its stopping, for ever. + let log = tokio::fs::OpenOptions::new() + .create(true) + .append(true) + .custom_flags(libc::O_NONBLOCK) + .open(&self.log) + .await + .map_err(|e| format!("{}: {e}", self.log.display()))?; + let regular = log.metadata().await.is_ok_and(|meta| meta.is_file()); + if !regular { + return Err(format!("{}: not a regular file", self.log.display())); + } + let log = log.into_std().await; + let mut command = tokio::process::Command::new(&self.spec.command); + command + .args(&self.spec.args) + .envs(self.environment.iter().map(|(k, v)| (k, v))) + .env("AGENT_COMPANION", &self.name) + .envs(&self.spec.env) + .stdin(std::process::Stdio::null()) + .stdout(log.try_clone().map_err(|e| e.to_string())?) + .stderr(log) + // Its own group, so stopping it stops what it started too. + .process_group(0) + .kill_on_drop(true); + if let Some(cwd) = &self.spec.cwd { + command.current_dir(cwd); + } + command + .spawn() + .map_err(|e| format!("{}: {e}", self.spec.command)) + } +} + +fn exited(exit: std::io::Result) -> Value { + use std::os::unix::process::ExitStatusExt; + match exit { + Ok(status) => json!({"code": status.code(), "signal": status.signal()}), + Err(error) => json!({"error": "wait_failed", "detail": error.to_string()}), + } +} + +/// SIGTERM to its group, then SIGKILL to what is left of it once its +/// grace is up, the companion itself gone or not. +async fn terminate(child: &mut tokio::process::Child) { + let Some(pid) = child.id() else { return }; + let group = -(pid as i32); + let deadline = tokio::time::Instant::now() + GRACE; + // SAFETY: a signal to the process group this companion leads. + unsafe { libc::kill(group, libc::SIGTERM) }; + let _ = tokio::time::timeout_at(deadline, child.wait()).await; + // SAFETY: signal 0 only asks whether any of the group lives. + while unsafe { libc::kill(group, 0) } == 0 && tokio::time::Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(50)).await; + } + // SAFETY: as above; a group already gone is no error to mind. + unsafe { libc::kill(group, libc::SIGKILL) }; + let _ = child.wait().await; +} + +fn epoch_ms() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_millis() as i64) +} + +/// The spec a stored row holds. +pub fn stored(name: &str, spec: &str) -> Result { + serde_json::from_str(spec).map_err(|_| Error::with("store_companion_invalid", name.to_owned())) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn spec(command: &str, args: &[&str]) -> Spec { + Spec { + command: command.into(), + args: args.iter().map(|a| a.to_string()).collect(), + env: BTreeMap::new(), + cwd: None, + } + } + + #[test] + fn a_companion_s_name_and_spec_are_checked() { + assert!(spec("sonny", &["lookout"]).check("lookout").is_ok()); + for name in ["", ".hidden", "a b", &"x".repeat(65)] { + assert_eq!( + spec("sonny", &[]).check(name).unwrap_err().code, + "invalid_companion" + ); + } + assert!(spec("", &[]).check("x").is_err()); + let mut relative = spec("sonny", &[]); + relative.cwd = Some("here".into()); + assert!(relative.check("x").is_err()); + let mut bad = spec("sonny", &[]); + bad.env.insert("A=B".into(), "1".into()); + assert!(bad.check("x").is_err()); + } + + async fn until(what: impl Fn() -> bool) { + let deadline = Instant::now() + Duration::from_secs(10); + while !what() { + assert!(Instant::now() < deadline, "timed out"); + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + #[tokio::test] + async fn a_companion_runs_again_when_it_exits_and_stops_with_the_daemon() { + let dir = std::env::temp_dir().join(format!("companions-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let store = dir.join("state.sqlite"); + let mut companions = + Companions::new(&store, vec![("AGENT_SOCKET".into(), "/s".into())], true); + let marks = dir.join("marks"); + let script = format!( + "echo \"$AGENT_COMPANION $AGENT_SOCKET\" >> {}; exit 3", + marks.display() + ); + companions.keep("once", spec("sh", &["-c", &script])); + until(|| std::fs::read_to_string(&marks).is_ok_and(|text| text.lines().count() >= 2)).await; + assert_eq!( + std::fs::read_to_string(&marks).unwrap().lines().next(), + Some("once /s") + ); + let shown = companions.shown("once").unwrap(); + assert!(shown["restarts"].as_u64().unwrap() >= 1, "{shown}"); + assert_eq!(shown["last_exit"]["code"], 3); + // The same spec again keeps the run it has. + companions.keep("long", spec("sleep", &["30"])); + until(|| companions.shown("long").unwrap()["running"] == true).await; + let pid = companions.shown("long").unwrap()["pid"].as_u64().unwrap() as u32; + assert!(companions.runs_pid(pid)); + companions.keep("long", spec("sleep", &["30"])); + assert_eq!(companions.shown("long").unwrap()["pid"], json!(pid)); + companions.stop().await; + // SAFETY: signal 0 only asks whether it lives. + assert_ne!(unsafe { libc::kill(pid as i32, 0) }, 0); + assert_eq!(companions.listed(), json!([])); + let _ = std::fs::remove_dir_all(&dir); + } + + #[tokio::test] + async fn a_log_that_is_no_file_fails_its_companion_at_once() { + let dir = std::env::temp_dir().join(format!("companions-fifo-{}", std::process::id())); + std::fs::create_dir_all(dir.join("state.sqlite.companions")).unwrap(); + let fifo = std::ffi::CString::new( + dir.join("state.sqlite.companions/piped.log") + .to_string_lossy() + .as_bytes(), + ) + .unwrap(); + // SAFETY: a FIFO made at a path this test owns. + assert_eq!(unsafe { libc::mkfifo(fifo.as_ptr(), 0o600) }, 0); + let mut companions = Companions::new(&dir.join("state.sqlite"), vec![], true); + companions.keep("piped", spec("sleep", &["30"])); + until(|| companions.shown("piped").unwrap()["last_exit"]["error"] == "spawn_failed").await; + tokio::time::timeout(Duration::from_secs(5), companions.stop()) + .await + .expect("stops"); + let _ = std::fs::remove_dir_all(&dir); + } + + #[tokio::test] + async fn a_removed_companion_s_group_is_stopped_before_the_daemon_is() { + let dir = std::env::temp_dir().join(format!("companions-rm-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let mut companions = Companions::new(&dir.join("state.sqlite"), vec![], true); + // It exits on SIGTERM; what it started does not. + let started = dir.join("started"); + let script = format!( + "(trap '' TERM; exec sleep 30) & echo $! > {}; wait", + started.display() + ); + companions.keep("stubborn", spec("sh", &["-c", &script])); + until(|| std::fs::read_to_string(&started).is_ok_and(|pid| pid.ends_with('\n'))).await; + let pid: i32 = std::fs::read_to_string(&started) + .unwrap() + .trim() + .parse() + .unwrap(); + assert!(companions.drop_one("stubborn")); + assert_eq!(companions.listed(), json!([])); + companions.stop().await; + // Gone, or a zombie no init here reaps, once SIGKILL lands. + let dead = || { + let state = std::fs::read_to_string(format!("/proc/{pid}/stat")).unwrap_or_default(); + // SAFETY: signal 0 only asks whether it lives. + (unsafe { libc::kill(pid, 0) }) != 0 || state.contains(") Z ") + }; + until(dead).await; + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/src/server/mod.rs b/src/server/mod.rs index 52ca2501d..273b10aff 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -1,5 +1,6 @@ //! One process hosts every bot. Client sessions (stdio or Unix socket) speak //! the same JSONL protocol; turns run as tasks; events fan out through a hub. +mod companions; mod handles; mod hub; mod session; @@ -488,6 +489,14 @@ enum Command { tag: String, lease: u64, }, + /// The daemon's companions: with `add`, keep one running under its + /// name for as long as the daemon runs (the same again changes + /// nothing, another spec runs anew); with `remove`, stop and forget + /// one; with neither, list them. + Companions { + add: Option, + remove: Option, + }, /// The daemon's live state for a fleet controller: sessions, turns, /// connections, pools, storage worker, and handle registry. Stats, @@ -889,6 +898,7 @@ fn gate( } struct Service { store: Store, + companions: companions::Companions, transport: Arc, providers: Arc>, registry: Registry, @@ -1253,6 +1263,16 @@ pub async fn run(config: Configuration) -> Result<()> { .socket .as_deref() .and_then(|socket| socket::capabilities(socket, sender.clone(), sessions_next.clone())); + // Where the daemon is, not a bot's shell's marks: a companion is no + // bot's command. + let mut companions = companions::Companions::new( + &config.store, + (environment.iter()) + .filter(|(key, _)| key != "AGENT_SHELL_CONTEXT") + .cloned() + .collect(), + config.socket.is_some(), + ); let registry = registry .with_credentials(credentials) .with_environment(environment) @@ -1275,6 +1295,9 @@ pub async fn run(config: Configuration) -> Result<()> { let mut sessions: HashMap = HashMap::new(); // The sandboxed bot each session on a bot's own socket answers as. let mut callers: HashMap = HashMap::new(); + // Sessions a companion opened, known once as each opens: what they + // send is no client's activity. + let mut companion_sessions: std::collections::HashSet = std::collections::HashSet::new(); let (stdout, stdout_writer) = Output::stdout(); let stdio_owner = config.socket.is_none(); let mut socket_owner = None; @@ -1361,8 +1384,13 @@ pub async fn run(config: Configuration) -> Result<()> { handles.wake(waiting.bot, waiting.turn); } } + // Companions start once the daemon listens, as they were last kept. + for (name, spec) in store.op("companions", |db| db.companions()).await? { + companions.keep(&name, companions::stored(&name, &spec)?); + } let mut service = Service { store, + companions, identity, transport, providers: Arc::new(providers), @@ -1428,11 +1456,15 @@ pub async fn run(config: Configuration) -> Result<()> { _ = terminate.recv() => break, _ = interrupt.recv() => break, _ = tokio::time::sleep_until(drain_until.unwrap_or(undrained)), if drain_until.is_some() => break, - _ = tokio::time::sleep(idle_exit.unwrap_or(Duration::MAX).min(Duration::from_secs(3600))), if idle_exit.is_some() => { - // Idle means no client, no live turn, and no running command. - // Parked turns are durable and resume on the next start. + // Due from the last client's activity, so what a companion + // sends, which wakes this loop, never puts it off. + _ = tokio::time::sleep_until(idle_due(last_activity, idle_exit)), if idle_exit.is_some() => { + // Idle means no client but companions, no live turn, and no + // running command. Parked turns are durable and resume on + // the next start. let running = service.store.op("running_processes", |db| db.running_processes()).await?; - if sessions.is_empty() && service.active.is_empty() && service.admissions.is_empty() && running == 0 { + let clients = sessions.keys().any(|id| !companion_sessions.contains(id)); + if !clients && service.active.is_empty() && service.admissions.is_empty() && running == 0 { if last_activity.elapsed() >= idle_exit.unwrap() { break; } } else { last_activity = std::time::Instant::now(); @@ -1471,17 +1503,26 @@ pub async fn run(config: Configuration) -> Result<()> { } message = inbound.recv() => { let Some(message) = message else { break }; - last_activity = std::time::Instant::now(); + // What a companion sends is no client's activity. + let companion = match &message { + Inbound::Open(_, _, _, peer) => peer.is_some_and(|pid| service.companions.runs_pid(pid)), + Inbound::Request(id, _) | Inbound::Closed(id) => companion_sessions.contains(id), + }; + if !companion { + last_activity = std::time::Instant::now(); + } match message { - Inbound::Open(id, output, caller) => { + Inbound::Open(id, output, caller, _) => { let _ = output.try_send(ready.clone()); sessions.insert(id, output); if let Some(caller) = caller { callers.insert(id, caller); } + if companion { companion_sessions.insert(id); } service.sessions = sessions.len(); } Inbound::Closed(id) => { sessions.remove(&id); callers.remove(&id); + companion_sessions.remove(&id); service.sessions = sessions.len(); service.hub.close_session(id); service.handles.close_session(id); @@ -1584,6 +1625,8 @@ pub async fn run(config: Configuration) -> Result<()> { Ok(()) } .await; + // Companions go with the daemon. + service.companions.stop().await; // The stdio owner's output is gone: no one is left to answer, and a // blocked stdout write cannot be cancelled. The next start ends the rest. if stdio_owner @@ -1711,6 +1754,8 @@ impl Service { } Command::ToolResult { .. } => return refuse("answer a declared tool call"), Command::Shutdown { .. } => return refuse("stop the daemon"), + // A companion runs as the daemon does, outside any sandbox. + Command::Companions { .. } => return refuse("manage the daemon's companions"), _ => {} } let me = self @@ -2951,6 +2996,32 @@ impl Service { }); Err(Error::new("deferred")) } + Command::Companions { add, remove } => match (add, remove) { + (Some(_), Some(_)) => { + fail_with("usage", "companions takes add or remove, not both") + } + (Some(added), None) => { + let (name, spec) = added.split(); + spec.check(&name)?; + let text = serde_json::to_string(&spec) + .map_err(|e| Error::with("invalid_companion", e.to_string()))?; + let kept = name.clone(); + let changed = store + .op("put_companion", move |db| db.put_companion(&kept, &text)) + .await?; + self.companions.keep(&name, spec); + Ok(json!({"companion": self.companions.shown(&name), "changed": changed})) + } + (None, Some(name)) => { + let gone = name.clone(); + let removed = store + .op("remove_companion", move |db| db.remove_companion(&gone)) + .await?; + self.companions.drop_one(&name); + Ok(json!({"name": name, "removed": removed})) + } + (None, None) => Ok(json!({"companions": self.companions.listed()})), + }, Command::Stats => { let (waiting, running, queued, paced, pending_bytes, approvals) = store .op("counts", |db| { @@ -2979,6 +3050,7 @@ impl Service { "paced_turns": paced, "approval_requests": approvals, "approvers": self.hub.served(), + "companions": self.companions.listed(), "queued_turns": queued, "pending_bytes": pending_bytes, "pending_limit": self.limits.pending, @@ -3640,6 +3712,18 @@ impl Service { } } +/// When the daemon next looks whether it has been idle `idle_exit` since +/// `last_activity`: then, or within the hour for a wait too long to add. +fn idle_due( + last_activity: std::time::Instant, + idle_exit: Option, +) -> tokio::time::Instant { + let last = tokio::time::Instant::from_std(last_activity); + idle_exit + .and_then(|idle| last.checked_add(idle)) + .unwrap_or_else(|| tokio::time::Instant::now() + Duration::from_secs(3600)) +} + /// Commit a finished turn. A group that cannot commit (a full disk, say) /// leaves nothing durable, so the completion is tried again, with backoff, /// while the bot stays busy and other bots go on. An interrupt that arrives @@ -4173,6 +4257,11 @@ mod tests { let ending: Arc = Arc::default(); let mut service = Service { store: store.clone(), + companions: companions::Companions::new( + std::path::Path::new("/nonexistent"), + Vec::new(), + false, + ), identity: 0, transport: Transport::new(64, 1).unwrap(), providers: Arc::new(HashMap::new()), @@ -4338,6 +4427,11 @@ mod tests { let (cancel, cancelled) = watch::channel(None); let mut service = Service { store: store.clone(), + companions: companions::Companions::new( + std::path::Path::new("/nonexistent"), + Vec::new(), + false, + ), identity: 0, transport: Transport::new(64, 1).unwrap(), providers: Arc::new(HashMap::new()), @@ -4546,6 +4640,11 @@ mod tests { let mut service = Service { identity: 0, store: store.clone(), + companions: companions::Companions::new( + std::path::Path::new("/nonexistent"), + Vec::new(), + false, + ), transport, providers: Arc::new(HashMap::from([("openai".to_owned(), provider)])), registry, @@ -4778,6 +4877,11 @@ mod tests { Service { identity: 0, store: store.clone(), + companions: companions::Companions::new( + std::path::Path::new("/nonexistent"), + Vec::new(), + false, + ), transport: Transport::new(64, 1).unwrap(), providers: Arc::new(HashMap::new()), registry: Registry::new("echo").unwrap(), diff --git a/src/server/session.rs b/src/server/session.rs index e0fe796fe..56b1a7d1e 100644 --- a/src/server/session.rs +++ b/src/server/session.rs @@ -12,8 +12,9 @@ use tokio::{ // box for every request. Open/Closed messages are infrequent control traffic. #[allow(clippy::large_enum_variant)] pub enum Inbound { - /// A connection, and the sandboxed bot whose own socket it came in on. - Open(u64, Output, Option), + /// A connection, the sandboxed bot whose own socket it came in on, and + /// the process at its other end, where the system says. + Open(u64, Output, Option, Option), Request(u64, Result), Closed(u64), } @@ -77,9 +78,14 @@ pub async fn accept( ) { while let Ok((stream, _)) = listener.accept().await { let id = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let peer = stream + .peer_cred() + .ok() + .and_then(|cred| cred.pid()) + .and_then(|pid| u32::try_from(pid).ok()); let (read, write) = stream.into_split(); if sender - .send(Inbound::Open(id, Output::writer(write), caller)) + .send(Inbound::Open(id, Output::writer(write), caller, peer)) .await .is_err() { diff --git a/src/store/db.rs b/src/store/db.rs index b6f1a5735..0318caca2 100644 --- a/src/store/db.rs +++ b/src/store/db.rs @@ -1571,7 +1571,7 @@ impl Database { /// Stored schema version, kept in `PRAGMA user_version`. Stores created /// before versioning and stores from newer binaries are rejected; an older /// versioned store is migrated forward, one version at a time, at open. - pub const SCHEMA: i32 = 61; + pub const SCHEMA: i32 = 62; /// Verbatim user prompts a compaction keeps: per-prompt text, and the /// total text plus `(ordinal, String)` entry metadata. Empty entries cost /// space too, so the retained list cannot grow with conversation length. @@ -1753,6 +1753,7 @@ impl Database { INSERT OR IGNORE INTO turn_sequence VALUES (1,0); CREATE TABLE IF NOT EXISTS checkpoints(bot TEXT NOT NULL REFERENCES bots(name), head INTEGER NOT NULL REFERENCES nodes(id), PRIMARY KEY(bot,head)); + CREATE TABLE IF NOT EXISTS companions(name TEXT PRIMARY KEY, spec TEXT NOT NULL); CREATE TABLE IF NOT EXISTS tools(turn INTEGER NOT NULL REFERENCES turns(id), call_id TEXT NOT NULL, status TEXT NOT NULL, custom INTEGER NOT NULL DEFAULT 0, PRIMARY KEY(turn,call_id)); CREATE TABLE IF NOT EXISTS approvals(id INTEGER PRIMARY KEY AUTOINCREMENT, turn INTEGER NOT NULL REFERENCES turns(id), @@ -4121,6 +4122,30 @@ impl Database { } /// Durable lineage, drawn at creation and copied with the database. The /// daemon also binds it to the physical file for its cache namespace. + /// Every companion, by name, as `put_companion` kept it. + pub fn companions(&self) -> Result> { + let mut statement = self + .conn + .prepare_cached("SELECT name,spec FROM companions ORDER BY name")?; + let rows = statement.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?; + Ok(rows.collect::>()?) + } + /// Keep a companion's spec under its name; whether it changed. + pub fn put_companion(&mut self, name: &str, spec: &str) -> Result { + let changed = self.conn.execute( + "INSERT INTO companions(name,spec) VALUES(?1,?2) + ON CONFLICT(name) DO UPDATE SET spec=?2 WHERE spec<>?2", + params![name, spec], + )?; + Ok(changed == 1) + } + /// Forget a companion; whether there was one. + pub fn remove_companion(&mut self, name: &str) -> Result { + Ok(self + .conn + .execute("DELETE FROM companions WHERE name=?", [name])? + == 1) + } pub fn store_identity(&self) -> Result { Ok(self .conn @@ -10433,6 +10458,8 @@ fn migrate(conn: &Connection, from: i32) -> Result<()> { "CREATE TABLE IF NOT EXISTS declined(node INTEGER PRIMARY KEY REFERENCES nodes(id), item BLOB NOT NULL);", )?; + // 61 -> 62: the daemon's companions, in a table the schema creates; + // there were none. Ok(()) } /// 53 -> 54: a turn's denial count became a window of its last decided diff --git a/tests/test_acp.py b/tests/test_acp.py index ae892fc49..11fde30f8 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -118,7 +118,9 @@ def setUp(self): super().setUp() self.sonny = self.binary.parent / 'sonny' self.store = self.path / 'state.sqlite' - self.daemon = SocketClient(self.binary, self.store, self.url, 'echo') + # Its companions, Sonny's lookout among them, have its environment. + self.daemon = SocketClient(self.binary, self.store, self.url, 'echo', + env={**clean_env(), 'HOME': str(self.path)}) self.addCleanup(self.daemon.close) self.t3 = http.server.ThreadingHTTPServer(('127.0.0.1', 0), T3Tools) self.t3.calls, self.t3.headers, self.t3.delegating = queue.Queue(), [], True @@ -130,12 +132,6 @@ def setUp(self): 'GIT_COMMITTER_NAME': 't', 'GIT_COMMITTER_EMAIL': 't@example.com', 'T3_ACP_MCP_ENDPOINT': f'http://127.0.0.1:{self.t3.server_port}/mcp', 'T3_ACP_MCP_AUTHORIZATION': 'Bearer thread-1'} - # The lookout the threads find running. - lookout = subprocess.Popen([str(self.sonny), 'lookout', '--store', str(self.store), - '--socket', str(self.daemon.socket_path)], - env=self.env, stdout=subprocess.PIPE, text=True) - self.addCleanup(lambda: (lookout.kill(), lookout.wait(timeout=5), lookout.stdout.close())) - self.assertEqual(json.loads(lookout.stdout.readline())['event'], 'watching') def repo(self): repo = self.path / 'calc' @@ -190,6 +186,10 @@ def test_a_thread_is_its_folders_lead_and_a_task_it_starts_opens_as_a_thread_of_ thread.send({'id': asked['id'], 'result': {'outcome': {'outcome': 'selected', 'optionId': 'allow'}}}) answered, updates = thread.answer_of(prompt) self.assertEqual(answered['result'], {'stopReason': 'end_turn'}) + # The lookout that answered it is the daemon's, not the thread's. + [lookout] = self.daemon.request('companions')['result']['companions'] + self.assertEqual((lookout['name'], lookout['command'], lookout['args'][0]), + ('sonny-lookout', str(self.sonny.resolve()), 'lookout')) # The task was T3 Code's child of this prompt before it was answered. delegated = self.t3.calls.get_nowait() kinds = [u['sessionUpdate'] for u in updates] diff --git a/tests/test_companions.py b/tests/test_companions.py new file mode 100644 index 000000000..0798c154a --- /dev/null +++ b/tests/test_companions.py @@ -0,0 +1,117 @@ +"""Companions: commands the daemon keeps running for as long as it runs.""" +import json +import os +import subprocess +import time +import unittest + +from bench.targets import clean_env +from tests.test_runtime import ModelFixture + + +def alive(pid): + try: + os.kill(pid, 0) + return True + except OSError: + return False + + +@unittest.skipUnless(os.environ.get('AGENT_TEST_RUNTIME') == '1', 'set AGENT_TEST_RUNTIME=1 after a Rust release build') +class CompanionTests(ModelFixture): + def setUp(self): + super().setUp() + self.store = self.path / 'state.sqlite' + self.daemon = ['--store', str(self.store), '--provider', f'openai=responses,{self.url}'] + self.addCleanup(lambda: subprocess.run([str(self.binary), 'shutdown', '--store', str(self.store)], + env=clean_env(), capture_output=True, timeout=35)) + + def agent(self, *args, check=True): + result = subprocess.run([str(self.binary), *args], env=clean_env(), capture_output=True, text=True, + timeout=30, cwd=self.path) + if check: + self.assertEqual(result.returncode, 0, result.stderr + result.stdout) + return result + + def listed(self): + return {c['name']: c for c in json.loads(self.agent('companion', *self.daemon, 'ls').stdout)['companions']} + + def until(self, what, timeout=15): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + value = what() + if value: + return value + time.sleep(0.1) + self.fail('timed out') + + def test_a_companion_lives_and_dies_with_the_daemon_and_comes_back_with_the_next(self): + marks = self.path / 'marks' + script = f'echo "$AGENT_COMPANION $AGENT_SOCKET $GREETING" >> {marks}; exec sleep 600' + added = json.loads(self.agent('companion', *self.daemon, 'add', 'helper', '--env', 'GREETING=hi', + '--cwd', str(self.path), '--', 'sh', '-c', script).stdout) + self.assertTrue(added['changed']) + self.assertEqual((added['companion']['command'], added['companion']['args']), ('sh', ['-c', script])) + pid = self.until(lambda: self.listed()['helper']['pid']) + socket = json.loads(self.agent('start', *self.daemon).stdout)['socket'] + self.until(lambda: marks.exists() and marks.read_text()) + self.assertEqual(marks.read_text().split(), ['helper', socket, 'hi']) + # The same spec again keeps the run it has. + again = json.loads(self.agent('companion', *self.daemon, 'add', 'helper', '--env', 'GREETING=hi', + '--cwd', str(self.path), '--', 'sh', '-c', script).stdout) + self.assertFalse(again['changed']) + self.assertEqual(self.listed()['helper']['pid'], pid) + self.assertEqual(json.loads(self.agent('stats', *self.daemon).stdout)['companions'][0]['pid'], pid) + # It stops with the daemon, and the next daemon starts it again. + self.agent('shutdown', '--store', str(self.store)) + self.assertFalse(alive(pid)) + self.agent('start', *self.daemon) + second = self.until(lambda: self.listed()['helper']['pid']) + self.assertNotEqual(second, pid) + self.assertTrue(alive(second)) + removed = json.loads(self.agent('companion', *self.daemon, 'rm', 'helper').stdout) + self.assertEqual(removed, {'name': 'helper', 'removed': True}) + self.until(lambda: not alive(second)) + self.assertEqual(self.listed(), {}) + self.assertFalse(json.loads(self.agent('companion', *self.daemon, 'rm', 'helper').stdout)['removed']) + + def test_a_companion_that_exits_starts_again(self): + self.agent('companion', *self.daemon, 'add', 'flaky', '--', 'sh', '-c', 'exit 7') + shown = self.until(lambda: (c := self.listed()['flaky'])['restarts'] >= 1 and c) + self.assertEqual(shown['last_exit']['code'], 7) + + def test_a_companion_s_session_keeps_no_daemon_from_going_idle(self): + daemon = json.loads(self.agent('start', *self.daemon, '--idle-exit', '2s').stdout)['pid'] + # It holds a session open for as long as it runs. + self.agent('companion', '--store', str(self.store), 'add', 'watcher', '--', + str(self.binary), 'follow', '--all', '--store', str(self.store)) + pid = self.until(lambda: self.listed()['watcher']['pid']) + # Asked nothing more, the daemon goes, and its companion with it. + self.until(lambda: not alive(daemon), timeout=20) + self.until(lambda: not alive(pid)) + + def test_a_companion_that_keeps_asking_keeps_no_daemon_from_going_idle(self): + daemon = json.loads(self.agent('start', *self.daemon, '--idle-exit', '2s').stdout)['pid'] + # One connection, a request every 0.3 s, from a process the + # companion's script starts rather than the companion itself. + asker = ('import json, os, socket, time\n' + 's = socket.socket(socket.AF_UNIX); s.connect(os.environ["AGENT_SOCKET"])\n' + 'r = s.makefile("r"); r.readline()\n' + 'n = 0\n' + 'while True:\n' + ' n += 1; s.sendall((json.dumps({"id": n, "op": "stats"}) + "\\n").encode()); r.readline()\n' + ' time.sleep(0.3)\n') + self.agent('companion', '--store', str(self.store), 'add', 'asker', '--env', f'ASKER={asker}', '--', + 'sh', '-c', 'python3 -c "$ASKER"; exit 1') + self.until(lambda: self.listed()['asker']['pid']) + self.until(lambda: not alive(daemon), timeout=20) + + def test_a_bad_companion_is_refused_by_name(self): + refused = self.agent('companion', *self.daemon, 'add', '.x', '--', 'true', check=False) + self.assertNotEqual(refused.returncode, 0) + self.assertEqual(json.loads(refused.stderr.strip().splitlines()[-1])['error'], 'invalid_companion') + self.assertNotEqual(self.agent('companion', *self.daemon, 'add', 'x', check=False).returncode, 0) + + +if __name__ == '__main__': + unittest.main()