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.
//!
//! Then, once each new shell shows its prompt: agents with a session
//! reference are resumed (`pi --session <file>`, run), and every other
//! saved command line is typed but not run, for the user to press Enter.

use std::time::{Duration, Instant};

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

/// What restore types into a rebuilt pane.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Typed {
    /// Resume the agent's session: typed and run.
    Resume(String),
    /// The saved command line: typed, not run.
    Retype(String),
}

/// The command that resumes an agent's session.
fn resume_command(agent: &str, session: &str) -> Option<String> {
    match agent {
        "pi" => Some(format!("pi --session {}", shell_quote(session))),
        _ => None,
    }
}

pub fn typed_for(pane: &PaneSnap) -> Option<Typed> {
    let typeable = |text: &String| !text.trim().is_empty() && !text.chars().any(char::is_control);
    if let Some(agent) = &pane.agent {
        let resume = agent.session.as_deref().and_then(|session| resume_command(&agent.agent, session));
        if let Some(command) = resume.filter(typeable) {
            return Some(Typed::Resume(command));
        }
    }
    pane.command.clone().filter(typeable).map(Typed::Retype)
}

/// A rebuilt pane and what it held.
struct Rebuilt {
    pane_id: String,
    label: String,
    pane: PaneSnap,
}

/// 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()))
}

use snapshot::is_shell;

/// 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,
    session: &str,
    rebuilt: &mut Vec<Rebuilt>,
) -> 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])?;
            }
        }
    }
    for (pane_id, pane) in panes.into_iter().zip(&window.panes) {
        let label = if window.panes.len() > 1 {
            format!("{session}:{}.{}", window.name, pane.index)
        } else {
            format!("{session}:{}", window.name)
        };
        rebuilt.push(Rebuilt { pane_id, label, pane: pane.clone() });
    }
    Ok(())
}

/// Create `session` under `name`; returns its session id.
fn build_session(session: &SessionSnap, name: &str, rebuilt: &mut Vec<Rebuilt>) -> 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, &session.name, rebuilt)?;
    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, &session.name, rebuilt)?;
    }
    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, rebuilt: &mut Vec<Rebuilt>) -> Result<(), String> {
    let temporary = format!("{}-restoring-{}", session.name, std::process::id());
    let new_id = build_session(session, &temporary, rebuilt)?;
    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(|_| ())
}

/// Wait for each new shell's prompt, then resume agents (spaced out) and
/// type the other command lines without running them; report which agents
/// did not start.
fn type_commands(rebuilt: &[Rebuilt]) {
    let work: Vec<(&Rebuilt, Typed)> =
        rebuilt.iter().filter_map(|r| typed_for(&r.pane).map(|typed| (r, typed))).collect();
    if work.is_empty() {
        return;
    }
    // Text typed before the prompt is drawn is echoed twice; wait for it
    // (the cursor leaves the top-left corner).
    let deadline = Instant::now() + Duration::from_secs(10);
    loop {
        let cursors = tmux::run(&["list-panes", "-a", "-F", "#{pane_id} #{cursor_x} #{cursor_y}"]).unwrap_or_default();
        let waiting = work.iter().any(|(r, _)| {
            cursors.lines().any(|line| line == format!("{} 0 0", r.pane_id))
        });
        if !waiting || Instant::now() > deadline {
            break;
        }
        std::thread::sleep(Duration::from_millis(100));
    }
    std::thread::sleep(Duration::from_millis(200));

    let mut resumed = Vec::new();
    for (r, typed) in &work {
        match typed {
            Typed::Resume(command) => {
                let _ = tmux::run(&["send-keys", "-t", &r.pane_id, "-l", command]);
                let _ = tmux::run(&["send-keys", "-t", &r.pane_id, "Enter"]);
                println!("resume   {}  {command}", r.label);
                resumed.push(*r);
                // Agents start one at a time.
                std::thread::sleep(Duration::from_millis(500));
            }
            Typed::Retype(command) => {
                let _ = tmux::run(&["send-keys", "-t", &r.pane_id, "-l", command]);
                println!("type     {}  {command}  (press Enter to run)", r.label);
            }
        }
    }
    for (r, _) in &work {
        if let Some(agent) = r.pane.agent.as_ref().filter(|agent| agent.state == "working") {
            let name = if agent.name.is_empty() { &agent.agent } else { &agent.name };
            println!("interrupted  {}  {name}: it was working, its turn was cut short", r.label);
        }
    }
    if resumed.is_empty() {
        return;
    }
    // An agent that started has replaced the shell in the foreground.
    let deadline = Instant::now() + Duration::from_secs(30);
    let mut not_running: Vec<&Rebuilt> = resumed.clone();
    while !not_running.is_empty() && Instant::now() < deadline {
        std::thread::sleep(Duration::from_millis(500));
        not_running.retain(|r| {
            tmux::run(&["display-message", "-p", "-t", &r.pane_id, "#{pane_current_command}"])
                .map(|command| is_shell(command.trim()))
                .unwrap_or(true)
        });
    }
    println!("resumed {} of {} agents", resumed.len() - not_running.len(), resumed.len());
    for r in not_running {
        println!("not running  {}: check the pane for the error", r.label);
    }
}

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;
    let mut rebuilt = Vec::new();
    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}");
            for window in &session.windows {
                for pane in &window.panes {
                    match typed_for(pane) {
                        Some(Typed::Resume(command)) => println!("  resume  {}  {command}", window.name),
                        Some(Typed::Retype(command)) => println!("  type    {}  {command}", window.name),
                        None => {}
                    }
                }
            }
            continue;
        }
        let result = match plan {
            Plan::Replace => replace_session(session, &mut rebuilt),
            _ => build_session(session, &name, &mut rebuilt).map(|_| ()),
        };
        match result {
            Ok(()) => println!("{verb:<8} {name}: {what}"),
            Err(err) => {
                failures += 1;
                println!("FAILED   {name}: {err}");
            }
        }
    }
    if !rebuilt.is_empty() {
        type_commands(&rebuilt);
    }
    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::{AgentSnap, 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 {
                    id: String::new(),
                    index: i as u32,
                    cwd: "/tmp".into(),
                    active: i == 0,
                    agent: None,
                    command: 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 resumes_agents_by_session_and_retypes_the_rest() {
        let mut pane = window(1).panes.remove(0);
        assert_eq!(typed_for(&pane), None, "a shell prompt types nothing");
        pane.command = Some("nvim 'my file.sh'".into());
        assert_eq!(typed_for(&pane), Some(Typed::Retype("nvim 'my file.sh'".into())));
        pane.command = Some("echo one\necho two".into());
        assert_eq!(typed_for(&pane), None, "a newline would run it");
        pane.agent = Some(AgentSnap {
            agent: "pi".into(),
            name: "x".into(),
            state: "idle".into(),
            session: Some("/s/a b.jsonl".into()),
        });
        assert_eq!(typed_for(&pane), Some(Typed::Resume("pi --session '/s/a b.jsonl'".into())));
        if let Some(agent) = pane.agent.as_mut() {
            agent.session = None;
        }
        pane.command = Some("pi --model x".into());
        assert_eq!(typed_for(&pane), Some(Typed::Retype("pi --model x".into())), "no session: retyped");
    }

    #[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"));
    }
}