use std::collections::HashMap;
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use serde::Deserialize;
use super::{ClientId, Event, Server};
use crate::agent::AgentState;
use crate::layout::PaneId;
use crate::pane::PaneSpec;
use crate::protocol::{PaneInfo, Placement};
use crate::servers;
const MAX_LINE: usize = 4 << 20;
const MAX_AGENTS: usize = 64;
const MAX_NAME: usize = 48;
const LIST_GAP: Duration = Duration::from_millis(100);
const HEALTHY: Duration = Duration::from_secs(60);
#[derive(Debug, Clone)]
pub struct RemoteAgent {
pub pane: u32,
pub name: String,
pub workspace: String,
pub agent: String,
pub state: AgentState,
focused: bool,
pub since: Instant,
}
impl RemoteAgent {
fn same(&self, other: &RemoteAgent) -> bool {
(self.pane, &self.name, &self.workspace, &self.agent, self.state, self.focused)
== (other.pane, &other.name, &other.workspace, &other.agent, other.state, other.focused)
}
}
#[derive(Deserialize)]
struct Listed {
id: u32,
#[serde(default)]
name: Option<String>,
#[serde(default)]
title: String,
#[serde(default)]
workspace: String,
#[serde(default)]
agent: Option<String>,
#[serde(default)]
state: String,
#[serde(default)]
focused: bool,
#[serde(default)]
session: Option<String>,
}
fn clean(text: &str, max: usize) -> String {
let invisible = |c: &char| matches!(c, '\u{200b}'..='\u{200f}' | '\u{202a}'..='\u{202e}' | '\u{2060}'..='\u{206f}' | '\u{feff}');
let text: String = text.chars().filter(|c| !invisible(c)).take(max * 4).map(crate::text::visible_char).collect();
crate::text::truncate(text.trim(), max)
}
fn parse_list(line: &str, now: Instant) -> Option<Vec<RemoteAgent>> {
let listed: Vec<Listed> = serde_json::from_str(line).ok()?;
let agents = listed
.into_iter()
.filter(|p| p.session.is_none())
.filter_map(|p| {
let agent = p.agent.filter(|a| !a.trim().is_empty())?;
let name = p.name.filter(|n| !n.trim().is_empty()).unwrap_or(p.title);
let name = if name.trim().is_empty() { agent.clone() } else { name };
Some(RemoteAgent {
pane: p.id,
name: clean(&name, MAX_NAME),
workspace: clean(&p.workspace, MAX_NAME),
agent: clean(&agent, MAX_NAME),
state: AgentState::parse(&p.state).unwrap_or_default(),
focused: p.focused,
since: now,
})
})
.take(MAX_AGENTS)
.collect();
Some(agents)
}
pub struct Watch {
pub server: servers::Server,
pub session: String,
pub agents: Vec<RemoteAgent>,
conn: u32,
child: Option<Child>,
stop_reading: Arc<AtomicBool>,
said: Arc<Mutex<String>>,
seen: HashMap<u32, AgentState>,
heard: bool,
started: Option<Instant>,
retry_at: Option<Instant>,
failures: u32,
warned: bool,
}
impl Watch {
fn new(server: servers::Server, session: String) -> Watch {
Watch {
server,
session,
agents: Vec::new(),
conn: 0,
child: None,
stop_reading: Default::default(),
said: Default::default(),
seen: HashMap::new(),
heard: false,
started: None,
retry_at: None,
failures: 0,
warned: false,
}
}
fn stop(&mut self) -> Option<std::process::ExitStatus> {
self.stop_reading.store(true, Ordering::SeqCst);
crate::exec::kill_group_nowait(self.child.take()?)
}
fn take(&mut self, mut list: Vec<RemoteAgent>) -> (bool, Vec<(RemoteAgent, Option<AgentState>)>) {
for a in &mut list {
if let Some(old) = self.agents.iter().find(|o| o.pane == a.pane && o.state == a.state) {
a.since = old.since;
}
}
let moved = if self.heard {
list.iter()
.map(|a| (a.clone(), self.seen.get(&a.pane).copied()))
.filter(|(a, before)| *before != Some(a.state))
.collect()
} else {
Vec::new()
};
self.seen = list.iter().map(|a| (a.pane, a.state)).collect();
self.heard = true;
let changed = self.agents.len() != list.len() || self.agents.iter().zip(&list).any(|(a, b)| !a.same(b));
self.agents = list;
(changed, moved)
}
fn lost(&mut self) -> bool {
!std::mem::take(&mut self.agents).is_empty()
}
fn back_off(&mut self) {
if self.started.take().is_some_and(|t| t.elapsed() >= HEALTHY) {
self.failures = 0;
}
self.failures += 1;
self.retry_at = Some(Instant::now() + Duration::from_secs((2u64 << self.failures.min(5)).min(60)));
}
}
fn watch_args(s: &servers::Server, known: &[servers::Known], session: &str) -> Vec<String> {
let mut args: Vec<String> = [
"-T",
"-o",
"BatchMode=yes",
"-o",
"ConnectTimeout=15",
"-o",
"ServerAliveInterval=15",
"-o",
"ServerAliveCountMax=3",
"-o",
"ControlMaster=no",
"-o",
"ControlPath=none",
]
.iter()
.map(|a| a.to_string())
.collect();
args.extend(servers::ssh_args(s, known, false));
args.push(s.destination());
args.push(crate::client::remote_kraal(&format!("-s {} agents --follow", crate::shell::quote(session))));
args
}
fn why_not(name: &str, said: &str, status: Option<std::process::ExitStatus>) -> String {
let said = clean(said, 200);
let fix = format!("`kraal remote-install {name}` puts this kraal there");
if said.contains("--follow") {
format!("the kraal on {name} is too old to list its agents: {fix}")
} else if status.and_then(|s| s.code()) == Some(127) || (said.contains("kraal") && said.contains("not found")) {
format!("no kraal on {name} to list its agents: {fix}")
} else if said.is_empty() {
let status = status.map(|s| s.to_string()).unwrap_or_else(|| "ended".into());
format!("the agents of {name} cannot be listed: ssh exited ({status})")
} else {
format!("the agents of {name} cannot be listed: {said}")
}
}
impl Server {
pub(super) fn check_remotes(&mut self) {
let wanted: Vec<(servers::Server, String)> = self
.known_servers
.iter()
.filter_map(|k| k.server.agents_session().map(|session| (k.server.clone(), session.to_string())))
.collect();
if wanted.is_empty() && self.remotes.is_empty() {
return;
}
let mut watches = std::mem::take(&mut self.remotes);
let mut changed = false;
watches.retain_mut(|w| {
let keep = wanted
.iter()
.any(|(s, session)| s.name == w.server.name && s.same_connection(&w.server) && *session == w.session);
if !keep {
w.stop();
changed |= w.lost();
}
keep
});
for (s, session) in wanted {
if !watches.iter().any(|w| w.server.name == s.name) {
watches.push(Watch::new(s, session));
}
}
let now = Instant::now();
for w in &mut watches {
if w.child.as_ref().is_some_and(crate::exec::has_exited) {
let status = w.stop();
changed |= w.lost();
if !w.heard && !std::mem::replace(&mut w.warned, true) {
let said = w.said.lock().map(|s| s.clone()).unwrap_or_default();
self.warnings.push(why_not(&w.server.name, &said, status));
}
w.back_off();
}
if w.child.is_none() && w.retry_at.is_none_or(|r| r <= now) {
w.retry_at = None;
self.start_watch(w);
}
}
self.remotes = watches;
if changed {
self.mark_dirty();
}
}
fn start_watch(&mut self, w: &mut Watch) {
w.conn = self.alloc_id();
let mut cmd = Command::new("ssh");
cmd.args(watch_args(&w.server, &self.known_servers, &w.session))
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
match crate::exec::own_group(&mut cmd).spawn() {
Ok(mut c) => {
w.stop_reading = Default::default();
w.said = Default::default();
if let Some(err) = c.stderr.take() {
let said = w.said.clone();
crate::exec::read_lines_until(err, w.stop_reading.clone(), move |line| {
if !line.trim().is_empty()
&& let Ok(mut s) = said.lock()
{
*s = line.chars().take(300).collect();
}
});
}
if let Some(out) = c.stdout.take() {
let (tx, conn) = (self.tx.clone(), w.conn);
let mut last: Option<Instant> = None;
crate::exec::read_long_lines_until(out, w.stop_reading.clone(), MAX_LINE, move |line| {
if let Some(list) = line.and_then(|l| parse_list(l, Instant::now())) {
if let Some(t) = last {
std::thread::sleep(LIST_GAP.saturating_sub(t.elapsed()));
}
last = Some(Instant::now());
let _ = tx.send(Event::RemoteAgents(conn, list));
}
});
}
w.child = Some(c);
w.started = Some(Instant::now());
}
Err(e) => {
if !std::mem::replace(&mut w.warned, true) {
self.warnings
.push(format!("the agents of {} cannot be listed: cannot run ssh: {e}", w.server.name));
}
w.back_off();
}
}
}
pub(super) fn remote_list(&mut self, conn: u32, list: Vec<RemoteAgent>) {
let Some(w) = self.remotes.iter_mut().find(|w| w.conn == conn && w.child.is_some()) else { return };
let (changed, announce) = w.take(list);
let (server, session) = (w.server.name.clone(), w.session.clone());
if changed {
self.mark_dirty();
}
for (a, before) in announce.iter().filter(|(a, _)| a.state != AgentState::Unknown) {
let fields = serde_json::json!({
"pane": a.pane,
"name": a.name,
"workspace": a.workspace,
"agent": a.agent,
"previous": before.map(|b| b.label()),
"server": server,
"session": session,
});
self.emit(&format!("agent.{}", a.state.label()), fields);
}
if self.attach_pane(&server, &session).is_some() {
return;
}
for (a, before) in announce {
let title = format!("{} Ā· {} on {server}", a.agent, a.name);
match a.state {
AgentState::Blocked if self.config.notifications.on_blocked => {
self.notify(&title, "needs your input", None)
}
AgentState::Done if self.config.notifications.on_done && before == Some(AgentState::Working) => {
self.notify(&title, "finished", None)
}
_ => {}
}
}
}
pub fn remote_agents(&self) -> impl Iterator<Item = (&str, &RemoteAgent)> {
self.remotes.iter().flat_map(|w| w.agents.iter().map(move |a| (w.server.name.as_str(), a)))
}
pub(super) fn remote_agent_infos(&self) -> Vec<PaneInfo> {
let info = |w: &Watch, a: &RemoteAgent| PaneInfo {
id: a.pane,
name: Some(a.name.clone()),
title: String::new(),
workspace: a.workspace.clone(),
tab: String::new(),
kind: "remote".into(),
command: Vec::new(),
cwd: None,
agent: Some(a.agent.clone()),
state: a.state,
exit_code: None,
suspended: false,
keyboard: 0,
server: Some(w.server.name.clone()),
session: Some(w.session.clone()),
production: None,
contexts: Vec::new(),
focused: a.focused,
floating: false,
cols: 0,
rows: 0,
};
self.remotes.iter().flat_map(|w| w.agents.iter().map(move |a| info(w, a))).collect()
}
pub(super) fn attach_pane(&self, server: &str, session: &str) -> Option<PaneId> {
let want = ["-s", session, "attach", "--host", server];
let runs =
|p: &crate::pane::Pane| p.spec.command.get(1..6).is_some_and(|a| a.iter().map(String::as_str).eq(want));
let live = || self.panes.values().filter(|p| p.exit_code.is_none() && !p.suspended);
live()
.find(|p| p.attachment().is_some_and(|a| a.server == server && a.session == session))
.or_else(|| live().find(|p| p.attachment().is_none() && runs(p)))
.map(|p| p.id)
}
pub(super) fn remote_shown(&self, client: ClientId, server: &str, pane: u32) -> bool {
let Some(w) = self.remotes.iter().find(|w| w.server.name == server) else { return false };
let Some(tab) = self.attach_pane(server, &w.session) else { return false };
if self.client_focus(client) != Some(tab) {
return false;
}
match self.panes.get(&tab).and_then(|p| p.attachment()).filter(|a| a.reports) {
Some(a) => a.shows == Some(pane),
None => w.agents.iter().any(|a| a.pane == pane && a.focused),
}
}
pub(super) fn go_to_remote(&mut self, client: ClientId, server: &str, pane: u32) {
let Some(w) = self.remotes.iter().find(|w| w.server.name == server) else {
self.toast(client, format!("{server} lists no agents any more"));
return;
};
let (target, session) = (w.server.clone(), w.session.clone());
if let Some(p) = self.attach_pane(server, &session) {
self.reveal(client, p);
let tab = self.panes.get_mut(&p);
match tab {
Some(t) if t.attachment().is_some_and(|a| a.reports) => {
t.write_input(format!("\x1b]7701;show;{pane}\x07").into_bytes());
}
_ => self.focus_remote(client, &target, &session, pane),
}
return;
}
let Some(ws) = self.clients.get(&client).map(|c| c.workspace) else { return };
let exe = std::env::current_exe().map(|p| p.display().to_string()).unwrap_or_else(|_| "kraal".into());
let command = [exe.as_str(), "-s", &session, "attach", "--host", server, "--pane", &pane.to_string()]
.map(String::from)
.to_vec();
let spec = PaneSpec { command, env_remove: vec!["KRAAL_PANE".into()], ..Default::default() };
if let Err(e) = self.spawn_pane_in(ws, None, spec, Placement::Tab, Some(server.to_string())) {
self.toast(client, format!("{e:#}"));
}
}
fn focus_remote(&mut self, client: ClientId, s: &servers::Server, session: &str, pane: u32) {
let mut args: Vec<String> =
["-T", "-o", "BatchMode=yes", "-o", "ConnectTimeout=15"].iter().map(|a| a.to_string()).collect();
args.extend(self.ssh_options(s));
args.push(s.destination());
args.push(crate::client::remote_kraal(&format!("-s {} focus {pane}", crate::shell::quote(session))));
let name = s.name.clone();
self.background(
move || {
let mut cmd = Command::new("ssh");
cmd.args(&args);
crate::exec::run(&mut cmd, Duration::from_secs(20), 64 << 10)
.map_err(|e| e.to_string())
.and_then(|out| if out.success() { Ok(()) } else { Err(crate::exec::last_stderr_line(&out)) })
},
move |server, result| {
if let Err(e) = result {
server.toast(client, format!("{name} did not show that agent: {}", clean(&e, 200)));
}
},
);
}
pub(super) fn stop_remotes(&mut self) {
for w in &mut self.remotes {
w.stop();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn lists_are_read_bounded_and_cleaned() {
let now = Instant::now();
let line = r#"[{"id":3,"name":null,"title":"ā³ Fix login","workspace":"api","agent":"claude","state":"blocked","focused":true,"cols":80},
{"id":4,"name":"evil\u001b]52;c;eA==\u0007\u202e","title":"","workspace":"w","agent":"codex","state":"sideways"},
{"id":5,"name":"shell","title":"","workspace":"w","agent":null,"state":"idle"},
{"id":6,"name":"","title":"","workspace":"w","agent":"gemini","state":"working"}]"#;
let list = parse_list(&line.replace('\n', ""), now).unwrap();
assert_eq!(list.len(), 3, "a pane without an agent is not one: {list:?}");
assert_eq!((list[0].pane, list[0].name.as_str(), list[0].state), (3, "ā³ Fix login", AgentState::Blocked));
assert!(list[0].focused);
assert!(!list[1].name.contains(['\u{1b}', '\u{7}', '\u{202e}']), "{:?}", list[1].name);
assert!(list[1].name.starts_with("evil"));
assert_eq!(list[1].state, AgentState::Unknown);
assert_eq!(list[2].name, "gemini");
let long = format!(r#"[{{"id":1,"name":"{}","agent":"claude","state":"done"}}]"#, "x".repeat(10_000));
assert_eq!(crate::text::columns(&parse_list(&long, now).unwrap()[0].name), MAX_NAME);
let many: Vec<String> =
(0..500).map(|i| format!(r#"{{"id":{i},"name":"a{i}","agent":"claude","state":"idle"}}"#)).collect();
assert_eq!(parse_list(&format!("[{}]", many.join(",")), now).unwrap().len(), MAX_AGENTS);
let relayed = r#"[{"id":7,"agent":"claude","state":"blocked","kind":"remote","server":"x","session":"main"}]"#;
assert!(parse_list(relayed, now).unwrap().is_empty());
assert!(parse_list("kraal: no kraal server", now).is_none());
assert!(parse_list(r#"{"id":1}"#, now).is_none());
assert_eq!(parse_list("[]", now).unwrap().len(), 0);
}
#[test]
fn only_states_seen_change_are_announced_even_across_a_reconnection() {
let start = Instant::now();
let agent = |pane: u32, state: AgentState| RemoteAgent {
pane,
name: format!("a{pane}"),
workspace: "w".into(),
agent: "claude".into(),
state,
focused: false,
since: Instant::now(),
};
let states = |moved: &[(RemoteAgent, Option<AgentState>)]| -> Vec<(u32, AgentState, Option<AgentState>)> {
moved.iter().map(|(a, before)| (a.pane, a.state, *before)).collect()
};
let mut w = Watch::new(servers::Server { name: "gpu".into(), ..Default::default() }, "main".into());
let (changed, moved) = w.take(vec![agent(1, AgentState::Blocked), agent(2, AgentState::Working)]);
assert!(changed && moved.is_empty());
let since = w.agents[0].since;
assert!(since >= start);
let (changed, moved) = w.take(vec![agent(1, AgentState::Blocked), agent(2, AgentState::Working)]);
assert!(!changed && moved.is_empty());
assert_eq!(w.agents[0].since, since, "still blocked since then");
let (changed, moved) = w.take(vec![agent(1, AgentState::Blocked), agent(2, AgentState::Done)]);
assert!(changed);
assert_eq!(states(&moved), [(2, AgentState::Done, Some(AgentState::Working))]);
assert!(w.lost() && w.agents.is_empty());
let (changed, moved) =
w.take(vec![agent(1, AgentState::Working), agent(2, AgentState::Done), agent(3, AgentState::Blocked)]);
assert!(changed);
assert_eq!(
states(&moved),
[(1, AgentState::Working, Some(AgentState::Blocked)), (3, AgentState::Blocked, None)]
);
}
#[test]
fn a_connection_that_keeps_dropping_is_tried_less_and_less_often() {
let mut w = Watch::new(servers::Server { name: "gpu".into(), ..Default::default() }, "main".into());
let waits = |w: &Watch| w.retry_at.unwrap().saturating_duration_since(Instant::now()).as_secs_f32();
for expected in [4.0, 8.0, 16.0, 32.0, 60.0, 60.0] {
w.take(Vec::new());
assert!(!w.lost());
w.back_off();
let wait = waits(&w);
assert!(wait > expected - 1.0 && wait <= expected, "{wait} s, not {expected}");
}
w.started = Some(Instant::now() - HEALTHY);
w.take(Vec::new());
w.back_off();
assert!(waits(&w) <= 4.0, "{}", waits(&w));
}
#[test]
fn stopping_a_watch_leaves_nothing_running() {
let mut w = Watch::new(servers::Server { name: "gpu".into(), ..Default::default() }, "main".into());
let mut cmd = Command::new("sh");
cmd.args(["-c", "sleep 30 & sleep 30"]);
let child = crate::exec::own_group(&mut cmd).spawn().unwrap();
let pid = child.id() as libc::pid_t;
w.child = Some(child);
w.stop();
assert!(w.child.is_none());
let start = Instant::now();
while unsafe { libc::kill(pid, 0) == 0 || libc::kill(-pid, 0) == 0 } {
assert!(start.elapsed() < Duration::from_secs(5), "still running");
std::thread::sleep(Duration::from_millis(10));
}
}
#[test]
fn the_watch_runs_kraal_there_without_ever_asking() {
let s = servers::Server { name: "gpu".into(), host: "gpu.lan".into(), port: Some(2222), ..Default::default() };
let args = watch_args(&s, &[], "main");
let at = |a: &str| args.iter().position(|x| x == a);
assert!(args.contains(&"BatchMode=yes".to_string()), "{args:?}");
assert!(args.contains(&"ControlPath=none".to_string()), "{args:?}");
assert!(at("gpu.lan").is_some_and(|i| i + 2 == args.len()), "the destination, then the command: {args:?}");
assert!(at("-p").is_some_and(|i| args[i + 1] == "2222"), "{args:?}");
let remote = args.last().unwrap();
assert!(remote.starts_with("sh -c "), "{remote}");
assert!(remote.contains("-s main agents --follow"), "{remote}");
}
#[test]
fn why_a_server_lists_nothing() {
let old = why_not("gpu", "error: unexpected argument '--follow' found", None);
assert!(old.contains("too old") && old.contains("kraal remote-install gpu"), "{old}");
let none = why_not("gpu", "sh: 1: exec: /home/u/.local/bin/kraal: not found", None);
assert!(none.starts_with("no kraal on gpu"), "{none}");
let ssh = why_not("gpu", "ssh: connect to host gpu.lan port 22: Connection refused\u{1b}[2J", None);
assert!(ssh.ends_with("Connection refusedā[2J"), "what the server said is shown, made visible: {ssh}");
}
#[test]
fn fuzz_remote_lists() {
crate::fuzz::run("remote_lists", 300, |rng| {
let now = Instant::now();
let text = |rng: &mut crate::fuzz::Rng| -> String {
serde_json::to_string(&String::from_utf8_lossy(&rng.text(200))).unwrap()
};
let mut items = Vec::new();
for _ in 0..rng.size(80) {
let (name, title, workspace) = (text(rng), text(rng), text(rng));
let agent = if rng.one_in(4) { "null".to_string() } else { text(rng) };
let state = *rng.pick(&["\"idle\"", "\"working\"", "\"blocked\"", "\"done\"", "\"\"", "3", "null"]);
let session = if rng.one_in(8) { ",\"session\":\"x\"" } else { "" };
items.push(format!(
"{{\"id\":{},\"name\":{name},\"title\":{title},\"workspace\":{workspace},\"agent\":{agent},\"state\":{state},\"focused\":{}{session}}}",
rng.number() % (u64::from(u32::MAX) + 2),
rng.one_in(2)
));
}
let mut line = format!("[{}]", items.join(",")).into_bytes();
if rng.one_in(3) {
rng.mutate(&mut line, &[b"[", b"]", b"{", b"}", b",", b"\"", b"\\u202e", b"\\u001b"]);
}
if let Some(list) = parse_list(&String::from_utf8_lossy(&line), now) {
assert!(list.len() <= MAX_AGENTS);
for a in &list {
for s in [&a.name, &a.workspace, &a.agent] {
assert!(crate::text::columns(s) <= MAX_NAME, "{s:?}");
assert!(!s.chars().any(|c| c.is_control() || ('\u{202a}'..='\u{202e}').contains(&c)), "{s:?}");
}
}
}
});
}
}