diff --git a/crates/goat-agent/Cargo.toml b/crates/goat-agent/Cargo.toml index 2d2d771..8da1c96 100644 --- a/crates/goat-agent/Cargo.toml +++ b/crates/goat-agent/Cargo.toml @@ -25,7 +25,7 @@ goat-tool-browser = { workspace = true } goat-skill = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true } +tokio = { workspace = true, features = ["process"] } tokio-util = { workspace = true } futures = { workspace = true } tracing = { workspace = true } diff --git a/crates/goat-agent/src/delegate.rs b/crates/goat-agent/src/delegate.rs index ccea2c4..ac97147 100644 --- a/crates/goat-agent/src/delegate.rs +++ b/crates/goat-agent/src/delegate.rs @@ -124,7 +124,7 @@ pub(crate) async fn run_delegation( account: env.target.account.clone(), effort, }; - let tool_defs = build_tool_defs(ctx, provider.as_ref(), Some(&spec.tools), false); + let tool_defs = build_tool_defs(ctx, provider.as_ref(), Some(&spec.tools), false, false); let mut conversation = Conversation::new(); conversation.push( Message::text( @@ -152,6 +152,7 @@ pub(crate) async fn run_delegation( tool_defs: &tool_defs, cwd: env.cwd, allow_delegate: false, + allow_ask: false, exec_policy: crate::agent::tighter(&env.exec_policy, &spec.exec_policy), }; let child_token = token.child_token(); diff --git a/crates/goat-agent/src/lib.rs b/crates/goat-agent/src/lib.rs index 5ea8b49..9e63d0c 100644 --- a/crates/goat-agent/src/lib.rs +++ b/crates/goat-agent/src/lib.rs @@ -27,6 +27,8 @@ mod conversation; mod delegate; mod instructions; mod persist; +mod process; +mod process_tools; mod prompt; mod rate_limit_cache; mod retry; @@ -45,6 +47,11 @@ pub async fn model_list_entries(credentials: &CredentialStore) -> Vec) { + while wake_rx.try_recv().is_ok() {} +} pub struct GoatAgent { registry: Registry, @@ -109,6 +116,8 @@ pub(crate) struct Ctx<'a> { pub(crate) instructions: Option<&'a str>, pub(crate) semaphore: &'a Arc, pub(crate) child_ids: &'a AtomicU64, + pub(crate) wake_ids: &'a AtomicU64, + pub(crate) processes: &'a Arc, pub(crate) asks: &'a Mutex>>>, pub(crate) rl_cache: &'a std::sync::Mutex, pub(crate) rl_path: Option<&'a std::path::Path>, @@ -205,6 +214,7 @@ pub(crate) struct LoopEnv<'a> { pub(crate) tool_defs: &'a [ToolDefinition], pub(crate) cwd: &'a Path, pub(crate) allow_delegate: bool, + pub(crate) allow_ask: bool, pub(crate) exec_policy: SandboxPolicy, } @@ -237,6 +247,9 @@ async fn run(agent: GoatAgent, mut ops: mpsc::Receiver, events: mpsc::Sender let session_date = prompt::current_utc_date(); let semaphore = Arc::new(Semaphore::new(delegate::MAX_CONCURRENT_AGENTS)); let child_ids = AtomicU64::new(CHILD_ID_BASE); + let wake_ids = AtomicU64::new(WAKE_ID_BASE); + let (wake_tx, mut wake_rx) = mpsc::channel::(256); + let processes = process::ProcessRegistry::new(events.clone(), wake_tx, Some(store.clone())); let asks: Mutex>>> = Mutex::new(HashMap::new()); let _ = events .send(Event::SkillsChanged { @@ -280,6 +293,8 @@ async fn run(agent: GoatAgent, mut ops: mpsc::Receiver, events: mpsc::Sender instructions: project_instructions.as_deref(), semaphore: &semaphore, child_ids: &child_ids, + wake_ids: &wake_ids, + processes: &processes, asks: &asks, rl_cache: &rl_cache, rl_path: rl_path.as_deref(), @@ -289,7 +304,24 @@ async fn run(agent: GoatAgent, mut ops: mpsc::Receiver, events: mpsc::Sender }; } - while let Some(op) = ops.recv().await { + loop { + let op = tokio::select! { + biased; + maybe_op = ops.recv() => match maybe_op { + Some(op) => op, + None => break, + }, + Some(_wake) = wake_rx.recv() => { + let ctx = ctx!(); + drain_extra_wakes(&mut wake_rx); + if let Flow::Shutdown = + turn::handle_wake(&ctx, &mut state, &mut ops).await + { + break; + } + continue; + } + }; match op { Op::SubmitMessage { id, @@ -307,6 +339,12 @@ async fn run(agent: GoatAgent, mut ops: mpsc::Receiver, events: mpsc::Sender } Op::Interrupt { .. } | Op::Answer { .. } | Op::DequeueMessage { .. } | Op::Clear {} => { } + Op::ProcessKill { process } => { + let _ = processes.kill(process).await; + } + Op::ProcessWatch { process, on } => { + let _ = processes.set_watch(process, on).await; + } Op::Compact { id, instructions } => { let ctx = ctx!(); if let Flow::Shutdown = @@ -331,6 +369,7 @@ async fn run(agent: GoatAgent, mut ops: mpsc::Receiver, events: mpsc::Sender state.thread_id, &mut state.target, &events, + &processes, ) .await; } @@ -414,6 +453,7 @@ async fn run(agent: GoatAgent, mut ops: mpsc::Receiver, events: mpsc::Sender Op::Shutdown {} => break, } } + processes.shutdown_all().await; mcp.shutdown().await; } diff --git a/crates/goat-agent/src/process.rs b/crates/goat-agent/src/process.rs new file mode 100644 index 0000000..9719e89 --- /dev/null +++ b/crates/goat-agent/src/process.rs @@ -0,0 +1,724 @@ +use std::{collections::HashMap, fmt::Write as _, path::Path, process::Stdio, sync::Arc}; + +use goat_protocol::{Event, ProcessExitReason, ProcessId, ProcessInfo, ProcessState}; +use tokio::{ + io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, + process::Command, + sync::{Mutex, mpsc}, +}; + +const RING_CAPACITY: usize = 2000; +const MAX_LIVE_PROCESSES: usize = 16; +const WATCH_FLOOD_LINES: usize = 500; + +pub(crate) struct Wake; + +struct Line { + stream: Stream, + text: String, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Stream { + Out, + Err, +} + +struct Entry { + command: String, + pgid: Option, + db_id: Option, + lines: std::collections::VecDeque, + dropped: usize, + read_cursor: usize, + watch_cursor: usize, + total: usize, + state: ProcessState, + exit_code: Option, + exit_observed: bool, + watched: bool, + watch_flooded: bool, + stdin: Option, + kill_pending: bool, +} + +impl Entry { + fn info(&self, id: ProcessId) -> ProcessInfo { + ProcessInfo { + id, + command: self.command.clone(), + state: self.state, + watched: self.watched, + exit_code: self.exit_code, + } + } +} + +struct Inner { + entries: HashMap, + next_id: u64, +} + +pub(crate) struct ProcessRegistry { + inner: Mutex, + events: mpsc::Sender, + wake_tx: mpsc::Sender, + store: Option, +} + +#[derive(Debug)] +pub(crate) enum SpawnError { + TooMany, + Spawn(String), +} + +impl std::fmt::Display for SpawnError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::TooMany => write!( + f, + "too many background processes are already running (limit {MAX_LIVE_PROCESSES}); stop one with ProcessKill first" + ), + Self::Spawn(msg) => write!(f, "failed to start process: {msg}"), + } + } +} + +pub(crate) struct Started { + pub(crate) id: ProcessId, + pub(crate) pgid: Option, +} + +impl ProcessRegistry { + pub(crate) fn new( + events: mpsc::Sender, + wake_tx: mpsc::Sender, + store: Option, + ) -> Arc { + Arc::new(Self { + inner: Mutex::new(Inner { + entries: HashMap::new(), + next_id: 1, + }), + events, + wake_tx, + store, + }) + } + + pub(crate) async fn spawn( + self: &Arc, + command: &str, + cwd: &Path, + watched: bool, + ) -> Result { + let id = { + let inner = self.inner.lock().await; + let live = inner + .entries + .values() + .filter(|e| e.state == ProcessState::Running) + .count(); + if live >= MAX_LIVE_PROCESSES { + return Err(SpawnError::TooMany); + } + ProcessId(inner.next_id) + }; + + let mut builder = shell_command(command); + builder + .current_dir(cwd) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + set_process_group(&mut builder); + + let mut child = builder + .spawn() + .map_err(|err| SpawnError::Spawn(err.to_string()))?; + + let pgid = child.id().and_then(|pid| i32::try_from(pid).ok()); + + let stdout = child.stdout.take(); + let stderr = child.stderr.take(); + let stdin = child.stdin.take(); + + { + let mut inner = self.inner.lock().await; + inner.next_id += 1; + inner.entries.insert( + id, + Entry { + command: command.to_owned(), + pgid, + db_id: None, + lines: std::collections::VecDeque::new(), + dropped: 0, + read_cursor: 0, + watch_cursor: 0, + total: 0, + state: ProcessState::Running, + exit_code: None, + exit_observed: false, + watched, + watch_flooded: false, + stdin, + kill_pending: false, + }, + ); + } + + let _ = self + .events + .send(Event::ProcessStarted { + process: id, + command: command.to_owned(), + watched, + }) + .await; + self.broadcast_list().await; + + if let Some(pipe) = stdout { + self.spawn_reader(id, pipe, Stream::Out); + } + if let Some(pipe) = stderr { + self.spawn_reader(id, pipe, Stream::Err); + } + self.spawn_waiter(id, child); + + Ok(Started { id, pgid }) + } + + pub(crate) async fn set_db_id(&self, id: ProcessId, db_id: i64) { + let mut inner = self.inner.lock().await; + if let Some(entry) = inner.entries.get_mut(&id) { + entry.db_id = Some(db_id); + } + } + + fn spawn_reader(self: &Arc, id: ProcessId, pipe: R, stream: Stream) + where + R: tokio::io::AsyncRead + Unpin + Send + 'static, + { + let registry = Arc::clone(self); + tokio::spawn(async move { + let mut reader = BufReader::new(pipe).lines(); + while let Ok(Some(line)) = reader.next_line().await { + registry.append_line(id, stream, line).await; + } + }); + } + + fn spawn_waiter(self: &Arc, id: ProcessId, mut child: tokio::process::Child) { + let registry = Arc::clone(self); + tokio::spawn(async move { + let status = child.wait().await; + let code = status.ok().and_then(|s| s.code()); + registry + .mark_exited(id, code, ProcessExitReason::Natural) + .await; + }); + } + + async fn append_line(self: &Arc, id: ProcessId, stream: Stream, text: String) { + let should_wake = { + let mut inner = self.inner.lock().await; + let Some(entry) = inner.entries.get_mut(&id) else { + return; + }; + if entry.lines.len() >= RING_CAPACITY { + entry.lines.pop_front(); + entry.dropped += 1; + if entry.read_cursor > 0 { + entry.read_cursor -= 1; + } + if entry.watch_cursor > 0 { + entry.watch_cursor -= 1; + } + } + entry.lines.push_back(Line { + stream, + text: text.clone(), + }); + entry.total += 1; + if entry.watched && !entry.watch_flooded { + let pending = entry.lines.len() - entry.watch_cursor; + if pending > WATCH_FLOOD_LINES { + entry.watched = false; + entry.watch_flooded = true; + } + } + entry.watched + }; + let _ = self + .events + .send(Event::ProcessOutput { + process: id, + chunk: text, + }) + .await; + if should_wake { + let _ = self.wake_tx.send(Wake).await; + } + } + + async fn mark_exited( + self: &Arc, + id: ProcessId, + code: Option, + natural: ProcessExitReason, + ) { + let (watched, reason, db_id) = { + let mut inner = self.inner.lock().await; + let Some(entry) = inner.entries.get_mut(&id) else { + return; + }; + if entry.state == ProcessState::Exited { + return; + } + entry.state = ProcessState::Exited; + entry.exit_code = code; + let reason = if entry.kill_pending { + ProcessExitReason::Killed + } else { + natural + }; + (entry.watched, reason, entry.db_id) + }; + if let (Some(store), Some(db_id)) = (self.store.as_ref(), db_id) { + let _ = store.finish_process(db_id, now_ms()).await; + } + let _ = self + .events + .send(Event::ProcessExited { + process: id, + code, + reason, + }) + .await; + self.broadcast_list().await; + if watched { + let _ = self.wake_tx.send(Wake).await; + } + } + + pub(crate) async fn read_new(&self, id: ProcessId) -> Option { + let mut inner = self.inner.lock().await; + let entry = inner.entries.get_mut(&id)?; + let chunk = collect_from(entry, entry.read_cursor); + entry.read_cursor = entry.lines.len(); + Some(ReadChunk { + text: chunk, + state: entry.state, + exit_code: entry.exit_code, + }) + } + + pub(crate) async fn take_pending_observations(&self) -> Vec<(ProcessId, Observation)> { + let mut inner = self.inner.lock().await; + let ids: Vec = inner.entries.keys().copied().collect(); + let mut out = Vec::new(); + for id in ids { + let Some(entry) = inner.entries.get_mut(&id) else { + continue; + }; + if !entry.watched { + continue; + } + let has_new = entry.lines.len() > entry.watch_cursor; + let exited_unseen = entry.state == ProcessState::Exited && !entry.exit_observed; + if !has_new && !exited_unseen { + continue; + } + let text = collect_from(entry, entry.watch_cursor); + entry.watch_cursor = entry.lines.len(); + entry.exit_observed = true; + out.push(( + id, + Observation { + command: entry.command.clone(), + output: text, + state: entry.state, + exit_code: entry.exit_code, + }, + )); + } + out.sort_by_key(|(id, _)| id.0); + out + } + + pub(crate) async fn write_stdin(&self, id: ProcessId, text: &str) -> Result<(), String> { + let mut stdin = { + let mut inner = self.inner.lock().await; + let entry = inner + .entries + .get_mut(&id) + .ok_or_else(|| format!("no process #{id}"))?; + if entry.state == ProcessState::Exited { + return Err(format!( + "process #{id} has exited; start it again with ProcessStart" + )); + } + entry + .stdin + .take() + .ok_or_else(|| format!("process #{id} does not accept input"))? + }; + let write = async { + stdin.write_all(text.as_bytes()).await?; + stdin.flush().await + }; + let result = write.await; + let mut inner = self.inner.lock().await; + if let Some(entry) = inner.entries.get_mut(&id) { + entry.stdin = Some(stdin); + } + result.map_err(|err| format!("failed to write to process #{id}: {err}")) + } + + pub(crate) async fn set_watch(&self, id: ProcessId, on: bool) -> Result<(), String> { + { + let mut inner = self.inner.lock().await; + let entry = inner + .entries + .get_mut(&id) + .ok_or_else(|| format!("no process #{id}"))?; + entry.watched = on; + if on { + entry.watch_flooded = false; + entry.watch_cursor = entry.lines.len(); + } + } + self.broadcast_list().await; + Ok(()) + } + + pub(crate) async fn kill(&self, id: ProcessId) -> Result<(), String> { + let pgid = { + let mut inner = self.inner.lock().await; + let entry = inner + .entries + .get_mut(&id) + .ok_or_else(|| format!("no process #{id}"))?; + if entry.state == ProcessState::Exited { + return Ok(()); + } + entry.kill_pending = true; + entry.pgid + }; + kill_group(pgid); + Ok(()) + } + + pub(crate) async fn list(&self) -> Vec { + let inner = self.inner.lock().await; + collect_infos(&inner) + } + + async fn broadcast_list(&self) { + let processes = { + let inner = self.inner.lock().await; + collect_infos(&inner) + }; + let _ = self + .events + .send(Event::ProcessListChanged { processes }) + .await; + } + + pub(crate) async fn shutdown_all(&self) { + let pgids: Vec> = { + let inner = self.inner.lock().await; + inner + .entries + .values() + .filter(|e| e.state == ProcessState::Running) + .map(|e| e.pgid) + .collect() + }; + for pgid in pgids { + kill_group(pgid); + } + } +} + +pub(crate) struct ReadChunk { + pub(crate) text: String, + pub(crate) state: ProcessState, + pub(crate) exit_code: Option, +} + +pub(crate) struct Observation { + pub(crate) command: String, + pub(crate) output: String, + pub(crate) state: ProcessState, + pub(crate) exit_code: Option, +} + +fn collect_from(entry: &Entry, cursor: usize) -> String { + let mut out = String::new(); + if cursor == 0 && entry.dropped > 0 { + let _ = writeln!(out, "[{} earlier lines dropped]", entry.dropped); + } + for line in entry.lines.iter().skip(cursor) { + if line.stream == Stream::Err { + out.push_str("[err] "); + } + out.push_str(&line.text); + out.push('\n'); + } + out +} + +fn collect_infos(inner: &Inner) -> Vec { + let mut infos: Vec = inner + .entries + .iter() + .map(|(id, entry)| entry.info(*id)) + .collect(); + infos.sort_by_key(|i| i.id.0); + infos +} + +fn shell_command(command: &str) -> Command { + #[cfg(windows)] + { + let mut builder = Command::new("cmd"); + builder.arg("/C").arg(command); + builder + } + #[cfg(not(windows))] + { + let mut builder = Command::new("sh"); + builder.arg("-c").arg(command); + builder + } +} + +fn set_process_group(builder: &mut Command) { + #[cfg(unix)] + builder.process_group(0); + #[cfg(windows)] + { + const CREATE_NEW_PROCESS_GROUP: u32 = 0x0000_0200; + builder.creation_flags(CREATE_NEW_PROCESS_GROUP); + } + #[cfg(not(any(unix, windows)))] + let _ = builder; +} + +fn kill_group(pgid: Option) { + let Some(pgid) = pgid else { + return; + }; + #[cfg(windows)] + { + let _ = std::process::Command::new("taskkill") + .arg("/F") + .arg("/T") + .arg("/PID") + .arg(pgid.to_string()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } + #[cfg(not(windows))] + { + let _ = std::process::Command::new("kill") + .arg("-KILL") + .arg(format!("-{pgid}")) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } +} + +fn now_ms() -> i64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX)) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use super::{ProcessRegistry, Wake}; + use goat_protocol::{Event, ProcessState}; + use std::time::Duration; + use tokio::sync::mpsc; + + #[cfg(not(windows))] + mod plat { + pub const TWO_ECHOES: &str = "echo hello; echo world"; + pub const ECHO_ONE: &str = "echo one"; + pub const ECHO_STDERR: &str = "echo oops 1>&2"; + pub const ECHO_PING: &str = "echo ping"; + pub const ECHO_QUIET: &str = "echo quiet"; + pub const SLEEP_LONG: &str = "sleep 30"; + pub const CAT: &str = "cat"; + pub const TRUE: &str = "true"; + } + + #[cfg(windows)] + mod plat { + pub const TWO_ECHOES: &str = "echo hello& echo world"; + pub const ECHO_ONE: &str = "echo one"; + pub const ECHO_STDERR: &str = "echo oops 1>&2"; + pub const ECHO_PING: &str = "echo ping"; + pub const ECHO_QUIET: &str = "echo quiet"; + pub const SLEEP_LONG: &str = "ping -n 31 127.0.0.1 >nul"; + pub const CAT: &str = "findstr \"^\""; + pub const TRUE: &str = "type nul"; + } + + fn harness() -> ( + std::sync::Arc, + mpsc::Receiver, + mpsc::Receiver, + ) { + let (event_tx, event_rx) = mpsc::channel(256); + let (wake_tx, wake_rx) = mpsc::channel(256); + let registry = ProcessRegistry::new(event_tx, wake_tx, None); + (registry, event_rx, wake_rx) + } + + async fn wait_until_exited(registry: &ProcessRegistry, id: goat_protocol::ProcessId) { + for _ in 0..200 { + let list = registry.list().await; + if list + .iter() + .any(|p| p.id == id && p.state == ProcessState::Exited) + { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("process did not exit in time"); + } + + #[tokio::test] + async fn spawn_reads_output_and_exits() { + let (registry, _events, _wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry + .spawn(plat::TWO_ECHOES, &cwd, false) + .await + .unwrap_or_else(|e| panic!("spawn failed: {e}")); + wait_until_exited(®istry, started.id).await; + let chunk = registry.read_new(started.id).await.unwrap(); + assert!(chunk.text.contains("hello"), "got: {}", chunk.text); + assert!(chunk.text.contains("world"), "got: {}", chunk.text); + assert_eq!(chunk.state, ProcessState::Exited); + assert_eq!(chunk.exit_code, Some(0)); + } + + #[tokio::test] + async fn read_new_is_cursor_based() { + let (registry, _events, _wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::ECHO_ONE, &cwd, false).await.unwrap(); + wait_until_exited(®istry, started.id).await; + let first = registry.read_new(started.id).await.unwrap(); + assert!(first.text.contains("one")); + let second = registry.read_new(started.id).await.unwrap(); + assert!( + !second.text.contains("one"), + "second read should be empty of old output" + ); + } + + #[tokio::test] + async fn stderr_is_tagged() { + let (registry, _events, _wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry + .spawn(plat::ECHO_STDERR, &cwd, false) + .await + .unwrap(); + wait_until_exited(®istry, started.id).await; + let chunk = registry.read_new(started.id).await.unwrap(); + assert!(chunk.text.contains("[err] oops"), "got: {}", chunk.text); + } + + #[tokio::test] + async fn watched_process_wakes_on_output() { + let (registry, _events, mut wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::ECHO_PING, &cwd, true).await.unwrap(); + let _woke = tokio::time::timeout(Duration::from_secs(5), wake.recv()) + .await + .expect("should wake") + .expect("wake channel open"); + let obs = registry.take_pending_observations().await; + assert!( + obs.iter() + .any(|(id, o)| *id == started.id && o.output.contains("ping")), + "got: {obs:?}", + obs = obs + .iter() + .map(|(_, o)| o.output.clone()) + .collect::>() + ); + } + + #[tokio::test] + async fn unwatched_process_does_not_wake() { + let (registry, _events, mut wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::ECHO_QUIET, &cwd, false).await.unwrap(); + wait_until_exited(®istry, started.id).await; + let result = tokio::time::timeout(Duration::from_millis(200), wake.recv()).await; + assert!(result.is_err(), "unwatched process must not wake the agent"); + } + + #[tokio::test] + async fn kill_terminates_running_process() { + let (registry, _events, _wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::SLEEP_LONG, &cwd, false).await.unwrap(); + let running = registry.list().await; + assert_eq!(running[0].state, ProcessState::Running); + registry.kill(started.id).await.unwrap(); + wait_until_exited(®istry, started.id).await; + } + + #[tokio::test] + async fn stdin_write_reaches_process() { + let (registry, _events, _wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::CAT, &cwd, false).await.unwrap(); + registry.write_stdin(started.id, "typed\n").await.unwrap(); + for _ in 0..200 { + let chunk = registry.read_new(started.id).await.unwrap(); + if chunk.text.contains("typed") { + registry.kill(started.id).await.unwrap(); + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("stdin was not echoed back"); + } + + #[tokio::test] + async fn write_to_exited_process_errors() { + let (registry, _events, _wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::TRUE, &cwd, false).await.unwrap(); + wait_until_exited(®istry, started.id).await; + let result = registry.write_stdin(started.id, "x\n").await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn watch_can_be_toggled() { + let (registry, _events, mut wake) = harness(); + let cwd = std::env::temp_dir(); + let started = registry.spawn(plat::SLEEP_LONG, &cwd, false).await.unwrap(); + registry.set_watch(started.id, true).await.unwrap(); + registry.write_stdin(started.id, "").await.ok(); + registry.set_watch(started.id, false).await.unwrap(); + registry.kill(started.id).await.unwrap(); + let _ = wake.try_recv(); + } +} diff --git a/crates/goat-agent/src/process_tools.rs b/crates/goat-agent/src/process_tools.rs new file mode 100644 index 0000000..2ca8b3f --- /dev/null +++ b/crates/goat-agent/src/process_tools.rs @@ -0,0 +1,322 @@ +use std::fmt::Write as _; + +use goat_tool::{SandboxPolicy, ToolOutput}; +use serde::Deserialize; +use tokio_util::sync::CancellationToken; + +use crate::{Ctx, LoopEnv}; + +pub(crate) const PROCESS_START_TOOL_NAME: &str = "ProcessStart"; +pub(crate) const PROCESS_OUTPUT_TOOL_NAME: &str = "ProcessOutput"; +pub(crate) const PROCESS_INPUT_TOOL_NAME: &str = "ProcessInput"; +pub(crate) const PROCESS_KILL_TOOL_NAME: &str = "ProcessKill"; +pub(crate) const PROCESS_LIST_TOOL_NAME: &str = "ProcessList"; +pub(crate) const PROCESS_WATCH_TOOL_NAME: &str = "ProcessWatch"; + +pub(crate) fn is_process_tool(name: &str) -> bool { + matches!( + name, + PROCESS_START_TOOL_NAME + | PROCESS_OUTPUT_TOOL_NAME + | PROCESS_INPUT_TOOL_NAME + | PROCESS_KILL_TOOL_NAME + | PROCESS_LIST_TOOL_NAME + | PROCESS_WATCH_TOOL_NAME + ) +} + +pub(crate) fn tool_defs() -> Vec { + vec![ + def( + PROCESS_START_TOOL_NAME, + "Start a long-running command in the background and return immediately with a process id. Use this for dev servers (pnpm dev, vite), watchers, or a poller that waits for a long task (e.g. `gh run watch`) instead of blocking on Bash. Output is buffered; read it later with ProcessOutput. Set watch=true to be woken when the process prints a matching line or exits (pipe the command through grep to keep events meaningful). The process keeps running across turns until it exits or you call ProcessKill.", + serde_json::json!({ + "type": "object", + "properties": { + "command": {"type": "string", "description": "shell command to run in the background"}, + "watch": {"type": "boolean", "description": "wake the agent on new output and on exit (default false)"} + }, + "required": ["command"] + }), + ), + def( + PROCESS_OUTPUT_TOOL_NAME, + "Read output produced by a background process since the last read (a moving cursor, not the whole history). Returns whether the process is still running or has exited with its code.", + serde_json::json!({ + "type": "object", + "properties": {"process": {"type": "string", "description": "process id from ProcessStart"}}, + "required": ["process"] + }), + ), + def( + PROCESS_INPUT_TOOL_NAME, + "Send keystrokes to a background process's stdin (e.g. answer an interactive prompt). Include a trailing newline to submit a line.", + serde_json::json!({ + "type": "object", + "properties": { + "process": {"type": "string", "description": "process id from ProcessStart"}, + "text": {"type": "string", "description": "raw bytes to write to stdin"} + }, + "required": ["process", "text"] + }), + ), + def( + PROCESS_KILL_TOOL_NAME, + "Terminate a background process (and its process group).", + serde_json::json!({ + "type": "object", + "properties": {"process": {"type": "string", "description": "process id from ProcessStart"}}, + "required": ["process"] + }), + ), + def( + PROCESS_LIST_TOOL_NAME, + "List background processes and their state (running or exited).", + serde_json::json!({"type": "object", "properties": {}}), + ), + def( + PROCESS_WATCH_TOOL_NAME, + "Turn push observation on or off for a background process. When on, new output and exit wake the agent; when off, output is only buffered for ProcessOutput.", + serde_json::json!({ + "type": "object", + "properties": { + "process": {"type": "string", "description": "process id from ProcessStart"}, + "on": {"type": "boolean", "description": "true to watch, false to stop watching"} + }, + "required": ["process", "on"] + }), + ), + ] +} + +fn def(name: &str, description: &str, schema: serde_json::Value) -> goat_provider::ToolDefinition { + goat_provider::ToolDefinition { + name: name.to_owned(), + description: description.to_owned(), + input_schema: schema, + } +} + +pub(crate) fn call_display(name: &str, input: &str) -> goat_protocol::ToolDisplay { + let detail = process_id_arg(input).map(|p| format!("#{p}")); + match (name, detail) { + (PROCESS_START_TOOL_NAME, _) => { + let cmd = serde_json::from_str::(input) + .map(|i| i.command) + .unwrap_or_default(); + goat_protocol::ToolDisplay::primary(format!("ProcessStart({})", flatten(&cmd))) + } + (_, Some(detail)) => goat_protocol::ToolDisplay::primary(format!("{name}({detail})")), + (_, None) => goat_protocol::ToolDisplay::primary(name.to_owned()), + } +} + +fn flatten(text: &str) -> String { + let flat = text.split_whitespace().collect::>().join(" "); + if flat.chars().count() > 60 { + let head: String = flat.chars().take(60).collect(); + format!("{head}…") + } else { + flat + } +} + +fn process_id_arg(input: &str) -> Option { + let value: serde_json::Value = serde_json::from_str(input).ok()?; + let raw = value.get("process")?; + raw.as_str() + .and_then(|s| s.parse().ok()) + .or_else(|| raw.as_u64()) +} + +#[derive(Deserialize)] +struct StartInput { + command: String, + #[serde(default)] + watch: bool, +} + +#[derive(Deserialize)] +struct ProcessRef { + process: goat_protocol::ProcessId, +} + +#[derive(Deserialize)] +struct InputArgs { + process: goat_protocol::ProcessId, + text: String, +} + +#[derive(Deserialize)] +struct WatchArgs { + process: goat_protocol::ProcessId, + on: bool, +} + +pub(crate) async fn run_process_tool( + ctx: &Ctx<'_>, + env: &LoopEnv<'_>, + name: &str, + input_json: &str, + token: &CancellationToken, +) -> Option> { + if token.is_cancelled() { + return None; + } + let result = match name { + PROCESS_START_TOOL_NAME => start(ctx, env, input_json).await, + PROCESS_OUTPUT_TOOL_NAME => output(ctx, input_json).await, + PROCESS_INPUT_TOOL_NAME => input(ctx, input_json).await, + PROCESS_KILL_TOOL_NAME => kill(ctx, input_json).await, + PROCESS_LIST_TOOL_NAME => Ok(list(ctx).await), + PROCESS_WATCH_TOOL_NAME => watch(ctx, input_json).await, + _ => Err(format!("unknown process tool: {name}")), + }; + Some(result) +} + +async fn start(ctx: &Ctx<'_>, env: &LoopEnv<'_>, input_json: &str) -> Result { + if !matches!(env.exec_policy, SandboxPolicy::Full) { + return Err( + "background processes are only available with full shell access, not while planning" + .to_owned(), + ); + } + let args: StartInput = + serde_json::from_str(input_json).map_err(|err| format!("invalid input: {err}"))?; + let started = ctx + .processes + .spawn(&args.command, env.cwd, args.watch) + .await + .map_err(|err| err.to_string())?; + if let Some(pgid) = started.pgid { + let db_id = ctx + .store + .create_process(goat_store::NewProcess { + pgid: i64::from(pgid), + command: args.command.clone(), + cwd: env.cwd.display().to_string(), + started_at: now_ms(), + }) + .await + .ok(); + if let Some(db_id) = db_id { + ctx.processes.set_db_id(started.id, db_id).await; + } + } + let id = started.id; + let watched = if args.watch { " (watched)" } else { "" }; + Ok(ToolOutput::text(format!( + "Started process #{id}{watched}. Read output with ProcessOutput(process={id}); stop with ProcessKill(process={id})." + )) + .with_summary(format!("#{id} {}", flatten(&args.command)))) +} + +async fn output(ctx: &Ctx<'_>, input_json: &str) -> Result { + let args: ProcessRef = + serde_json::from_str(input_json).map_err(|err| format!("invalid input: {err}"))?; + let chunk = ctx + .processes + .read_new(args.process) + .await + .ok_or_else(|| format!("no process #{}", args.process))?; + let status = match chunk.state { + goat_protocol::ProcessState::Running => "running".to_owned(), + goat_protocol::ProcessState::Exited => match chunk.exit_code { + Some(code) => format!("exited (code {code})"), + None => "exited".to_owned(), + }, + }; + let body = if chunk.text.trim().is_empty() { + format!("[no new output] process #{} is {status}", args.process) + } else { + format!( + "{}\n[process #{} is {status}]", + chunk.text.trim_end(), + args.process + ) + }; + Ok(ToolOutput::text(crate::tools_exec::cap_tool_result(body))) +} + +async fn input(ctx: &Ctx<'_>, input_json: &str) -> Result { + let args: InputArgs = + serde_json::from_str(input_json).map_err(|err| format!("invalid input: {err}"))?; + ctx.processes.write_stdin(args.process, &args.text).await?; + Ok(ToolOutput::text(format!( + "Wrote to process #{}.", + args.process + ))) +} + +async fn kill(ctx: &Ctx<'_>, input_json: &str) -> Result { + let args: ProcessRef = + serde_json::from_str(input_json).map_err(|err| format!("invalid input: {err}"))?; + ctx.processes.kill(args.process).await?; + Ok(ToolOutput::text(format!( + "Killed process #{}.", + args.process + ))) +} + +async fn watch(ctx: &Ctx<'_>, input_json: &str) -> Result { + let args: WatchArgs = + serde_json::from_str(input_json).map_err(|err| format!("invalid input: {err}"))?; + ctx.processes.set_watch(args.process, args.on).await?; + let state = if args.on { "watching" } else { "not watching" }; + Ok(ToolOutput::text(format!( + "Now {state} process #{}.", + args.process + ))) +} + +async fn list(ctx: &Ctx<'_>) -> ToolOutput { + let processes = ctx.processes.list().await; + if processes.is_empty() { + return ToolOutput::text("No background processes.".to_owned()); + } + let mut out = String::from("Background processes:\n"); + for p in &processes { + let state = match p.state { + goat_protocol::ProcessState::Running => "running".to_owned(), + goat_protocol::ProcessState::Exited => match p.exit_code { + Some(code) => format!("exited({code})"), + None => "exited".to_owned(), + }, + }; + let watched = if p.watched { " watched" } else { "" }; + let _ = writeln!(out, " #{} [{state}{watched}] {}", p.id, p.command); + } + ToolOutput::text(out) +} + +pub(crate) async fn roster_message(ctx: &Ctx<'_>) -> Option { + let processes = ctx.processes.list().await; + let running: Vec<_> = processes + .iter() + .filter(|p| p.state == goat_protocol::ProcessState::Running) + .collect(); + if running.is_empty() { + return None; + } + let mut text = String::from( + "\nProcesses running now (read with ProcessOutput, stop with ProcessKill):\n", + ); + for p in running { + let watched = if p.watched { " watched" } else { "" }; + let _ = writeln!(text, " #{}{watched} — {}", p.id, p.command); + } + text.push_str(""); + Some(goat_provider::Message::text( + goat_provider::MessageRole::User, + text, + )) +} + +fn now_ms() -> i64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX)) + .unwrap_or_default() +} diff --git a/crates/goat-agent/src/rounds.rs b/crates/goat-agent/src/rounds.rs index 9c7f0d2..4cb3ca9 100644 --- a/crates/goat-agent/src/rounds.rs +++ b/crates/goat-agent/src/rounds.rs @@ -428,8 +428,18 @@ pub(crate) async fn core_loop( } } } - let round = - crate::retry::run_round_with_retry(ctx, run, env, conversation.messages(), token).await; + let roster = if run.is_top() { + crate::process_tools::roster_message(ctx).await + } else { + None + }; + let round = if let Some(roster) = roster { + let mut messages = conversation.messages().to_vec(); + messages.push(roster); + crate::retry::run_round_with_retry(ctx, run, env, &messages, token).await + } else { + crate::retry::run_round_with_retry(ctx, run, env, conversation.messages(), token).await + }; match &round.end { RoundEnd::Cancelled => return LoopOutcome::Cancelled, RoundEnd::Failed(StreamError::ContextOverflow { .. }) if !compacted_for_overflow => { diff --git a/crates/goat-agent/src/tools_exec.rs b/crates/goat-agent/src/tools_exec.rs index ebdfbdc..a3d02e5 100644 --- a/crates/goat-agent/src/tools_exec.rs +++ b/crates/goat-agent/src/tools_exec.rs @@ -65,6 +65,9 @@ pub(crate) fn call_display(tools: &ToolRegistry, name: &str, input: &str) -> Too AGENT_TOOL_NAME => agent_call_display(input), ASK_TOOL_NAME => ask_call_display(input), WEB_SEARCH_TOOL_NAME => web_search_display(input), + _ if crate::process_tools::is_process_tool(name) => { + crate::process_tools::call_display(name, input) + } _ => tools.get(name).map_or_else( || goat_tool::display::generic_named(name, input), |tool| tool.display_input(input), @@ -135,7 +138,9 @@ async fn execute_tool( .await .map(ToolOutput::text), ) - } else if prep.name == ASK_TOOL_NAME && env.allow_delegate { + } else if crate::process_tools::is_process_tool(prep.name) && env.allow_delegate { + crate::process_tools::run_process_tool(ctx, env, prep.name, prep.input_json, token).await + } else if prep.name == ASK_TOOL_NAME && env.allow_ask { Some(run_ask(ctx, run, prep.input_json, ToolCallId(prep.tui_id), token).await) } else if prep.name == AGENT_TOOL_NAME && env.allow_delegate { let permit = tokio::select! { @@ -272,6 +277,7 @@ pub(crate) fn build_tool_defs( provider: &dyn Provider, selection: Option<&ToolSelection>, allow_delegate: bool, + allow_ask: bool, ) -> Vec { if !provider.capabilities().tools { return Vec::new(); @@ -294,6 +300,9 @@ pub(crate) fn build_tool_defs( if !ctx.agents.is_empty() { defs.push(agent_tool_def(ctx)); } + defs.extend(crate::process_tools::tool_defs()); + } + if allow_ask { defs.push(ask_tool_def()); } defs diff --git a/crates/goat-agent/src/turn.rs b/crates/goat-agent/src/turn.rs index 785facc..e1726cf 100644 --- a/crates/goat-agent/src/turn.rs +++ b/crates/goat-agent/src/turn.rs @@ -1,3 +1,5 @@ +use std::fmt::Write as _; + use goat_protocol::{Event, InputAttachment, ModelTarget, Op, TaskId}; use goat_provider::{ContentBlock, Message, MessageRole, Provider, ToolDefinition}; use goat_store::Store; @@ -39,9 +41,13 @@ pub(crate) fn user_message(text: &str, attachments: &[InputAttachment]) -> Messa } } -fn top_regime(ctx: &Ctx<'_>, provider: &dyn Provider) -> (Vec, SandboxPolicy) { +fn top_regime( + ctx: &Ctx<'_>, + provider: &dyn Provider, + allow_ask: bool, +) -> (Vec, SandboxPolicy) { ( - build_tool_defs(ctx, provider, None, true), + build_tool_defs(ctx, provider, None, true, allow_ask), SandboxPolicy::Full, ) } @@ -111,8 +117,15 @@ pub(crate) async fn handle_idle_op( thread_id: Option, target: &mut Option, events: &mpsc::Sender, + processes: &std::sync::Arc, ) { match op { + Op::ProcessKill { process } => { + let _ = processes.kill(process).await; + } + Op::ProcessWatch { process, on } => { + let _ = processes.set_watch(process, on).await; + } Op::SelectModel { target: chosen } => { if let Some(tid) = thread_id && let Err(err) = store @@ -160,6 +173,56 @@ enum TurnFlow { Shutdown, } +pub(crate) async fn handle_wake( + ctx: &Ctx<'_>, + state: &mut SessionState, + ops: &mut mpsc::Receiver, +) -> Flow { + let observations = ctx.processes.take_pending_observations().await; + if observations.is_empty() { + return Flow::Continue; + } + let mut body = String::from( + "\nBackground processes you are watching produced output or exited. React if needed; otherwise keep waiting.\n", + ); + for (id, obs) in &observations { + let status = match obs.state { + goat_protocol::ProcessState::Running => "running".to_owned(), + goat_protocol::ProcessState::Exited => match obs.exit_code { + Some(code) => format!("exited(code {code})"), + None => "exited".to_owned(), + }, + }; + let _ = write!(body, "\n[process #{id} · {} · {status}]\n", obs.command); + if obs.output.trim().is_empty() { + body.push_str("(no new output)\n"); + } else { + body.push_str(obs.output.trim_end()); + body.push('\n'); + } + } + body.push_str(""); + + let wake_id = TaskId( + ctx.wake_ids + .fetch_add(1, std::sync::atomic::Ordering::Relaxed), + ); + run_turn_chain( + ctx, + crate::UserInput { + id: wake_id, + text: body, + display: Some("(process activity)".to_owned()), + attachments: Vec::new(), + }, + std::collections::VecDeque::new(), + state, + ops, + false, + ) + .await +} + pub(crate) async fn handle_turn( ctx: &Ctx<'_>, id: TaskId, @@ -180,6 +243,7 @@ pub(crate) async fn handle_turn( std::collections::VecDeque::new(), state, ops, + true, ) .await } @@ -190,11 +254,13 @@ async fn run_turn_chain( seed: std::collections::VecDeque, state: &mut SessionState, ops: &mut mpsc::Receiver, + allow_ask: bool, ) -> Flow { let mut next = Some((input, seed)); let mut pending: Vec = Vec::new(); while let Some((turn_input, turn_seed)) = next.take() { - let (flow, deferred) = run_one_turn(ctx, turn_input, turn_seed, state, ops).await; + let (flow, deferred) = + run_one_turn(ctx, turn_input, turn_seed, state, ops, allow_ask).await; pending.extend(deferred); match flow { TurnFlow::Shutdown => return Flow::Shutdown, @@ -238,6 +304,7 @@ async fn drain_deferred( state.thread_id, &mut state.target, ctx.events, + ctx.processes, ) .await; } @@ -361,7 +428,7 @@ pub(crate) async fn handle_shell( ); drop(steering); if let Some(next_input) = captured.pop_front() { - return Box::pin(run_turn_chain(ctx, next_input, captured, state, ops)).await; + return Box::pin(run_turn_chain(ctx, next_input, captured, state, ops, true)).await; } Flow::Continue } @@ -413,7 +480,7 @@ pub(crate) async fn handle_compact( return Flow::Shutdown; } let cwd = resolve_thread_cwd(ctx, state.thread_id).await; - let (tool_defs, exec_policy) = top_regime(ctx, provider.as_ref()); + let (tool_defs, exec_policy) = top_regime(ctx, provider.as_ref(), true); let ids = crate::TurnIds { stored_thread: state.thread_id, turn_db_id: None, @@ -427,6 +494,7 @@ pub(crate) async fn handle_compact( tool_defs: &tool_defs, cwd: &cwd, allow_delegate: true, + allow_ask: true, exec_policy, }; let token = CancellationToken::new(); @@ -530,7 +598,7 @@ pub(crate) async fn handle_compact( ); drop(steering); if let Some(next_input) = captured.pop_front() { - return Box::pin(run_turn_chain(ctx, next_input, captured, state, ops)).await; + return Box::pin(run_turn_chain(ctx, next_input, captured, state, ops, true)).await; } Flow::Continue } @@ -542,6 +610,7 @@ async fn run_one_turn( seed: std::collections::VecDeque, state: &mut SessionState, ops: &mut mpsc::Receiver, + allow_ask: bool, ) -> (TurnFlow, Vec) { let id = input.id; let text = input.text; @@ -612,7 +681,7 @@ async fn run_one_turn( } let cwd = resolve_thread_cwd(ctx, ids.stored_thread).await; - let (tool_defs, exec_policy) = top_regime(ctx, provider.as_ref()); + let (tool_defs, exec_policy) = top_regime(ctx, provider.as_ref(), allow_ask); let steering: crate::SteeringQueue = std::sync::Mutex::new(seed); let run = Run::top(id, &ids, &steering); let env = crate::LoopEnv { @@ -621,6 +690,7 @@ async fn run_one_turn( tool_defs: &tool_defs, cwd: &cwd, allow_delegate: true, + allow_ask, exec_policy, }; let token = CancellationToken::new(); diff --git a/crates/goat-client/src/lib.rs b/crates/goat-client/src/lib.rs index b7228ba..ca4cc52 100644 --- a/crates/goat-client/src/lib.rs +++ b/crates/goat-client/src/lib.rs @@ -233,7 +233,9 @@ async fn run_connection( } Op::Interrupt { .. } | Op::Answer { .. } - | Op::DequeueMessage { .. } => { + | Op::DequeueMessage { .. } + | Op::ProcessKill { .. } + | Op::ProcessWatch { .. } => { let mut op = op; shared.idmap.lock().await.translate_outbound(&mut op); ClientFrame::Control { session, op } diff --git a/crates/goat-daemon/src/lib.rs b/crates/goat-daemon/src/lib.rs index 5b65e39..7827c02 100644 --- a/crates/goat-daemon/src/lib.rs +++ b/crates/goat-daemon/src/lib.rs @@ -34,7 +34,9 @@ pub struct RemoteSettings { pub async fn serve(config: DaemonConfig) -> Result<(), DaemonError> { let listener = bind(&config.socket_path)?; + let db_path = config.db_path.clone(); sweep_orphaned_turns(&config.db_path).await; + sweep_orphaned_processes(&config.db_path).await; let manager = Manager::new(config.auth_path, config.db_path); let shutdown = tokio_util::sync::CancellationToken::new(); tracing::info!(socket = %config.socket_path.display(), "daemon listening"); @@ -49,6 +51,10 @@ pub async fn serve(config: DaemonConfig) -> Result<(), DaemonError> { tracing::info!("daemon shutting down"); break; } + () = shutdown_signal() => { + tracing::info!("received termination signal, shutting down"); + break; + } accepted = listener.accept() => match accepted { Ok(stream) => { let manager = manager.clone(); @@ -62,10 +68,33 @@ pub async fn serve(config: DaemonConfig) -> Result<(), DaemonError> { } } + manager.shutdown_all_sessions().await; + sweep_orphaned_processes(&db_path).await; transport::cleanup(&config.socket_path); Ok(()) } +#[cfg(unix)] +async fn shutdown_signal() { + use tokio::signal::unix::{SignalKind, signal}; + let (Ok(mut term), Ok(mut int)) = ( + signal(SignalKind::terminate()), + signal(SignalKind::interrupt()), + ) else { + std::future::pending::<()>().await; + return; + }; + tokio::select! { + _ = term.recv() => {} + _ = int.recv() => {} + } +} + +#[cfg(not(unix))] +async fn shutdown_signal() { + let _ = tokio::signal::ctrl_c().await; +} + fn spawn_remote( manager: &Manager, shutdown: &tokio_util::sync::CancellationToken, @@ -115,3 +144,49 @@ async fn sweep_orphaned_turns(db_path: &Path) { Err(err) => tracing::warn!(%err, "failed to sweep orphaned turns"), } } + +async fn sweep_orphaned_processes(db_path: &Path) { + let Ok(store) = goat_store::Store::open(db_path) else { + return; + }; + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX)); + match store.take_orphan_processes(now).await { + Ok(orphans) => { + for orphan in &orphans { + kill_process_group(orphan.pgid); + } + if !orphans.is_empty() { + tracing::info!( + count = orphans.len(), + "killed orphaned background processes" + ); + } + } + Err(err) => tracing::warn!(%err, "failed to sweep orphaned processes"), + } +} + +fn kill_process_group(pgid: i64) { + #[cfg(windows)] + { + let _ = std::process::Command::new("taskkill") + .arg("/F") + .arg("/T") + .arg("/PID") + .arg(pgid.to_string()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status(); + } + #[cfg(not(windows))] + if let Ok(pgid) = i32::try_from(pgid) { + let _ = std::process::Command::new("kill") + .arg("-KILL") + .arg(format!("-{pgid}")) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status(); + } +} diff --git a/crates/goat-daemon/src/manager.rs b/crates/goat-daemon/src/manager.rs index 756a6f4..10a48e4 100644 --- a/crates/goat-daemon/src/manager.rs +++ b/crates/goat-daemon/src/manager.rs @@ -201,6 +201,7 @@ impl Manager { snapshot: None, tokens: 0, open_asks: 0, + live_processes: 0, thread_id, awaits_restore: thread_id.is_some(), ready, @@ -369,6 +370,24 @@ impl Manager { tracing::info!(session = session.0, "evicted idle session with no windows"); } + pub(crate) async fn shutdown_all_sessions(&self) { + let all_ops: Vec> = { + let table = self.inner.sessions.lock().await; + let mut ops = Vec::new(); + for live in table.values() { + let inner = live.inner.lock().await; + ops.push(inner.ops.clone()); + } + ops + }; + for ops in &all_ops { + let _ = ops.send(Op::Shutdown {}).await; + } + if !all_ops.is_empty() { + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + } + } + pub(crate) async fn rebind( &self, client: ClientId, diff --git a/crates/goat-daemon/src/session.rs b/crates/goat-daemon/src/session.rs index 10762b5..1e65dec 100644 --- a/crates/goat-daemon/src/session.rs +++ b/crates/goat-daemon/src/session.rs @@ -27,6 +27,7 @@ pub(crate) struct SessionInner { pub(crate) snapshot: Option, pub(crate) tokens: u64, pub(crate) open_asks: usize, + pub(crate) live_processes: usize, pub(crate) thread_id: Option, pub(crate) awaits_restore: bool, pub(crate) ready: Arc, @@ -125,6 +126,10 @@ impl SessionInner { Event::AskDismissed { .. } => { self.open_asks = self.open_asks.saturating_sub(1); } + Event::ProcessStarted { .. } => self.live_processes += 1, + Event::ProcessExited { .. } => { + self.live_processes = self.live_processes.saturating_sub(1); + } Event::Usage { usage, .. } => { self.tokens = self .tokens @@ -225,6 +230,7 @@ impl SessionInner { self.subscribers.is_empty() && self.pending_attaches == 0 && self.open_asks == 0 + && self.live_processes == 0 && matches!(self.state, SessionLiveState::Idle {}) } } @@ -306,6 +312,7 @@ mod tests { snapshot: None, tokens: 0, open_asks: 0, + live_processes: 0, thread_id: None, awaits_restore: false, ready: std::sync::Arc::new(tokio::sync::Notify::new()), @@ -334,6 +341,30 @@ mod tests { assert!(inner.evictable()); } + #[test] + fn live_process_blocks_eviction() { + let mut inner = blank_inner(); + assert!(inner.evictable()); + inner.record_and_fanout(Event::ProcessStarted { + process: goat_protocol::ProcessId(1), + command: "pnpm dev".to_owned(), + watched: false, + }); + assert!( + !inner.evictable(), + "a live background process must keep the session alive after the window closes" + ); + inner.record_and_fanout(Event::ProcessExited { + process: goat_protocol::ProcessId(1), + code: Some(0), + reason: goat_protocol::ProcessExitReason::Natural, + }); + assert!( + inner.evictable(), + "once the process exits the session may be evicted" + ); + } + #[test] fn upsert_replaces_sender_for_same_client() { let mut subs: Vec = Vec::new(); diff --git a/crates/goat-protocol/src/event.rs b/crates/goat-protocol/src/event.rs index 56abc23..65619ac 100644 --- a/crates/goat-protocol/src/event.rs +++ b/crates/goat-protocol/src/event.rs @@ -1,8 +1,9 @@ use serde::{Deserialize, Serialize}; use crate::{ - AccountEntry, InputAttachment, LoginProvider, ModelEntry, ModelTarget, RateLimitSnapshot, - SkillInfo, TaskId, ThreadSummary, ToolCall, ToolCallId, ToolOutcome, TranscriptEntry, Usage, + AccountEntry, InputAttachment, LoginProvider, ModelEntry, ModelTarget, ProcessExitReason, + ProcessId, ProcessInfo, RateLimitSnapshot, SkillInfo, TaskId, ThreadSummary, ToolCall, + ToolCallId, ToolOutcome, TranscriptEntry, Usage, }; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] @@ -157,6 +158,30 @@ pub enum Event { ThreadBound { thread_id: i64, }, + ProcessStarted { + process: ProcessId, + command: String, + watched: bool, + }, + ProcessOutput { + process: ProcessId, + chunk: String, + }, + ProcessExited { + process: ProcessId, + #[serde(default, skip_serializing_if = "Option::is_none")] + code: Option, + reason: ProcessExitReason, + }, + ProcessObserved { + id: TaskId, + process: ProcessId, + command: String, + output: String, + }, + ProcessListChanged { + processes: Vec, + }, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] diff --git a/crates/goat-protocol/src/lib.rs b/crates/goat-protocol/src/lib.rs index 0b8a303..11294bd 100644 --- a/crates/goat-protocol/src/lib.rs +++ b/crates/goat-protocol/src/lib.rs @@ -142,4 +142,70 @@ mod tests { let back: TaskId = serde_json::from_str(&json).unwrap(); assert_eq!(back, big); } + + #[test] + fn process_id_serializes_as_string() { + assert_eq!( + serde_json::to_string(&super::ProcessId(7)).unwrap(), + r#""7""# + ); + } + + #[test] + fn op_process_kill_roundtrips() { + let op = Op::ProcessKill { + process: super::ProcessId(3), + }; + let json = serde_json::to_string(&op).unwrap(); + assert_eq!(json, r#"{"type":"ProcessKill","process":"3"}"#); + let back: Op = serde_json::from_str(&json).unwrap(); + assert_eq!(back, op); + } + + #[test] + fn op_process_watch_roundtrips() { + let op = Op::ProcessWatch { + process: super::ProcessId(4), + on: true, + }; + let json = serde_json::to_string(&op).unwrap(); + let back: Op = serde_json::from_str(&json).unwrap(); + assert_eq!(back, op); + } + + #[test] + fn event_process_started_roundtrips() { + let ev = Event::ProcessStarted { + process: super::ProcessId(1), + command: "pnpm dev".to_owned(), + watched: false, + }; + let json = serde_json::to_string(&ev).unwrap(); + let back: Event = serde_json::from_str(&json).unwrap(); + assert_eq!(back, ev); + } + + #[test] + fn event_process_exited_omits_code_when_absent() { + let ev = Event::ProcessExited { + process: super::ProcessId(1), + code: None, + reason: super::ProcessExitReason::Killed, + }; + let json = serde_json::to_string(&ev).unwrap(); + assert!(!json.contains("code")); + let back: Event = serde_json::from_str(&json).unwrap(); + assert_eq!(back, ev); + } + + #[test] + fn transcript_entry_process_roundtrips() { + let entry = TranscriptEntry::Process { + command: "pnpm dev".to_owned(), + output: "ready".to_owned(), + }; + let json = serde_json::to_string(&entry).unwrap(); + let back: TranscriptEntry = serde_json::from_str(&json).unwrap(); + assert_eq!(back, entry); + } } diff --git a/crates/goat-protocol/src/op.rs b/crates/goat-protocol/src/op.rs index a74ed84..c4931b9 100644 --- a/crates/goat-protocol/src/op.rs +++ b/crates/goat-protocol/src/op.rs @@ -1,6 +1,6 @@ use serde::{Deserialize, Serialize}; -use crate::{InputAttachment, LoginCredential, ModelTarget, TaskId, ToolCallId}; +use crate::{InputAttachment, LoginCredential, ModelTarget, ProcessId, TaskId, ToolCallId}; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] #[serde(tag = "type")] @@ -57,5 +57,12 @@ pub enum Op { DequeueMessage { id: TaskId, }, + ProcessKill { + process: ProcessId, + }, + ProcessWatch { + process: ProcessId, + on: bool, + }, Shutdown {}, } diff --git a/crates/goat-protocol/src/types.rs b/crates/goat-protocol/src/types.rs index 9266309..122a770 100644 --- a/crates/goat-protocol/src/types.rs +++ b/crates/goat-protocol/src/types.rs @@ -125,6 +125,65 @@ impl JsonSchema for ToolCallId { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct ProcessId(pub u64); + +impl Serialize for ProcessId { + fn serialize(&self, s: S) -> Result { + id_serde::serialize(&self.0, s) + } +} + +impl<'de> Deserialize<'de> for ProcessId { + fn deserialize>(d: D) -> Result { + id_serde::deserialize(d).map(Self) + } +} + +impl fmt::Display for ProcessId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl JsonSchema for ProcessId { + fn schema_name() -> std::borrow::Cow<'static, str> { + "ProcessId".into() + } + fn schema_id() -> std::borrow::Cow<'static, str> { + concat!(module_path!(), "::ProcessId").into() + } + fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + id_json_schema(generator) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum ProcessState { + Running, + Exited, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, schemars::JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum ProcessExitReason { + Natural, + Killed, + Timeout, + Shutdown, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] +pub struct ProcessInfo { + pub id: ProcessId, + pub command: String, + pub state: ProcessState, + pub watched: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub exit_code: Option, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] pub struct ToolDisplay { pub primary: String, @@ -283,6 +342,10 @@ pub enum TranscriptEntry { command: String, output: String, }, + Process { + command: String, + output: String, + }, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)] diff --git a/crates/goat-store/src/lib.rs b/crates/goat-store/src/lib.rs index bdff1f9..0e00bc2 100644 --- a/crates/goat-store/src/lib.rs +++ b/crates/goat-store/src/lib.rs @@ -12,8 +12,8 @@ use models::thread_from_row; use schema::migrate; pub use models::{ - Compaction, NewCompaction, NewMessage, NewThread, NewToolCall, NewTurn, OpenPrompt, StoreError, - StoredMessage, Thread, + Compaction, NewCompaction, NewMessage, NewProcess, NewThread, NewToolCall, NewTurn, OpenPrompt, + OrphanProcess, StoreError, StoredMessage, Thread, }; const READER_POOL_MAX: usize = 4; @@ -387,6 +387,60 @@ impl Store { .await } + pub async fn create_process(&self, process: NewProcess) -> Result { + self.run(move |conn| { + conn.execute( + "INSERT INTO processes (pgid, command, cwd, status, started_at) + VALUES (?1, ?2, ?3, 'running', ?4)", + params![ + process.pgid, + process.command, + process.cwd, + process.started_at, + ], + )?; + Ok(conn.last_insert_rowid()) + }) + .await + } + + pub async fn finish_process(&self, id: i64, finished_at: i64) -> Result<(), StoreError> { + self.run(move |conn| { + conn.execute( + "UPDATE processes SET status = 'dead', finished_at = ?2 WHERE id = ?1", + params![id, finished_at], + )?; + Ok(()) + }) + .await + } + + pub async fn take_orphan_processes( + &self, + finished_at: i64, + ) -> Result, StoreError> { + self.run(move |conn| { + let orphans = { + let mut stmt = conn + .prepare("SELECT id, pgid, command FROM processes WHERE status = 'running'")?; + let rows = stmt.query_map([], |row| { + Ok(OrphanProcess { + id: row.get(0)?, + pgid: row.get(1)?, + command: row.get(2)?, + }) + })?; + rows.collect::>>()? + }; + conn.execute( + "UPDATE processes SET status = 'dead', finished_at = ?1 WHERE status = 'running'", + params![finished_at], + )?; + Ok(orphans) + }) + .await + } + pub async fn create_message(&self, message: NewMessage) -> Result { self.run(move |conn| { conn.execute( @@ -502,7 +556,7 @@ impl Store { #[cfg(test)] #[cfg(test)] mod tests { - use super::{NewMessage, NewThread, NewToolCall, NewTurn, Store}; + use super::{NewMessage, NewProcess, NewThread, NewToolCall, NewTurn, Store}; fn sample_thread() -> NewThread { NewThread { @@ -841,7 +895,7 @@ mod tests { }) .await .unwrap(); - assert_eq!(has_index, (1, 9)); + assert_eq!(has_index, (1, 10)); let _ = std::fs::remove_file(&path); } @@ -873,7 +927,7 @@ mod tests { }) .await .unwrap(); - assert_eq!(result, (0, 9)); + assert_eq!(result, (0, 10)); let _ = std::fs::remove_file(&path); } @@ -910,7 +964,7 @@ mod tests { }) .await .unwrap(); - assert_eq!(result, (0, 9)); + assert_eq!(result, (0, 10)); let _ = std::fs::remove_file(&path); } @@ -967,4 +1021,37 @@ mod tests { let _ = std::fs::remove_file(path.with_extension("db-wal")); let _ = std::fs::remove_file(path.with_extension("db-shm")); } + + #[tokio::test] + async fn processes_roundtrip_and_sweep_orphans() { + let store = Store::open_in_memory().unwrap(); + let alive = store + .create_process(NewProcess { + pgid: 4242, + command: "pnpm dev".into(), + cwd: "/tmp/project".into(), + started_at: 100, + }) + .await + .unwrap(); + let finished = store + .create_process(NewProcess { + pgid: 4343, + command: "gh run watch".into(), + cwd: "/tmp/project".into(), + started_at: 101, + }) + .await + .unwrap(); + store.finish_process(finished, 150).await.unwrap(); + + let orphans = store.take_orphan_processes(200).await.unwrap(); + assert_eq!(orphans.len(), 1); + assert_eq!(orphans[0].id, alive); + assert_eq!(orphans[0].pgid, 4242); + assert_eq!(orphans[0].command, "pnpm dev"); + + let again = store.take_orphan_processes(300).await.unwrap(); + assert!(again.is_empty()); + } } diff --git a/crates/goat-store/src/models.rs b/crates/goat-store/src/models.rs index 66c0bad..f604b1d 100644 --- a/crates/goat-store/src/models.rs +++ b/crates/goat-store/src/models.rs @@ -107,6 +107,21 @@ pub struct NewToolCall { pub started_at: i64, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NewProcess { + pub pgid: i64, + pub command: String, + pub cwd: String, + pub started_at: i64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OrphanProcess { + pub id: i64, + pub pgid: i64, + pub command: String, +} + pub(crate) fn thread_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { Ok(Thread { id: row.get(0)?, diff --git a/crates/goat-store/src/schema.rs b/crates/goat-store/src/schema.rs index b2545fe..ff805f2 100644 --- a/crates/goat-store/src/schema.rs +++ b/crates/goat-store/src/schema.rs @@ -2,7 +2,7 @@ use rusqlite::Connection; use crate::StoreError; -pub(crate) const LATEST_VERSION: i64 = 9; +pub(crate) const LATEST_VERSION: i64 = 10; pub(crate) const SCHEMA_V1: &str = "\ CREATE TABLE threads ( @@ -96,6 +96,18 @@ pub(crate) const SCHEMA_V8: &str = "ALTER TABLE threads DROP COLUMN mode;"; pub(crate) const SCHEMA_V9: &str = "DROP TABLE IF EXISTS session_events;"; +pub(crate) const SCHEMA_V10: &str = "\ +CREATE TABLE processes ( + id INTEGER PRIMARY KEY, + pgid INTEGER NOT NULL, + command TEXT NOT NULL, + cwd TEXT NOT NULL, + status TEXT NOT NULL, + started_at INTEGER NOT NULL, + finished_at INTEGER +); +CREATE INDEX idx_processes_status ON processes(status);"; + pub(crate) fn migrate(conn: &mut Connection) -> Result<(), StoreError> { let mut version: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?; if version > LATEST_VERSION { @@ -112,6 +124,7 @@ pub(crate) fn migrate(conn: &mut Connection) -> Result<(), StoreError> { 6 => SCHEMA_V7, 7 => SCHEMA_V8, 8 => SCHEMA_V9, + 9 => SCHEMA_V10, _ => return Err(StoreError::UnknownVersion(version)), }; let next = version + 1; diff --git a/crates/goat-tui/src/app/engine.rs b/crates/goat-tui/src/app/engine.rs index 49864d5..635e664 100644 --- a/crates/goat-tui/src/app/engine.rs +++ b/crates/goat-tui/src/app/engine.rs @@ -74,7 +74,8 @@ impl App { } => { self.transcript.push_compaction(tokens_before, tokens_after); } - TranscriptEntry::Shell { command, output } => { + TranscriptEntry::Shell { command, output } + | TranscriptEntry::Process { command, output } => { let id = TaskId(0); self.transcript.push_shell(id, command); self.transcript.finish_shell(id, output); @@ -104,6 +105,16 @@ impl App { } } EngineEvent::LoginProviders { .. } | EngineEvent::ThreadBound { .. } => {} + EngineEvent::ProcessListChanged { processes } => { + self.processes = processes; + self.dirty = true; + } + EngineEvent::ProcessStarted { .. } + | EngineEvent::ProcessOutput { .. } + | EngineEvent::ProcessExited { .. } + | EngineEvent::ProcessObserved { .. } => { + self.dirty = true; + } EngineEvent::CompactionStarted { id } => { if self.agent_index(id).is_none() { self.turn.compacting = true; diff --git a/crates/goat-tui/src/app/mod.rs b/crates/goat-tui/src/app/mod.rs index 1ae8477..a2bb4db 100644 --- a/crates/goat-tui/src/app/mod.rs +++ b/crates/goat-tui/src/app/mod.rs @@ -135,6 +135,7 @@ pub struct App { pub(crate) focused: bool, pub(crate) notification_pending: Option, pub(crate) picker: Option, + pub(crate) processes: Vec, } #[derive(Default)] @@ -224,6 +225,7 @@ impl App { focused: true, notification_pending: None, picker: None, + processes: Vec::new(), } } @@ -1305,6 +1307,35 @@ impl App { Some(format!("{running} agents · {}", parts.join(", "))) } + pub(crate) fn process_summary(&self) -> Option { + let running: Vec<&goat_protocol::ProcessInfo> = self + .processes + .iter() + .filter(|p| p.state == goat_protocol::ProcessState::Running) + .collect(); + if running.is_empty() { + return None; + } + let mut shown: Vec = running + .iter() + .take(3) + .map(|p| { + let watch = if p.watched { "*" } else { "" }; + let cmd: String = p + .command + .split_whitespace() + .take(2) + .collect::>() + .join(" "); + format!("#{}{watch} {cmd}", p.id) + }) + .collect(); + if running.len() > 3 { + shown.push(format!("+{}", running.len() - 3)); + } + Some(shown.join(", ")) + } + pub(crate) fn current_context_window(&self) -> Option { let model = self.model.as_ref()?; self.context_window @@ -2723,4 +2754,34 @@ mod tests { assert!(ops.is_empty()); assert!(!app.take_dirty()); } + + #[test] + fn process_list_updates_summary_and_ignores_exited() { + let mut app = App::new(Theme::dark()); + assert!(app.process_summary().is_none()); + app.on_engine(EngineEvent::ProcessListChanged { + processes: vec![ + goat_protocol::ProcessInfo { + id: goat_protocol::ProcessId(1), + command: "pnpm dev".to_owned(), + state: goat_protocol::ProcessState::Running, + watched: false, + exit_code: None, + }, + goat_protocol::ProcessInfo { + id: goat_protocol::ProcessId(2), + command: "gh run watch".to_owned(), + state: goat_protocol::ProcessState::Exited, + watched: true, + exit_code: Some(0), + }, + ], + }); + let summary = app.process_summary().expect("running process shown"); + assert!(summary.contains("#1"), "got: {summary}"); + assert!( + !summary.contains("#2"), + "exited process must not show: {summary}" + ); + } } diff --git a/crates/goat-tui/src/view.rs b/crates/goat-tui/src/view.rs index c08b584..0f23a5a 100644 --- a/crates/goat-tui/src/view.rs +++ b/crates/goat-tui/src/view.rs @@ -329,7 +329,7 @@ fn render_toasts(frame: &mut Frame, area: Rect, app: &App, theme: Theme) { } fn footer_visible(app: &App) -> bool { - app.quit_armed() || app.is_busy() || app.clear_armed() + app.quit_armed() || app.is_busy() || app.clear_armed() || app.process_summary().is_some() } fn render_selection(frame: &mut Frame, app: &mut App, theme: Theme) { @@ -799,6 +799,14 @@ fn render_footer(frame: &mut Frame, area: Rect, app: &App, theme: Theme) { ])), inner, ); + } else if let Some(summary) = app.process_summary() { + frame.render_widget( + Paragraph::new(Line::from(vec![ + Span::styled(symbols::ui::BULLET, theme.hint_key()), + Span::styled(format!(" {summary}"), theme.muted()), + ])), + inner, + ); } }