agent-berth 2.2.1

Monitor and resume coding agents
Documentation
use std::io::Write;
use std::path::Path;
use std::process::{Command, Stdio};

use anyhow::{Context, Result, bail};

const DEFAULT_CONFIG: &str = ".tmux-up.conf";

#[cfg(windows)]
pub(crate) const PREVIEW_SHELL: &str = "powershell.exe -NoLogo -NoProfile -NonInteractive -Command";

pub fn session_name(root: &Path) -> String {
    let mut session = root
        .file_name()
        .map(|name| name.to_string_lossy().into_owned())
        .unwrap_or_else(|| "agent-berth".into());
    if session.starts_with('.') {
        session.remove(0);
    }
    session.replace(['.', ':'], "_")
}

pub fn ensure_session(root: &Path) -> Result<String> {
    if !root.is_dir() {
        bail!("{} does not exist", root.display());
    }
    let session = session_name(root);
    let target = format!("={session}");
    if tmux(["has-session", "-t", &target])?.success() {
        return Ok(session);
    }
    let status = tmux([
        "new",
        "-d",
        "-c",
        &root.display().to_string(),
        "-s",
        &session,
    ])?;
    if !status.success() {
        bail!("tmux new-session failed");
    }
    apply_config(root, &target)?;
    Ok(session)
}

pub fn new_window(session: &str, name: &str, cwd: &Path, cmd: &[String]) -> Result<()> {
    spawn_window(session, name, cwd, cmd)?;
    Ok(())
}

/// Like [`new_window`] but returns the id of the created pane.
pub fn spawn_window(session: &str, name: &str, cwd: &Path, cmd: &[String]) -> Result<String> {
    let target = format!("={session}");
    let mut args = vec![
        "new-window".into(),
        "-P".into(),
        "-F".into(),
        "#{pane_id}".into(),
        "-t".into(),
        target,
    ];
    args.extend(window_command(name, cwd, cmd));
    let pane_id = spawn(&args)?.context("tmux new-window failed for {name}")?;
    Ok(pane_id)
}

/// Create a window running `cmd` in the default session — the current one
/// inside tmux, the most recently used otherwise — creating a session when
/// none exists, and return the new pane id.
pub fn spawn_window_default(name: &str, cwd: &Path, cmd: &[String]) -> Result<String> {
    let tail = window_command(name, cwd, cmd);
    let mut args = vec![
        "new-window".into(),
        "-P".into(),
        "-F".into(),
        "#{pane_id}".into(),
    ];
    args.extend(passthrough_env());
    args.extend(tail.iter().cloned());
    if let Some(pane_id) = spawn(&args)? {
        return Ok(pane_id);
    }
    // No session exists yet; create one whose first window runs the command.
    let mut args = vec![
        "new-session".into(),
        "-d".into(),
        "-P".into(),
        "-F".into(),
        "#{pane_id}".into(),
    ];
    args.extend(passthrough_env());
    args.extend(tail);
    spawn(&args)?.context("tmux could not create a window")
}

fn window_command(name: &str, cwd: &Path, cmd: &[String]) -> Vec<String> {
    let mut args = vec![
        "-n".into(),
        name.into(),
        "-c".into(),
        cwd.display().to_string(),
        "--".into(),
    ];
    args.extend(cmd.iter().cloned());
    args
}

/// `-e` flags forwarding AGENT_BERTH_* overrides to the new window: the tmux
/// server passes its own (startup-time) environment to the processes it
/// spawns, not the caller's.
fn passthrough_env() -> Vec<String> {
    let mut vars: Vec<String> = std::env::vars()
        .filter(|(key, _)| key.starts_with("AGENT_BERTH_"))
        .map(|(key, value)| format!("{key}={value}"))
        .collect();
    vars.sort();
    vars.into_iter()
        .flat_map(|var| ["-e".into(), var])
        .collect()
}

fn spawn(args: &[String]) -> Result<Option<String>> {
    let output = command()
        .args(args)
        .output()
        .context("tmux is not available")?;
    if !output.status.success() {
        return Ok(None);
    }
    let pane_id = String::from_utf8_lossy(&output.stdout).trim().to_string();
    if pane_id.is_empty() {
        bail!("tmux did not report a pane id");
    }
    Ok(Some(pane_id))
}

pub fn attach_or_switch(session: &str) -> Result<()> {
    let target = format!("={session}");
    if std::env::var_os("TMUX").is_some() {
        let status = tmux(["switch-client", "-t", &target])?;
        if !status.success() {
            bail!("tmux switch-client failed");
        }
        return Ok(());
    }
    let status = command()
        .args(["attach", "-t", &target])
        .status()
        .context("tmux attach")?;
    if !status.success() {
        bail!("tmux attach failed");
    }
    Ok(())
}

pub fn available() -> bool {
    crate::paths::on_path("tmux")
}

const PANE_FORMAT: &str = "#{pane_id}\t#{pane_pid}\t#{session_name}\t#{window_index}\t#{window_name}\t#{pane_current_path}";

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Pane {
    pub id: String,
    pub pid: u32,
    pub session: String,
    pub window: String,
    pub window_name: String,
    pub path: String,
}

impl Pane {
    fn parse(line: &str) -> Option<Self> {
        let mut fields = line.split('\t');
        let id = fields.next()?.trim().to_string();
        let pid = fields.next()?.trim().parse().ok()?;
        let session = fields.next()?.to_string();
        let window = fields.next()?.to_string();
        let window_name = fields.next()?.to_string();
        let path = fields.next()?.to_string();
        if id.is_empty() {
            return None;
        }
        Some(Self {
            id,
            pid,
            session,
            window,
            window_name,
            path,
        })
    }
}

pub fn list_panes(all: bool) -> Result<Vec<Pane>> {
    let scope = if all { "-a" } else { "-s" };
    let output = command()
        .args(["list-panes", scope, "-F", PANE_FORMAT])
        .output()
        .context("tmux is not available")?;
    if !output.status.success() {
        bail!("tmux list-panes failed");
    }
    let text = String::from_utf8_lossy(&output.stdout);
    Ok(text.lines().filter_map(Pane::parse).collect())
}

pub fn find_pane(id: &str) -> Result<Option<Pane>> {
    Ok(list_panes(true)?.into_iter().find(|pane| pane.id == id))
}

/// Visible content of a pane, including ANSI colors and attributes.
pub fn capture_pane(id: &str) -> Result<String> {
    let output = command()
        .args(["capture-pane", "-p", "-e", "-t", id])
        .output()
        .context("tmux is not available")?;
    if !output.status.success() {
        bail!("tmux capture-pane failed for {id}");
    }
    Ok(String::from_utf8_lossy(&output.stdout).into_owned())
}

pub fn attach_pane(pane: &Pane) -> Result<()> {
    let session = format!("={}", pane.session);
    let window = format!("{session}:{}", pane.window);
    if std::env::var_os("TMUX").is_some() {
        run(["switch-client", "-t", &session])?;
        run(["select-window", "-t", &window])?;
        run(["select-pane", "-t", &pane.id])?;
        return Ok(());
    }
    run(["select-window", "-t", &window])?;
    run(["select-pane", "-t", &pane.id])?;
    let status = command()
        .args(["attach", "-t", &session])
        .status()
        .context("tmux attach")?;
    if !status.success() {
        bail!("tmux attach failed");
    }
    Ok(())
}

fn run<I, S>(args: I) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: AsRef<std::ffi::OsStr>,
{
    if !tmux(args)?.success() {
        bail!("tmux command failed");
    }
    Ok(())
}

/// fzf preview command for `{1}`, the pane id, including the active socket and config.
///
/// The output is trimmed to the preview height so fzf, which opens previews
/// scrolled to the top, shows the pane bottom (last message and status bar).
pub fn preview_command() -> String {
    let mut parts = vec!["tmux".to_string(), "-u".to_string()];
    if let Some(socket) = crate::paths::env_path("AGENT_BERTH_TMUX_SOCKET") {
        parts.push("-L".into());
        parts.push(quote_arg(&socket.display().to_string()));
    }
    if let Some(config) = crate::paths::env_path("AGENT_BERTH_TMUX_CONFIG") {
        parts.push("-f".into());
        parts.push(quote_arg(&config.display().to_string()));
    }
    parts.push("capture-pane -p -e -t {1}".into());
    let capture = parts.join(" ");
    #[cfg(windows)]
    {
        format!(
            "[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false); \
             $lines = if ($env:FZF_PREVIEW_LINES) {{ [int]$env:FZF_PREVIEW_LINES }} else {{ 40 }}; \
             {capture} | Select-Object -Last $lines"
        )
    }
    #[cfg(not(windows))]
    {
        format!("{capture} | tail -n \"${{FZF_PREVIEW_LINES:-40}}\"")
    }
}

#[cfg(windows)]
fn quote_arg(value: &str) -> String {
    format!("'{}'", value.replace('\'', "''"))
}

#[cfg(not(windows))]
fn quote_arg(value: &str) -> String {
    if value.chars().any(char::is_whitespace) {
        format!("\"{}\"", value.replace('"', "\\\""))
    } else {
        value.to_string()
    }
}

fn apply_config(root: &Path, target: &str) -> Result<()> {
    let config = root.join(DEFAULT_CONFIG);
    let mut commands = if config.is_file() {
        let mut text = std::fs::read_to_string(&config)?;
        if !text.ends_with('\n') {
            text.push('\n');
        }
        text.push('\n');
        text
    } else {
        let editor = std::env::var("EDITOR").unwrap_or_else(|_| "vim".into());
        format!("send '{editor}' Enter\nneww -n shell\nselectw -t 1\n\n")
    };
    commands.push_str("detach-client\n");
    let mut child = command()
        .args(["-C", "attach", "-t", target])
        .stdin(Stdio::piped())
        .stdout(Stdio::null())
        .stderr(Stdio::inherit())
        .env("TMUX", "")
        .spawn()
        .context("tmux control mode")?;
    if let Some(stdin) = child.stdin.as_mut() {
        stdin.write_all(commands.as_bytes())?;
    }
    let status = child.wait()?;
    if !status.success() {
        bail!("tmux control mode failed");
    }
    Ok(())
}

#[cfg(test)]
#[path = "tmux_tests.rs"]
mod tests;

fn tmux<I, S>(args: I) -> Result<std::process::ExitStatus>
where
    I: IntoIterator<Item = S>,
    S: AsRef<std::ffi::OsStr>,
{
    command()
        .args(args)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .status()
        .context("tmux is not available")
}

fn command() -> Command {
    let mut command = Command::new("tmux");
    // Force UTF-8 output; a non-UTF-8 locale makes tmux replace the tabs in
    // format output with underscores.
    command.arg("-u");
    if let Some(socket) = crate::paths::env_path("AGENT_BERTH_TMUX_SOCKET") {
        command.arg("-L").arg(socket);
    }
    if let Some(config) = crate::paths::env_path("AGENT_BERTH_TMUX_CONFIG") {
        command.arg("-f").arg(config);
    }
    command
}