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

//! `dingger restore`: rebuild tmux sessions from a snapshot.
//!
//! Only adds what is missing. A session that exists is left alone, except a
//! bare one (one window, one pane, only a shell): that is what a terminal
//! creates on startup (Ghostty's `tmux new-session -As 0`), so the saved
//! session replaces it and its clients move over without detaching.
//!
//! Each session is rebuilt with its windows (index, name, automatic-rename),
//! panes (directory, in order), exact layout, zoom, and active window and
//! pane. Everything is addressed by tmux id, so names with `.` or `:` work.

use crate::snapshot::{self, SessionSnap, Snapshot, WindowSnap};
use crate::tmux;

/// What restore does with a saved session.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Plan {
    Create,
    /// Replace an existing bare session of the same name.
    Replace,
    /// It exists and has been used: leave it alone.
    Keep,
}

/// A session on the running server.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Existing {
    pub name: String,
    /// One window with one pane that runs only a shell.
    pub bare: bool,
}

pub fn plan(snapshot: &Snapshot, existing: &[Existing]) -> Vec<(String, Plan)> {
    snapshot
        .sessions
        .iter()
        .map(|session| {
            let plan = match existing.iter().find(|e| e.name == session.name) {
                None => Plan::Create,
                // A saved single shell has nothing to add over a bare one.
                Some(e) if e.bare && !saved_bare(session) => Plan::Replace,
                Some(_) => Plan::Keep,
            };
            (session.name.clone(), plan)
        })
        .collect()
}

/// Sessions restore would create or replace.
pub fn missing(snapshot: &Snapshot, existing: &[Existing]) -> Vec<String> {
    plan(snapshot, existing)
        .into_iter()
        .filter(|(_, plan)| *plan != Plan::Keep)
        .map(|(name, _)| name)
        .collect()
}

fn saved_bare(session: &SessionSnap) -> bool {
    matches!(&session.windows[..], [window] if matches!(&window.panes[..], [pane] if pane.agent.is_none()))
}

fn is_shell(command: &str) -> bool {
    matches!(
        command.trim_start_matches('-'),
        "zsh" | "bash" | "fish" | "sh" | "dash" | "ksh" | "tcsh" | "csh" | "nu" | "elvish" | "xonsh"
    )
}

/// The running server's sessions; empty when tmux is not running.
pub fn existing_sessions() -> Vec<Existing> {
    let Ok(out) = tmux::run(&[
        "list-panes",
        "-a",
        "-F",
        "#{session_name}\t#{session_windows}\t#{window_panes}\t#{pane_current_command}",
    ]) else {
        return Vec::new();
    };
    let mut sessions: Vec<Existing> = Vec::new();
    for line in out.lines() {
        let f: Vec<&str> = line.split('\t').collect();
        if f.len() < 4 || sessions.iter().any(|s| s.name == f[0]) {
            continue;
        }
        sessions.push(Existing {
            name: f[0].to_string(),
            bare: f[1] == "1" && f[2] == "1" && is_shell(f[3]),
        });
    }
    sessions
}

fn first_line(out: String) -> String {
    out.lines().next().unwrap_or_default().trim().to_string()
}

/// Build one window's panes, layout, and options. `pane` is its first pane.
fn build_window(window_id: &str, first_pane: String, window: &WindowSnap) -> Result<(), String> {
    let width = window.width.max(10).to_string();
    let height = window.height.max(5).to_string();
    // The saved size, so the layout applies exactly; unset afterwards so the
    // window follows its clients again.
    tmux::run(&["resize-window", "-t", window_id, "-x", &width, "-y", &height])?;
    let mut panes = vec![first_pane];
    for pane in window.panes.iter().skip(1) {
        // Split the newest pane: tmux lists a split right after its source,
        // so panes keep the saved order. Rebalance to make room for the next.
        let last = panes.last().cloned().unwrap_or_default();
        let id = first_line(tmux::run(&[
            "split-window", "-d", "-t", &last, "-c", &pane.cwd, "-P", "-F", "#{pane_id}",
        ])?);
        panes.push(id);
        let _ = tmux::run(&["select-layout", "-t", window_id, "tiled"]);
    }
    if panes.len() > 1 {
        if let Err(err) = tmux::run(&["select-layout", "-t", window_id, &window.layout]) {
            tracing::warn!(%err, "saved layout did not apply; keeping a tiled layout");
        }
    }
    tmux::run(&["set-option", "-w", "-u", "-t", window_id, "window-size"])?;
    if window.automatic_rename {
        tmux::run(&["set-option", "-w", "-t", window_id, "automatic-rename", "on"])?;
    }
    if let Some(position) = window.panes.iter().position(|pane| pane.active) {
        if let Some(id) = panes.get(position) {
            tmux::run(&["select-pane", "-t", id])?;
            if window.zoomed && panes.len() > 1 {
                tmux::run(&["resize-pane", "-Z", "-t", id])?;
            }
        }
    }
    Ok(())
}

/// Create `session` under `name`; returns its session id.
fn build_session(session: &SessionSnap, name: &str) -> Result<String, String> {
    let Some(first) = session.windows.first() else {
        return Err(format!("{}: no windows saved", session.name));
    };
    let cwd = first.panes.first().map(|pane| pane.cwd.as_str()).unwrap_or("~");
    let created = first_line(tmux::run(&[
        "new-session", "-d", "-s", name, "-n", &first.name, "-c", cwd,
        "-x", &first.width.max(10).to_string(), "-y", &first.height.max(5).to_string(),
        "-P", "-F", "#{session_id}\t#{window_id}\t#{window_index}\t#{pane_id}",
    ])?);
    let ids: Vec<&str> = created.split('\t').collect();
    let [session_id, first_window, first_index, first_pane] = ids[..] else {
        return Err(format!("unexpected new-session output: {created}"));
    };
    let session_id = session_id.to_string();
    if first_index != first.index.to_string() {
        tmux::run(&["move-window", "-s", first_window, "-t", &format!("{session_id}:{}", first.index)])?;
    }
    build_window(first_window, first_pane.to_string(), first)?;
    for window in session.windows.iter().skip(1) {
        let cwd = window.panes.first().map(|pane| pane.cwd.as_str()).unwrap_or("~");
        let created = first_line(tmux::run(&[
            "new-window", "-d", "-t", &format!("{session_id}:{}", window.index), "-n", &window.name,
            "-c", cwd, "-P", "-F", "#{window_id}\t#{pane_id}",
        ])?);
        let Some((window_id, pane_id)) = created.split_once('\t') else {
            return Err(format!("unexpected new-window output: {created}"));
        };
        build_window(window_id, pane_id.to_string(), window)?;
    }
    if let Some(active) = session.windows.iter().find(|window| window.active) {
        tmux::run(&["select-window", "-t", &format!("{session_id}:{}", active.index)])?;
    }
    Ok(session_id)
}

/// Replace a bare session: build the saved one beside it, move its clients
/// over, then remove it and take its name.
fn replace_session(session: &SessionSnap) -> Result<(), String> {
    let temporary = format!("{}-restoring-{}", session.name, std::process::id());
    let new_id = build_session(session, &temporary)?;
    let clients = tmux::run(&["list-clients", "-F", "#{client_name}\t#{session_name}"]).unwrap_or_default();
    for line in clients.lines() {
        if let Some((client, attached)) = line.split_once('\t') {
            if attached == session.name {
                tmux::run(&["switch-client", "-c", client, "-t", &new_id])?;
            }
        }
    }
    tmux::run(&["kill-session", "-t", &format!("={}", session.name)])?;
    tmux::run(&["rename-session", "-t", &new_id, &session.name]).map(|_| ())
}

fn windows_and_panes(session: &SessionSnap) -> String {
    let panes: usize = session.windows.iter().map(|window| window.panes.len()).sum();
    format!("{} windows, {panes} panes", session.windows.len())
}

/// The snapshot to restore: `from`, or the restore candidate.
fn source(from: Option<&str>) -> Result<Snapshot, String> {
    match from {
        Some(path) => snapshot::load(std::path::Path::new(path)).ok_or_else(|| format!("cannot read {path}")),
        None => snapshot::candidate()
            .map(|candidate| candidate.snapshot)
            .ok_or_else(|| "nothing to restore: no saved layout from an earlier tmux server".to_string()),
    }
}

pub fn list() -> Result<(), String> {
    match snapshot::candidate() {
        Some(candidate) => {
            let missing = missing(&candidate.snapshot, &existing_sessions());
            println!(
                "restore candidate ({}), saved {}: {}; missing now: {}",
                candidate.status,
                format_time(candidate.snapshot.saved_at),
                candidate.snapshot.summary(),
                if missing.is_empty() { "none".to_string() } else { missing.join(", ") }
            );
        }
        None => println!("no restore candidate"),
    }
    for path in snapshot::history().iter().rev() {
        if let Some(snapshot) = snapshot::load(path) {
            println!("{}  {}  {}", path.display(), format_time(snapshot.saved_at), snapshot.summary());
        }
    }
    Ok(())
}

fn format_time(ms: u64) -> String {
    std::process::Command::new("date")
        .args(["-r", &(ms / 1000).to_string(), "+%Y-%m-%d %H:%M"])
        .output()
        .map(|out| String::from_utf8_lossy(&out.stdout).trim().to_string())
        .unwrap_or_else(|_| ms.to_string())
}

/// Restore missing sessions (all, or only `only`).
pub fn run(from: Option<&str>, only: &[String], dry_run: bool) -> Result<(), String> {
    let snapshot = source(from)?;
    let existing = existing_sessions();
    let mut failures = 0;
    for (name, plan) in plan(&snapshot, &existing) {
        if !only.is_empty() && !only.contains(&name) {
            continue;
        }
        let Some(session) = snapshot.sessions.iter().find(|s| s.name == name) else {
            continue;
        };
        let what = windows_and_panes(session);
        let verb = match plan {
            Plan::Keep => {
                println!("keep     {name}: it exists");
                continue;
            }
            Plan::Create => "create",
            Plan::Replace => "replace",
        };
        if dry_run {
            println!("{verb:<8} {name}: {what}");
            continue;
        }
        let result = match plan {
            Plan::Replace => replace_session(session),
            _ => build_session(session, &name).map(|_| ()),
        };
        match result {
            Ok(()) => println!("{verb:<8} {name}: {what}"),
            Err(err) => {
                failures += 1;
                println!("FAILED   {name}: {err}");
            }
        }
    }
    if dry_run || from.is_some() {
        return if failures == 0 { Ok(()) } else { Err(format!("{failures} sessions failed")) };
    }
    if missing(&snapshot, &existing_sessions()).is_empty() {
        let _ = snapshot::set_candidate_status("restored");
    }
    if failures == 0 { Ok(()) } else { Err(format!("{failures} sessions failed")) }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::snapshot::{PaneSnap, ServerId, SessionSnap};

    fn window(panes: usize) -> WindowSnap {
        WindowSnap {
            index: 0,
            name: "w".into(),
            layout: String::new(),
            width: 80,
            height: 24,
            active: true,
            zoomed: false,
            automatic_rename: false,
            panes: (0..panes)
                .map(|i| PaneSnap { index: i as u32, cwd: "/tmp".into(), active: i == 0, agent: None })
                .collect(),
        }
    }

    fn snapshot(names: &[&str]) -> Snapshot {
        Snapshot {
            version: 1,
            saved_at: 1,
            server: ServerId { pid: 1, start_time: 1 },
            sessions: names
                .iter()
                .map(|name| SessionSnap { name: name.to_string(), windows: vec![window(2)] })
                .collect(),
        }
    }

    #[test]
    fn creates_missing_replaces_bare_and_keeps_used_sessions() {
        let existing = [
            Existing { name: "0".into(), bare: true },
            Existing { name: "keep".into(), bare: false },
            Existing { name: "new".into(), bare: true },
        ];
        assert_eq!(
            plan(&snapshot(&["0", "keep", "work"]), &existing),
            vec![("0".into(), Plan::Replace), ("keep".into(), Plan::Keep), ("work".into(), Plan::Create)]
        );
        assert_eq!(missing(&snapshot(&["0", "keep", "work"]), &existing), vec!["0", "work"]);
        assert!(missing(&snapshot(&["keep"]), &existing).is_empty());
    }

    #[test]
    fn a_saved_single_shell_does_not_replace_a_bare_session() {
        let mut saved = snapshot(&["0"]);
        saved.sessions[0].windows = vec![window(1)];
        let existing = [Existing { name: "0".into(), bare: true }];
        assert_eq!(plan(&saved, &existing), vec![("0".into(), Plan::Keep)]);
        assert!(missing(&saved, &existing).is_empty(), "or the offer would never go away");
    }

    #[test]
    fn shells_are_recognized() {
        assert!(is_shell("zsh") && is_shell("-zsh") && is_shell("bash") && is_shell("fish"));
        assert!(!is_shell("nvim") && !is_shell("node") && !is_shell("tmux"));
    }
}