diff --git a/crates/maestro-daemon/src/attach.rs b/crates/maestro-daemon/src/attach.rs index 84a83dc..c53fb89 100644 --- a/crates/maestro-daemon/src/attach.rs +++ b/crates/maestro-daemon/src/attach.rs @@ -6,29 +6,26 @@ use axum::extract::ws::{Message, WebSocket}; use futures_util::{SinkExt, StreamExt}; use maestro_protocol::ClientControl; +use tokio::sync::broadcast; use tokio::sync::broadcast::error::RecvError; use uuid::Uuid; use crate::session::SessionManager; /// Pump a WebSocket against a session's PTY, behind `WS /api/attach/:id`. +/// +/// The subscription is acquired by the handler *before* the upgrade and handed +/// in here, so there is no window in which the session can disappear between +/// the existence check and the subscribe. pub(crate) async fn serve( socket: WebSocket, manager: SessionManager, id: Uuid, + replay: Vec, + mut rx: broadcast::Receiver>, ) -> Result<(), axum::Error> { let (mut sink, mut stream) = socket.split(); - let (replay, mut rx) = match manager.subscribe_pty_with_replay(id).await { - Some(pair) => pair, - None => { - // session was killed between the existence check and now — let the - // client see a clean close instead of a dropped socket. - let _ = sink.send(Message::Close(None)).await; - return Ok(()); - } - }; - // NAP-229: send buffered scrollback up-front so the client sees the last // screenful immediately instead of a blank terminal. The replay lock the // bridge held during snapshot+subscribe guarantees no byte appears both diff --git a/crates/maestro-daemon/src/local_api.rs b/crates/maestro-daemon/src/local_api.rs index c456364..d7bb438 100644 --- a/crates/maestro-daemon/src/local_api.rs +++ b/crates/maestro-daemon/src/local_api.rs @@ -177,21 +177,25 @@ pub fn state_for( // --------------------------------------------------------------------------- #[allow(clippy::result_large_err)] +/// Subscribing *before* the upgrade makes the 404 and the subscription one +/// atomic step. Checking existence and then subscribing inside the callback +/// left a window in which the liveness watcher could reap the session, so a +/// client that passed the check still got a socket that immediately closed. async fn attach_ws( Path(id): Path, State(s): State, ws: axum::extract::WebSocketUpgrade, ) -> Result { - if !s.manager.has(id).await { + let Some((replay, rx)) = s.manager.subscribe_pty_with_replay(id).await else { return Err(error_response( StatusCode::NOT_FOUND, "unknown_session", "no such session on this daemon", )); - } + }; let manager = s.manager.clone(); Ok(ws.on_upgrade(move |socket| async move { - let _ = crate::attach::serve(socket, manager, id).await; + let _ = crate::attach::serve(socket, manager, id, replay, rx).await; })) } diff --git a/crates/maestro-daemon/src/proc.rs b/crates/maestro-daemon/src/proc.rs index afe991b..28cba42 100644 --- a/crates/maestro-daemon/src/proc.rs +++ b/crates/maestro-daemon/src/proc.rs @@ -1,8 +1,8 @@ //! Best-effort `/proc` inspection for the activity watcher (NAP-340, Linux). //! -//! Decides whether a session is **blocked waiting on the user**: walk from the -//! pane's process down to the deepest descendant (the agent/shell actually on -//! the terminal) and check it's parked waiting on terminal input. +//! Decides whether a session is **blocked waiting on the user**: ask the kernel +//! which process group owns the terminal, then check whether any live member of +//! that group is parked on the terminal. //! //! `/proc//syscall` would say so exactly, but it's ptrace-gated and //! unreadable for our own sessions on a default kernel, so we combine three @@ -16,55 +16,120 @@ //! Every step degrades to `false` on a read failure, so a non-Linux host or a //! denied `/proc` read falls back to `Idle`. -/// Whether the process tree under `pane_pid` looks blocked on terminal input -/// from `tty` (e.g. `/dev/pts/3`) — the signal for `AwaitingInput`. -pub fn awaiting_input(pane_pid: i32, tty: &str) -> bool { - let leaf = deepest_descendant(pane_pid); +use std::collections::HashMap; + +/// One walk of `/proc`, shared by every session in an activity pass. +/// +/// The foreground process group is the unit that matters, and its members are +/// not reliably reachable from the pane: the group leader can exit while others +/// run on, and a stage that forks a terminal-reading child and exits leaves +/// that child reparented to init. Walking the pane's process tree misses both. +/// +/// Snapshotting once per pass costs one `/proc` walk however many sessions the +/// daemon owns — cheaper than the per-session `tmux list-panes` fork that +/// already runs on the same tick. +#[derive(Debug, Default)] +pub struct ProcGroups { + stats: HashMap, + by_pgid: HashMap>, +} + +impl ProcGroups { + pub fn snapshot() -> Self { + let mut stats = HashMap::new(); + let mut by_pgid: HashMap> = HashMap::new(); + let Ok(entries) = std::fs::read_dir("/proc") else { + return Self::default(); + }; + for pid in entries + .flatten() + .filter_map(|e| e.file_name().to_str()?.parse::().ok()) + { + // A process can exit mid-walk; skipping it is the correct answer. + if let Some(stat) = read_stat(pid) { + by_pgid.entry(stat.pgrp).or_default().push(pid); + stats.insert(pid, stat); + } + } + Self { stats, by_pgid } + } + + /// The process group the terminal considers foreground, per `tpgid`. + fn foreground_pgid(&self, pane_pid: i32) -> Option { + let tpgid = self.stats.get(&pane_pid)?.tpgid; + // -1 when the process has no controlling terminal. + (tpgid > 0).then_some(tpgid) + } + + fn members(&self, pgid: i32) -> &[i32] { + self.by_pgid.get(&pgid).map_or(&[], Vec::as_slice) + } + + fn state(&self, pid: i32) -> Option { + self.stats.get(&pid).map(|s| s.state) + } +} + +/// Whether the session on `tty` (e.g. `/dev/pts/3`) is blocked waiting on the +/// user — the signal for `AwaitingInput`. +/// +/// True when *any* member of the terminal's foreground group is parked on it: a +/// group can hold a busy or timer-sleeping stage alongside the one actually +/// reading, and either can be the reparented one. +pub fn awaiting_input(pane_pid: i32, tty: &str, groups: &ProcGroups) -> bool { + let Some(pgid) = groups.foreground_pgid(pane_pid) else { + return false; + }; + groups + .members(pgid) + .iter() + .any(|pid| parked_on_tty(*pid, tty, groups)) +} + +/// Whether `pid` is parked on `tty` rather than busy or sleeping on a timer. +fn parked_on_tty(pid: i32, tty: &str, groups: &ProcGroups) -> bool { // Interruptible sleep is where a blocking read parks; anything else // (running, uninterruptible) is busy, not waiting on us. - if read_state(leaf) != Some('S') { + if groups.state(pid) != Some('S') { return false; } // A pure timer sleep (`sleep`, `usleep`) also sits in 'S' with the terminal // inherited — exclude it so it reads as Idle, not AwaitingInput. - if is_timer_sleep(leaf) { + if is_timer_sleep(pid) { return false; } - fd_points_at_tty(leaf, tty) + fd_points_at_tty(pid, tty) } -/// Follow the (last) child link down the process tree to the deepest -/// descendant. Bounded so a pathological/cyclic tree can't loop. -fn deepest_descendant(pid: i32) -> i32 { - let mut cur = pid; - for _ in 0..64 { - match last_child(cur) { - Some(child) if child != cur => cur = child, - _ => break, - } - } - cur +/// The fields of `/proc//stat` this module needs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Stat { + state: char, + /// The process group this process belongs to. + pgrp: i32, + /// Foreground process group of the controlling terminal. + tpgid: i32, } -/// The most-recently-created child of `pid`, via -/// `/proc//task//children` (a space-separated pid list). The last -/// entry is the newest child — typically the foreground one an agent spawned. -fn last_child(pid: i32) -> Option { - let contents = std::fs::read_to_string(format!("/proc/{pid}/task/{pid}/children")).ok()?; - contents.split_whitespace().last()?.parse().ok() +fn read_stat(pid: i32) -> Option { + parse_stat(&std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?) } -/// Read the process state char from `/proc//stat`. -fn read_state(pid: i32) -> Option { - parse_stat_state(&std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?) -} - -/// Extract the state field from a `/proc//stat` line. The format is -/// `pid (comm) STATE ...`; `comm` can itself contain spaces and parens, so the -/// state is the first non-space char *after the final `)`*. -fn parse_stat_state(stat: &str) -> Option { +/// Parse the fields after `comm` from a `/proc//stat` line. +/// +/// The format is `pid (comm) state ppid pgrp session tty_nr tpgid ...`, and +/// `comm` can itself contain spaces and parens — so everything is indexed from +/// the *final* `)`, never by splitting the whole line. +fn parse_stat(stat: &str) -> Option { let close = stat.rfind(')')?; - stat[close + 1..].trim_start().chars().next() + let mut rest = stat[close + 1..].split_whitespace(); + let state = rest.next()?.chars().next()?; + let _ppid = rest.next()?; + let pgrp = rest.next()?.parse().ok()?; + let _session = rest.next()?; + let _tty_nr = rest.next()?; + let tpgid = rest.next()?.parse().ok()?; + Some(Stat { state, pgrp, tpgid }) } /// Whether `pid`'s `wchan` (the kernel symbol it's blocked in) is a timer @@ -94,19 +159,51 @@ fn fd_points_at_tty(pid: i32, tty: &str) -> bool { mod tests { use super::*; - #[test] - fn parse_stat_state_handles_simple_and_tricky_comm() { - assert_eq!(parse_stat_state("1234 (bash) S 1 1234 ..."), Some('S')); - assert_eq!(parse_stat_state("1234 (sleep) S 1 ..."), Some('S')); - assert_eq!(parse_stat_state("99 (cat) R 1 ..."), Some('R')); - // comm with spaces and a paren — only the final ')' delimits it. - assert_eq!(parse_stat_state("42 (weird (name)) D 1 42 ..."), Some('D')); + /// `pid (comm) state ppid pgrp session tty_nr tpgid …` + fn stat_line(comm: &str, state: char, tpgid: i32) -> String { + stat_line_pgrp(comm, state, 1234, tpgid) + } + + fn stat_line_pgrp(comm: &str, state: char, pgrp: i32, tpgid: i32) -> String { + format!("1234 ({comm}) {state} 1000 {pgrp} 1234 34816 {tpgid} 4194304 …") } #[test] - fn parse_stat_state_rejects_garbage() { - assert_eq!(parse_stat_state("not a stat line"), None); - assert_eq!(parse_stat_state("123 (comm)"), None); + fn parse_stat_reads_state_pgrp_and_tpgid() { + let s = parse_stat(&stat_line_pgrp("bash", 'S', 500, 4321)).unwrap(); + assert_eq!(s.state, 'S'); + assert_eq!(s.pgrp, 500); + assert_eq!(s.tpgid, 4321); + + let r = parse_stat(&stat_line_pgrp("cat", 'R', 7, 99)).unwrap(); + assert_eq!(r.state, 'R'); + assert_eq!(r.pgrp, 7); + assert_eq!(r.tpgid, 99); + } + + /// `comm` is attacker-adjacent: it is the process name, it can contain + /// spaces and parens, and splitting the whole line on whitespace would put + /// every later field at the wrong index. + #[test] + fn parse_stat_indexes_from_the_final_paren() { + let s = parse_stat(&stat_line_pgrp("weird (name) with spaces", 'D', 42, 777)).unwrap(); + assert_eq!(s.state, 'D'); + assert_eq!(s.pgrp, 42); + assert_eq!(s.tpgid, 777); + } + + #[test] + fn parse_stat_carries_the_no_terminal_sentinel() { + // -1 means no controlling terminal; `foreground_pid` rejects it. + assert_eq!(parse_stat(&stat_line("daemon", 'S', -1)).unwrap().tpgid, -1); + } + + #[test] + fn parse_stat_rejects_garbage() { + assert!(parse_stat("not a stat line").is_none()); + assert!(parse_stat("123 (comm)").is_none()); + // Truncated before tpgid. + assert!(parse_stat("123 (comm) S 1 1 1").is_none()); } #[test] @@ -121,9 +218,46 @@ mod tests { assert!(!is_timer_wchan("0")); } + /// The snapshot indexes by process group, not by parentage — that is the + /// whole point, since a group member need not be under the pane. Checked + /// against this process's own group, which certainly has a live member. + #[test] + fn snapshot_indexes_this_process_under_its_group() { + // SAFETY: argument-less query of the caller's own process group. + let own = unsafe { libc::getpgid(0) }; + assert!(own > 0, "expected a real process group, got {own}"); + + let groups = ProcGroups::snapshot(); + let me = std::process::id() as i32; + assert!( + groups.members(own).contains(&me), + "snapshot missed this process ({me}) in group {own}: {:?}", + groups.members(own) + ); + assert!( + groups.state(me).is_some(), + "a process in the snapshot must carry its state" + ); + } + + #[test] + fn snapshot_has_no_members_for_a_group_that_cannot_exist() { + assert!(ProcGroups::snapshot().members(i32::MAX).is_empty()); + } + + #[test] + fn a_pid_with_no_controlling_terminal_has_no_foreground_group() { + let groups = ProcGroups::snapshot(); + assert_eq!(groups.foreground_pgid(i32::MAX), None); + } + #[test] fn awaiting_input_is_false_for_nonexistent_pid() { - // A pid that can't exist → every read fails → false (fall back to Idle). - assert!(!awaiting_input(i32::MAX, "/dev/pts/0")); + // A pid that can't exist → not in the snapshot → false (falls back to Idle). + assert!(!awaiting_input( + i32::MAX, + "/dev/pts/0", + &ProcGroups::snapshot() + )); } } diff --git a/crates/maestro-daemon/src/session.rs b/crates/maestro-daemon/src/session.rs index 1019dcb..2372aef 100644 --- a/crates/maestro-daemon/src/session.rs +++ b/crates/maestro-daemon/src/session.rs @@ -595,6 +595,11 @@ async fn run_activity_pass(inner: &Arc) { return; } + // One `/proc` walk for the whole pass, rather than one per session. Reading + // it is pure filesystem work, so it goes on the blocking pool alongside the + // tmux calls below. + let groups = tmux_blocking(crate::proc::ProcGroups::snapshot).await; + for id in ids { // The tmux pane lookup is a subprocess — keep it off the runtime worker // and outside the sessions lock. @@ -611,7 +616,7 @@ async fn run_activity_pass(inner: &Arc) { } }; // `/proc` inspection is cheap, but still done off the async lock. - let awaiting = crate::proc::awaiting_input(pane.pane_pid, &pane.tty); + let awaiting = crate::proc::awaiting_input(pane.pane_pid, &pane.tty, &groups); let updated = { let mut sessions = inner.sessions.lock().await; diff --git a/crates/maestro-daemon/tests/activity.rs b/crates/maestro-daemon/tests/activity.rs new file mode 100644 index 0000000..d7fe201 --- /dev/null +++ b/crates/maestro-daemon/tests/activity.rs @@ -0,0 +1,196 @@ +//! Activity detection against real processes on a real tty. +//! +//! The review noted this pipeline had no end-to-end coverage — `activity::decide` +//! was unit-tested, but nothing exercised the `/proc` inspection that feeds it. +//! These use a real tmux pane because that is the only way to get a controlling +//! terminal with a genuine foreground process group. + +mod common; + +use std::time::Duration; + +use common::TmuxFixture; +use maestro_daemon::proc::{awaiting_input, ProcGroups}; +use tempfile::TempDir; +use uuid::Uuid; + +/// Poll `awaiting_input` until it reports `want`, so we don't race the shell +/// reaching its blocking read. Reports what it last saw — the useful half of +/// the failure. +fn wait_for_awaiting(fx: &TmuxFixture, id: Uuid, want: bool, label: &str) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + let mut last = None; + while std::time::Instant::now() < deadline { + if let Ok(Some(pane)) = fx.tmux.pane_info(id) { + // Fresh each poll: the tree is still settling. + let got = awaiting_input(pane.pane_pid, &pane.tty, &ProcGroups::snapshot()); + if got == want { + return; + } + last = Some(got); + } + std::thread::sleep(Duration::from_millis(25)); + } + panic!("{label}: wanted awaiting_input={want}, last saw {last:?}"); +} + +#[test] +fn an_idle_shell_is_awaiting_input() { + let fx = TmuxFixture::new(); + let cwd = TempDir::new().unwrap(); + let id = Uuid::new_v4(); + fx.tmux + .new_session(id, cwd.path(), "bash --noprofile --norc") + .expect("new_session"); + + wait_for_awaiting(&fx, id, true, "an idle shell sits in read()"); +} + +/// The case the old heuristic got wrong. It walked `children` to the deepest +/// descendant, so it landed on the background `sleep` — which is not on the +/// terminal — and reported the session as *not* awaiting input, while the +/// foreground process was blocked on `read`. +#[test] +fn a_foreground_reader_with_a_background_helper_is_awaiting_input() { + let fx = TmuxFixture::new(); + let cwd = TempDir::new().unwrap(); + let script = cwd.path().join("agentish.sh"); + std::fs::write(&script, "#!/bin/bash\nsleep 300 &\nread -r line\n").unwrap(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).unwrap(); + } + + let id = Uuid::new_v4(); + fx.tmux + .new_session( + id, + cwd.path(), + &format!("bash --noprofile --norc {}", script.display()), + ) + .expect("new_session"); + + // Wait for the helper itself. Asserting before it exists would pass + // against a deepest-descendant walk too, because until then the newest + // process in the tree *is* the one blocked on read. + let pane = fx.tmux.pane_info(id).expect("pane_info").expect("pane"); + wait_until( + || subtree_has_comm(pane.pane_pid, "sleep"), + "the background helper to be spawned", + ); + + wait_for_awaiting( + &fx, + id, + true, + "a foreground reader must win over a newer background child", + ); +} + +/// A pure timer sleep also parks in 'S' with the terminal inherited; it must +/// read as idle, not as waiting on the user. +#[test] +fn a_foreground_sleep_is_not_awaiting_input() { + let fx = TmuxFixture::new(); + let cwd = TempDir::new().unwrap(); + let id = Uuid::new_v4(); + fx.tmux + .new_session(id, cwd.path(), "sleep 300") + .expect("new_session"); + + wait_for_awaiting(&fx, id, false, "a timer sleep is not user input"); +} + +/// Whether any process under `root` (inclusive) has this `comm`. +fn subtree_has_comm(root: i32, comm: &str) -> bool { + let mut queue = vec![root]; + while let Some(pid) = queue.pop() { + if std::fs::read_to_string(format!("/proc/{pid}/comm")).is_ok_and(|c| c.trim() == comm) { + return true; + } + queue.extend( + std::fs::read_to_string(format!("/proc/{pid}/task/{pid}/children")) + .map(|s| { + s.split_whitespace() + .filter_map(|x| x.parse::().ok()) + .collect::>() + }) + .unwrap_or_default(), + ); + } + false +} + +/// Block until `cond` holds, so a test asserts on the state it means to +/// exercise rather than on whatever the tree looked like first. +fn wait_until(cond: impl Fn() -> bool, what: &str) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while std::time::Instant::now() < deadline { + if cond() { + return; + } + std::thread::sleep(Duration::from_millis(20)); + } + panic!("timed out waiting for {what}"); +} + +/// Read `pgrp` and `tpgid` for a pid straight from `/proc`, so the test can +/// wait for the exact condition it means to exercise rather than a proxy. +fn pgrp_and_tpgid(pid: i32) -> Option<(i32, i32)> { + let stat = std::fs::read_to_string(format!("/proc/{pid}/stat")).ok()?; + let close = stat.rfind(')')?; + let f: Vec<&str> = stat[close + 1..].split_whitespace().collect(); + Some((f.get(2)?.parse().ok()?, f.get(5)?.parse().ok()?)) +} + +/// `tpgid` names a process *group*, and its leader can exit while another +/// member still holds the terminal — a foreground pipeline whose first stage +/// finished. Reading `/proc/` directly then finds nothing and reports +/// idle, so the group has to be searched for a live member. +/// +/// Needs an interactive shell: job control is what puts a pipeline in its own +/// process group, and that is how the daemon spawns commands +/// (`$SHELL -lic ''`). +#[test] +fn a_foreground_pipeline_whose_leader_exited_is_still_awaiting_input() { + let fx = TmuxFixture::new(); + let cwd = TempDir::new().unwrap(); + let id = Uuid::new_v4(); + fx.tmux + .new_session(id, cwd.path(), "bash --noprofile --norc -i") + .expect("new_session"); + wait_for_awaiting(&fx, id, true, "the shell must settle first"); + + // `true` leads the group and exits immediately; the right-hand stage + // survives, blocked reading the terminal. + fx.tmux + .send_keys(id, b"true | { read -r line < /dev/tty; }\n") + .expect("send_keys"); + + // Wait for the state this test exists for: the foreground group is no + // longer the shell's own, and its leader is gone. Asserting before that + // would pass against a leader-only lookup, because the shell is still + // foreground and alive. + let pane = fx.tmux.pane_info(id).expect("pane_info").expect("pane"); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + let mut reached = false; + while std::time::Instant::now() < deadline { + if let Some((pgrp, tpgid)) = pgrp_and_tpgid(pane.pane_pid) { + if tpgid > 0 && tpgid != pgrp && pgrp_and_tpgid(tpgid).is_none() { + reached = true; + break; + } + } + std::thread::sleep(Duration::from_millis(20)); + } + assert!( + reached, + "never observed a foreground group whose leader had exited" + ); + + assert!( + awaiting_input(pane.pane_pid, &pane.tty, &ProcGroups::snapshot()), + "a surviving pipeline member still holds the tty" + ); +}