diff options
Diffstat (limited to 'crates/dinggerd/src/snapshot.rs')
| -rw-r--r-- | crates/dinggerd/src/snapshot.rs | 172 |
1 files changed, 160 insertions, 12 deletions
diff --git a/crates/dinggerd/src/snapshot.rs b/crates/dinggerd/src/snapshot.rs index db5ba35..07316e8 100644 --- a/crates/dinggerd/src/snapshot.rs +++ b/crates/dinggerd/src/snapshot.rs @@ -9,6 +9,11 @@ //! last [`HISTORY`] files are kept. //! - Nothing is saved while tmux is not running or has no sessions, so a //! dying server cannot overwrite the last good snapshot. +//! - Agents carry their session reference (reported by their extension), +//! other panes the command line running in them: the zsh hook's text +//! (`@dingger_cmd`), else the foreground program's arguments. When an +//! agent or program exits, the snapshot keeps it for [`GRACE_MS`], so a +//! shutdown that stops programs before tmux still restores them. //! - `restore.json`: when a different tmux server is running (or none), the //! previous server's last snapshot becomes the restore candidate, until it //! is restored or dismissed. It is pinned before the new server's own @@ -24,6 +29,9 @@ use crate::tmux; /// Snapshot files kept. const HISTORY: usize = 50; +/// How long a pane's agent or command outlives its exit in snapshots. +const GRACE_MS: u64 = 60_000; + pub fn state_dir() -> PathBuf { std::env::var_os("DINGGER_STATE_DIR") .map(PathBuf::from) @@ -80,12 +88,19 @@ pub struct WindowSnap { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PaneSnap { + /// tmux pane id on the saved server. + #[serde(default)] + pub id: String, pub index: u32, pub cwd: String, pub active: bool, /// The agent running in the pane, if any. #[serde(default, skip_serializing_if = "Option::is_none")] pub agent: Option<AgentSnap>, + /// The command line running in the pane (not for agents with a + /// session): retyped on restore, not run. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub command: Option<String>, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -97,6 +112,9 @@ pub struct AgentSnap { /// Display status when saved (`working` means a turn was interrupted). #[serde(default)] pub state: String, + /// The agent's own session reference (pi: its session file). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session: Option<String>, } /// One window's layout: session, index, name, layout, and pane directories. @@ -173,15 +191,29 @@ pub fn capture(now_ms: u64) -> Option<Snapshot> { "#{@agent}", "#{@agent_name}", "#{@agent_state}", + "#{pane_id}", + "#{pane_pid}", + "#{pane_current_command}", + "#{@agent_session}", + "#{@dingger_cmd}", + "#{pane_dead}", ] .join("\t"); let out = tmux::run(&["list-panes", "-a", "-F", &format]).ok()?; + let lines: Vec<Vec<&str>> = out + .lines() + .map(|line| line.split('\t').collect::<Vec<_>>()) + .filter(|f| f.len() >= 21) + .collect(); + // What runs in the foreground of each pane's terminal. + let pane_pids: Vec<u32> = lines + .iter() + .filter(|f| f[20] != "1") + .filter_map(|f| f[16].parse().ok()) + .collect(); + let foreground = foreground_groups(&pane_pids); let mut sessions: Vec<SessionSnap> = Vec::new(); - for line in out.lines() { - let f: Vec<&str> = line.split('\t').collect(); - if f.len() < 15 { - continue; - } + for f in lines { let number = |s: &str| s.parse::<u32>().unwrap_or(0); if sessions.last().map(|s| s.name != f[0]).unwrap_or(true) { sessions.push(SessionSnap { name: f[0].to_string(), windows: Vec::new() }); @@ -202,20 +234,82 @@ pub fn capture(now_ms: u64) -> Option<Snapshot> { }); } let window = session.windows.last_mut()?; + let agent = (!f[12].is_empty()).then(|| AgentSnap { + agent: f[12].to_string(), + name: f[13].to_string(), + state: f[14].to_string(), + session: (!f[18].is_empty()).then(|| f[18].to_string()), + }); + let hook_command = (!f[19].is_empty()).then(|| f[19].to_string()); + let pane_pid: Option<u32> = f[16].parse().ok(); + let leader = pane_pid.and_then(|pid| foreground.get(&pid).copied()); + // At a prompt: the pane's own shell is in the foreground. Anything + // else is a running program, whatever its name. + let at_prompt = leader.is_some() && leader == pane_pid && is_shell(f[17]); + let command = if f[20] == "1" || leader.is_none() || at_prompt { + None // the hook's text may be stale at a prompt + } else if let Some(agent) = &agent { + // Resumed by session; otherwise retyped as it was started. + agent.session.is_none().then(|| hook_command.unwrap_or_else(|| agent.agent.clone())) + } else { + hook_command.or_else(|| leader.and_then(command_line)) + }; window.panes.push(PaneSnap { + id: f[15].to_string(), index: number(f[9]), cwd: f[11].to_string(), active: f[10] == "1", - agent: (!f[12].is_empty()).then(|| AgentSnap { - agent: f[12].to_string(), - name: f[13].to_string(), - state: f[14].to_string(), - }), + agent, + command, }); } (!sessions.is_empty()).then_some(Snapshot { version: 1, saved_at: now_ms, server, sessions }) } +pub fn is_shell(command: &str) -> bool { + matches!( + command.trim_start_matches('-'), + "zsh" | "bash" | "fish" | "sh" | "dash" | "ksh" | "tcsh" | "csh" | "nu" | "elvish" | "xonsh" + ) +} + +/// Quote for a POSIX shell, only when needed. +pub fn shell_quote(word: &str) -> String { + let plain = !word.is_empty() + && word.chars().all(|c| c.is_ascii_alphanumeric() || "_@%+=:,./-~".contains(c)); + if plain { + word.to_string() + } else { + format!("'{}'", word.replace('\'', "'\\''")) + } +} + +/// For each pane process: the leader of its terminal's foreground process +/// group (itself, for a shell at its prompt). One `ps` call. +fn foreground_groups(pane_pids: &[u32]) -> std::collections::HashMap<u32, u32> { + let mut groups = std::collections::HashMap::new(); + if pane_pids.is_empty() { + return groups; + } + let list = pane_pids.iter().map(u32::to_string).collect::<Vec<_>>().join(","); + let Ok(out) = std::process::Command::new("ps").args(["-o", "pid=,tpgid=", "-p", &list]).output() else { + return groups; + }; + for line in String::from_utf8_lossy(&out.stdout).lines() { + let mut fields = line.split_whitespace().filter_map(|f| f.parse::<i64>().ok()); + if let (Some(pid), Some(group)) = (fields.next(), fields.next()) { + groups.insert(pid as u32, if group > 0 { group as u32 } else { pid as u32 }); + } + } + groups +} + +/// A process's command line, quoted for retyping. +fn command_line(pid: u32) -> Option<String> { + let argv = crate::platform::process_argv(pid).filter(|argv| !argv.is_empty())?; + Some(argv.iter().map(|word| shell_quote(word)).collect::<Vec<_>>().join(" ")) +} + fn write_atomic(path: &Path, value: &impl Serialize) -> Result<(), String> { let json = serde_json::to_string_pretty(value).map_err(|err| err.to_string())?; if let Some(parent) = path.parent() { @@ -304,6 +398,8 @@ pub fn should_pin(latest: &Snapshot, running: Option<&ServerId>, current: Option pub struct Recorder { latest: Option<(PathBuf, Snapshot)>, loaded: bool, + /// Per pane id: its last agent or command, and when it was last seen. + recent: std::collections::HashMap<String, (Option<AgentSnap>, Option<String>, u64)>, } impl Recorder { @@ -324,9 +420,10 @@ impl Recorder { } } } - let Some(fresh) = running.and_then(|_| capture(now_ms)) else { + let Some(mut fresh) = running.and_then(|_| capture(now_ms)) else { return; }; + self.remember_recent(&mut fresh, now_ms); let action = save_action(self.latest.as_ref().map(|(_, s)| s), &fresh); let path = match (&action, &self.latest) { (SaveAction::Nothing, _) => return, @@ -348,10 +445,54 @@ impl Recorder { } } +impl Recorder { + /// Keep a pane's agent or command for [`GRACE_MS`] after it exits. + fn remember_recent(&mut self, snapshot: &mut Snapshot, now_ms: u64) { + let mut seen = std::collections::HashSet::new(); + for pane in snapshot.sessions.iter_mut().flat_map(|s| &mut s.windows).flat_map(|w| &mut w.panes) { + seen.insert(pane.id.clone()); + if pane.agent.is_some() || pane.command.is_some() { + self.recent.insert(pane.id.clone(), (pane.agent.clone(), pane.command.clone(), now_ms)); + } else if let Some((agent, command, at)) = self.recent.get(&pane.id) { + if now_ms.saturating_sub(*at) < GRACE_MS { + pane.agent = agent.clone(); + pane.command = command.clone(); + } + } + } + self.recent.retain(|id, (_, _, at)| seen.contains(id) && now_ms.saturating_sub(*at) < GRACE_MS); + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn quotes_only_when_needed() { + assert_eq!(shell_quote("nvim"), "nvim"); + assert_eq!(shell_quote("~/src/a.sh"), "~/src/a.sh"); + assert_eq!(shell_quote("my file"), "'my file'"); + assert_eq!(shell_quote("it's"), "'it'\\''s'"); + assert_eq!(shell_quote(""), "''"); + } + + #[test] + fn an_exited_agent_is_kept_for_the_grace_period() { + let mut recorder = Recorder::default(); + let mut with_agent = snap(1, &[("a", "L")], "working"); + with_agent.sessions[0].windows[0].panes[0].id = "%1".into(); + recorder.remember_recent(&mut with_agent, 1_000); + let mut gone = with_agent.clone(); + gone.sessions[0].windows[0].panes[0].agent = None; + let mut soon = gone.clone(); + recorder.remember_recent(&mut soon, 1_000 + GRACE_MS - 1); + assert!(soon.sessions[0].windows[0].panes[0].agent.is_some(), "kept within the grace period"); + let mut later = gone.clone(); + recorder.remember_recent(&mut later, 1_000 + GRACE_MS); + assert!(later.sessions[0].windows[0].panes[0].agent.is_none(), "dropped after it"); + } + fn snap(pid: u32, windows: &[(&str, &str)], state: &str) -> Snapshot { Snapshot { version: 1, @@ -372,10 +513,17 @@ mod tests { zoomed: false, automatic_rename: false, panes: vec![PaneSnap { + id: String::new(), index: 0, cwd: "/tmp".into(), active: true, - agent: Some(AgentSnap { agent: "pi".into(), name: String::new(), state: state.into() }), + agent: Some(AgentSnap { + agent: "pi".into(), + name: String::new(), + state: state.into(), + session: None, + }), + command: None, }], }) .collect(), |