path: root/crates/dinggerd/src/snapshot.rs

//! Snapshots of the tmux server's layout, for restore after a crash or
//! shutdown.
//!
//! The watcher captures sessions, windows, and panes every tick and saves a
//! snapshot when something changed:
//!
//! - `snapshots/<saved ms>.json`: a new file when the layout changes, the
//!   current file rewritten when only details change (agent states); the
//!   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
//!   snapshots start, so they never replace it.

use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};
use tracing::{info, warn};

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)
        .or_else(|| {
            std::env::var_os("HOME")
                .map(|home| PathBuf::from(home).join(".local").join("state").join("dingger"))
        })
        .unwrap_or_else(|| PathBuf::from(".dingger"))
}

fn snapshot_dir() -> PathBuf {
    state_dir().join("snapshots")
}

fn candidate_path() -> PathBuf {
    state_dir().join("restore.json")
}

/// A tmux server instance: its pid and start time (pids are reused).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ServerId {
    pub pid: u32,
    pub start_time: u64,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Snapshot {
    pub version: u32,
    /// Unix milliseconds.
    pub saved_at: u64,
    pub server: ServerId,
    pub sessions: Vec<SessionSnap>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionSnap {
    pub name: String,
    pub windows: Vec<WindowSnap>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WindowSnap {
    pub index: u32,
    pub name: String,
    /// tmux's layout string (`select-layout` reapplies it).
    pub layout: String,
    pub width: u32,
    pub height: u32,
    pub active: bool,
    pub zoomed: bool,
    pub automatic_rename: bool,
    pub panes: Vec<PaneSnap>,
}

#[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: retyped on restore, not run
    /// (an agent's, when its session cannot be resumed).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub command: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentSnap {
    /// Harness, e.g. `pi`.
    pub agent: String,
    #[serde(default)]
    pub name: String,
    /// 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.
type WindowKey = (String, u32, String, String, Vec<(u32, String)>);

impl Snapshot {
    /// The layout alone: what decides between a new file and a rewrite.
    fn layout_key(&self) -> Vec<WindowKey> {
        self.sessions
            .iter()
            .flat_map(|session| {
                session.windows.iter().map(|window| {
                    (
                        session.name.clone(),
                        window.index,
                        window.name.clone(),
                        window.layout.clone(),
                        window.panes.iter().map(|pane| (pane.index, pane.cwd.clone())).collect(),
                    )
                })
            })
            .collect()
    }

    /// Same content, ignoring when it was saved.
    fn same_content(&self, other: &Snapshot) -> bool {
        self.server == other.server && self.sessions == other.sessions
    }

    pub fn summary(&self) -> String {
        let windows: usize = self.sessions.iter().map(|s| s.windows.len()).sum();
        let panes: usize = self.sessions.iter().flat_map(|s| &s.windows).map(|w| w.panes.len()).sum();
        let agents = self
            .sessions
            .iter()
            .flat_map(|s| &s.windows)
            .flat_map(|w| &w.panes)
            .filter(|p| p.agent.is_some())
            .count();
        format!(
            "{} sessions, {windows} windows, {panes} panes, {agents} agents",
            self.sessions.len()
        )
    }
}

/// The running server, or `None` when tmux is not running.
pub fn server_id() -> Option<ServerId> {
    let out = tmux::run(&["display-message", "-p", "#{pid}\t#{start_time}"]).ok()?;
    let (pid, start) = out.trim().split_once('\t')?;
    Some(ServerId {
        pid: pid.parse().ok()?,
        start_time: start.parse().ok()?,
    })
}

/// Capture the running server's layout; `None` when tmux is not running or
/// has no sessions.
pub fn capture(now_ms: u64) -> Option<Snapshot> {
    let server = server_id()?;
    let format = [
        "#{session_name}",
        "#{window_index}",
        "#{window_name}",
        "#{window_layout}",
        "#{window_width}",
        "#{window_height}",
        "#{window_active}",
        "#{window_zoomed_flag}",
        "#{automatic-rename}",
        "#{pane_index}",
        "#{pane_active}",
        "#{pane_current_path}",
        "#{@agent}",
        "#{@agent_name}",
        "#{@agent_state}",
        "#{pane_id}",
        "#{pane_pid}",
        "#{pane_current_command}",
        "#{@agent_session}",
        "#{@dingger_cmd}",
        "#{pane_dead}",
        "#{?session_grouped,#{session_group},}",
    ]
    .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() >= 22)
        .collect();
    // One session per group: the others show the same windows, and saving
    // them would restore every window, and resume every agent, twice. The
    // first listed (tmux lists sessions by name) stands for its group.
    let mut group_of_kept: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
    let lines: Vec<Vec<&str>> = lines
        .into_iter()
        .filter(|f| f[21].is_empty() || *group_of_kept.entry(f[21]).or_insert(f[0]) == f[0])
        .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 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() });
        }
        let session = sessions.last_mut()?;
        let index = number(f[1]);
        if session.windows.last().map(|w| w.index != index).unwrap_or(true) {
            session.windows.push(WindowSnap {
                index,
                name: f[2].to_string(),
                layout: f[3].to_string(),
                width: number(f[4]),
                height: number(f[5]),
                active: f[6] == "1",
                zoomed: f[7] == "1",
                automatic_rename: f[8] == "1",
                panes: Vec::new(),
            });
        }
        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 {
            // Kept even with a session: restore retypes it when the session
            // cannot be resumed (it never got a message, or its file is gone).
            Some(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,
            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() {
        std::fs::create_dir_all(parent).map_err(|err| format!("mkdir {}: {err}", parent.display()))?;
    }
    let tmp = path.with_extension("tmp");
    std::fs::write(&tmp, json).map_err(|err| format!("write {}: {err}", tmp.display()))?;
    std::fs::rename(&tmp, path).map_err(|err| format!("rename {}: {err}", path.display()))
}

fn read_json<T: for<'de> Deserialize<'de>>(path: &Path) -> Option<T> {
    serde_json::from_str(&std::fs::read_to_string(path).ok()?).ok()
}

/// Saved snapshot files, oldest first.
pub fn history() -> Vec<PathBuf> {
    let mut files: Vec<PathBuf> = std::fs::read_dir(snapshot_dir())
        .map(|entries| {
            entries
                .filter_map(|entry| entry.ok().map(|entry| entry.path()))
                .filter(|path| path.extension().is_some_and(|ext| ext == "json"))
                .collect()
        })
        .unwrap_or_default();
    // Names are zero-padded milliseconds: lexical order is time order.
    files.sort();
    files
}

pub fn load(path: &Path) -> Option<Snapshot> {
    read_json(path)
}

/// What to do with a fresh capture, given the latest saved snapshot.
#[derive(Debug, PartialEq, Eq)]
pub enum SaveAction {
    Nothing,
    /// Rewrite the latest file: only details changed.
    Rewrite,
    /// A new history file: the layout changed, or a different server.
    New,
}

pub fn save_action(latest: Option<&Snapshot>, fresh: &Snapshot) -> SaveAction {
    match latest {
        None => SaveAction::New,
        Some(latest) if latest.same_content(fresh) => SaveAction::Nothing,
        Some(latest) if latest.server == fresh.server && latest.layout_key() == fresh.layout_key() => {
            SaveAction::Rewrite
        }
        Some(_) => SaveAction::New,
    }
}

/// The restore candidate and what happened to it.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Candidate {
    /// `offered`, `restored`, or `dismissed`.
    pub status: String,
    pub snapshot: Snapshot,
}

pub fn candidate() -> Option<Candidate> {
    read_json(&candidate_path())
}

pub fn set_candidate_status(status: &str) -> Result<(), String> {
    let Some(mut candidate) = candidate() else {
        return Err("no restore candidate".into());
    };
    candidate.status = status.to_string();
    write_atomic(&candidate_path(), &candidate)
}

/// Whether the latest snapshot should become the restore candidate: it
/// belongs to a server that is no longer running, and is not the candidate
/// already (whatever became of it).
pub fn should_pin(latest: &Snapshot, running: Option<&ServerId>, current: Option<&Candidate>) -> bool {
    running != Some(&latest.server)
        && current.map(|c| c.snapshot.server != latest.server).unwrap_or(true)
}

/// Keeps snapshots current; called by the watcher every loop, whether or not
/// tmux is running.
#[derive(Default)]
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 {
    pub fn tick(&mut self, now_ms: u64) {
        if !self.loaded {
            self.loaded = true;
            self.latest = history().last().and_then(|path| Some((path.clone(), load(path)?)));
        }
        let running = server_id();
        // Before saving anything for a new server: pin the old one's layout.
        if let Some((_, latest)) = self.latest.as_ref().filter(|(_, s)| running.as_ref() != Some(&s.server)) {
            let current = candidate();
            if should_pin(latest, running.as_ref(), current.as_ref()) {
                let candidate = Candidate { status: "offered".into(), snapshot: latest.clone() };
                match write_atomic(&candidate_path(), &candidate) {
                    Ok(()) => info!(snapshot = %latest.summary(), "restore candidate saved"),
                    Err(err) => warn!(%err, "failed to save the restore candidate"),
                }
            }
        }
        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,
            (SaveAction::Rewrite, Some((path, _))) => path.clone(),
            _ => snapshot_dir().join(format!("{:015}.json", fresh.saved_at)),
        };
        if let Err(err) = write_atomic(&path, &fresh) {
            warn!(%err, "failed to save snapshot");
            return;
        }
        if action == SaveAction::New {
            info!(snapshot = %fresh.summary(), "snapshot saved");
            let files = history();
            for old in files.iter().take(files.len().saturating_sub(HISTORY)) {
                let _ = std::fs::remove_file(old);
            }
        }
        self.latest = Some((path, fresh));
    }
}

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,
            saved_at: 1,
            server: ServerId { pid, start_time: 100 },
            sessions: vec![SessionSnap {
                name: "s".into(),
                windows: windows
                    .iter()
                    .enumerate()
                    .map(|(i, (name, layout))| WindowSnap {
                        index: i as u32,
                        name: name.to_string(),
                        layout: layout.to_string(),
                        width: 80,
                        height: 24,
                        active: i == 0,
                        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(),
                                session: None,
                            }),
                            command: None,
                        }],
                    })
                    .collect(),
            }],
        }
    }

    #[test]
    fn saves_new_files_for_layout_changes_and_rewrites_for_details() {
        let base = snap(1, &[("a", "L1")], "idle");
        assert_eq!(save_action(None, &base), SaveAction::New);
        let mut later = base.clone();
        later.saved_at = 99;
        assert_eq!(save_action(Some(&base), &later), SaveAction::Nothing, "time alone is no change");
        assert_eq!(save_action(Some(&base), &snap(1, &[("a", "L1")], "working")), SaveAction::Rewrite);
        assert_eq!(save_action(Some(&base), &snap(1, &[("a", "L2")], "idle")), SaveAction::New);
        assert_eq!(save_action(Some(&base), &snap(1, &[("a", "L1"), ("b", "L1")], "idle")), SaveAction::New);
        assert_eq!(save_action(Some(&base), &snap(2, &[("a", "L1")], "idle")), SaveAction::New, "new server");
    }

    #[test]
    fn pins_the_previous_servers_snapshot_once() {
        let old = snap(1, &[("a", "L1")], "idle");
        let old_server = old.server.clone();
        let new_server = ServerId { pid: 2, start_time: 200 };
        assert!(!should_pin(&old, Some(&old_server), None), "its server still runs");
        assert!(should_pin(&old, Some(&new_server), None), "a new server");
        assert!(should_pin(&old, None, None), "no server");
        let pinned = Candidate { status: "dismissed".into(), snapshot: old.clone() };
        assert!(!should_pin(&old, None, Some(&pinned)), "already the candidate, whatever its status");
        let older = Candidate { status: "restored".into(), snapshot: snap(0, &[], "") };
        assert!(should_pin(&old, None, Some(&older)), "replaces an older candidate");
    }
}