maestro/crates/maestro-daemon/src/pty.rs

//! Per-session PTY byte bridge.
//!
//! tmux writes pane output to a named pipe via `pipe-pane 'cat >> <fifo>'`.
//! The daemon reads from that FIFO and broadcasts the bytes to all attached
//! WS clients via a [`tokio::sync::broadcast`] channel.
//!
//! Implementation notes:
//!
//! - The FIFO is opened `O_RDWR | O_NONBLOCK`. `O_RDWR` keeps the file's
//!   writer count non-zero from our side, so `read()` never returns EOF when
//!   tmux's `cat` exits (e.g., during a pipe-pane re-issue or a brief tmux
//!   hiccup). `O_NONBLOCK` is essential: without it, `tokio::fs::File`'s
//!   reads route through `spawn_blocking` and block one of those threads in
//!   `read()` until data arrives — which never happens for idle sessions.
//!   A pile of such threads also prevents the tokio runtime from shutting
//!   down cleanly between tests.
//! - We wrap the non-blocking file descriptor in [`AsyncFd`] and use
//!   `try_io` to read directly off the FD, leaving the blocking pool alone.
//! - NAP-229: every byte we broadcast is also appended to a per-bridge ring
//!   buffer (default 64 KB). New subscribers receive the buffer as a
//!   one-shot prefix before joining the live broadcast, so attaching to an
//!   idle session immediately shows the last screenful instead of a blank
//!   terminal.

use std::collections::VecDeque;
use std::ffi::CString;
use std::fs::File;
use std::io::{self, Read};
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::OpenOptionsExt;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Instant;

use tokio::io::unix::AsyncFd;
use tokio::io::Interest;
use tokio::sync::broadcast;
use tokio::task::JoinHandle;

/// How many recent PTY frames each new subscriber inherits before starting to
/// receive live frames. A small buffer is fine — clients render a fresh
/// terminal anyway. Sourced from [`maestro_config::PTY_CAPACITY`]
/// (NAP-273).
pub const PTY_BROADCAST_CAPACITY: usize = maestro_config::PTY_CAPACITY;

/// Default replay-buffer size — roughly one screenful for a typical 24×80
/// terminal plus reasonable scrollback. Override via
/// [`PtyBridge::start_with_replay`] or the `[daemon] pty_replay_bytes` config.
pub const DEFAULT_REPLAY_BYTES: usize = 64 * 1024;

/// Maximum bytes read per syscall; chosen to match typical pipe buffer sizes
/// and avoid copying huge `Vec`s through the channel.
const READ_BUF_BYTES: usize = 4096;

/// Owns a FIFO + the reader task that pumps bytes into a broadcast channel.
/// Dropping the bridge aborts the reader and closes the channel; the FIFO
/// file is left on disk for `SessionManager::kill` to remove.
#[derive(Debug)]
pub struct PtyBridge {
    fifo_path: PathBuf,
    /// `None` once [`Self::seal`] has run: live subscribers see the channel
    /// close, but the replay buffer stays readable.
    sender: Option<broadcast::Sender<Vec<u8>>>,
    replay: ReplayBuffer,
    /// Monotonic timestamp of the last non-empty read off the FIFO, bumped by
    /// `run_reader`. The activity watcher reads it to gauge PTY quiescence
    /// without touching the byte stream (NAP-340).
    last_activity: Arc<Mutex<Instant>>,
    task: JoinHandle<()>,
}

impl PtyBridge {
    /// Create the FIFO (if it doesn't exist), open it non-blocking, and start
    /// the reader task with the default replay buffer size.
    pub async fn start(fifo_path: PathBuf) -> io::Result<Self> {
        Self::start_with_replay(fifo_path, DEFAULT_REPLAY_BYTES).await
    }

    /// Variant that lets the caller pick a custom replay-buffer cap (bytes).
    /// Pass `0` to disable replay (subscribers see only live bytes).
    pub async fn start_with_replay(fifo_path: PathBuf, replay_bytes: usize) -> io::Result<Self> {
        if let Some(parent) = fifo_path.parent() {
            tokio::fs::create_dir_all(parent).await?;
        }
        mkfifo(&fifo_path)?;

        let file = std::fs::OpenOptions::new()
            .read(true)
            .write(true)
            .custom_flags(libc::O_NONBLOCK)
            .open(&fifo_path)?;
        let async_fd = AsyncFd::with_interest(file, Interest::READABLE)?;

        let (sender, _) = broadcast::channel::<Vec<u8>>(PTY_BROADCAST_CAPACITY);
        let replay = ReplayBuffer::new(replay_bytes);
        let tx = sender.clone();
        let replay_for_task = replay.clone();
        let last_activity = Arc::new(Mutex::new(Instant::now()));
        let last_activity_for_task = last_activity.clone();
        let task = tokio::spawn(async move {
            run_reader(async_fd, tx, replay_for_task, last_activity_for_task).await;
        });
        Ok(Self {
            fifo_path,
            sender: Some(sender),
            replay,
            last_activity,
            task,
        })
    }

    /// How long since the session last produced PTY output. The activity
    /// watcher treats a short value as `Active` (NAP-340).
    pub fn idle_for(&self) -> std::time::Duration {
        self.last_activity
            .lock()
            .expect("last_activity poisoned")
            .elapsed()
    }

    pub fn fifo_path(&self) -> &Path {
        &self.fifo_path
    }

    /// Subscribe to live PTY frames. Use [`Self::subscribe_with_replay`] for
    /// new attaches that want the buffered scrollback first.
    pub fn subscribe(&self) -> broadcast::Receiver<Vec<u8>> {
        match &self.sender {
            Some(tx) => tx.subscribe(),
            None => {
                let (tx, rx) = broadcast::channel(1);
                drop(tx);
                rx
            }
        }
    }

    /// Atomically snapshot the replay buffer and return a fresh subscriber
    /// (NAP-229). Send the returned `Vec<u8>` to the client first; subsequent
    /// live bytes arrive on the receiver. Holding the replay lock across
    /// `subscribe()` keeps the reader task from publishing in between, so
    /// the seam has no gap and no duplicates.
    pub fn subscribe_with_replay(&self) -> (Vec<u8>, broadcast::Receiver<Vec<u8>>) {
        let guard = self.replay.lock_for_snapshot();
        let snapshot = guard.snapshot();
        let rx = match &self.sender {
            Some(tx) => tx.subscribe(),
            // Sealed: hand back the scrollback plus a receiver that is already
            // closed, so the client renders the final output and then sees the
            // stream end rather than hanging.
            None => {
                let (tx, rx) = broadcast::channel(1);
                drop(tx);
                rx
            }
        };
        drop(guard);
        (snapshot, rx)
    }

    /// Stop reading and close the live stream, keeping the replay buffer.
    ///
    /// Attached clients see the channel close — without that a `maestro attach`
    /// sits in raw mode against a shell that has exited (NAP-226) — while the
    /// buffered output survives for anyone who attaches later.
    pub fn seal(&mut self) {
        self.task.abort();
        self.sender = None;
    }
}

impl Drop for PtyBridge {
    fn drop(&mut self) {
        self.task.abort();
    }
}

// ---------------------------------------------------------------------------
// ReplayBuffer
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
struct ReplayBuffer {
    inner: Arc<Mutex<ReplayInner>>,
}

#[derive(Debug)]
struct ReplayInner {
    buf: VecDeque<u8>,
    cap: usize,
}

impl ReplayBuffer {
    fn new(cap: usize) -> Self {
        Self {
            inner: Arc::new(Mutex::new(ReplayInner {
                buf: VecDeque::with_capacity(cap.min(READ_BUF_BYTES * 2)),
                cap,
            })),
        }
    }

    /// Append bytes, dropping oldest as needed to stay under `cap`.
    /// Only used by tests — production calls inline this under
    /// `run_reader`'s lock so the broadcast send stays atomic with the push.
    #[cfg(test)]
    fn push(&self, bytes: &[u8]) {
        let mut s = self.inner.lock().expect("ReplayBuffer poisoned");
        if s.cap == 0 {
            return;
        }
        // If the incoming chunk alone exceeds the cap, take only its tail —
        // anything in the buffer would be evicted before this finishes anyway.
        let take_from = bytes.len().saturating_sub(s.cap);
        let bytes = &bytes[take_from..];
        s.buf.extend(bytes);
        while s.buf.len() > s.cap {
            s.buf.pop_front();
        }
    }

    /// Acquire the lock for a snapshot. The caller must call `snapshot()` on
    /// the returned guard before releasing the lock — that's the
    /// atomicity that prevents a race with concurrent appends.
    fn lock_for_snapshot(&self) -> ReplayGuard<'_> {
        ReplayGuard {
            inner: self.inner.lock().expect("ReplayBuffer poisoned"),
        }
    }
}

struct ReplayGuard<'a> {
    inner: std::sync::MutexGuard<'a, ReplayInner>,
}

impl ReplayGuard<'_> {
    fn snapshot(&self) -> Vec<u8> {
        let (a, b) = self.inner.buf.as_slices();
        let mut out = Vec::with_capacity(a.len() + b.len());
        out.extend_from_slice(a);
        out.extend_from_slice(b);
        out
    }
}

async fn run_reader(
    async_fd: AsyncFd<File>,
    tx: broadcast::Sender<Vec<u8>>,
    replay: ReplayBuffer,
    last_activity: Arc<Mutex<Instant>>,
) {
    let mut buf = [0u8; READ_BUF_BYTES];
    loop {
        let mut guard = match async_fd.readable().await {
            Ok(g) => g,
            Err(e) => {
                tracing::warn!(error = ?e, "fifo readability wait failed");
                return;
            }
        };
        // `&File: Read`, so we can read through the immutable ref the AsyncFd
        // hands us inside try_io.
        let result = guard.try_io(|inner| {
            let mut file: &File = inner.get_ref();
            file.read(&mut buf)
        });
        match result {
            Ok(Ok(0)) => {
                // 0 bytes despite holding a write fd; just yield and retry.
                tokio::task::yield_now().await;
            }
            Ok(Ok(n)) => {
                // Mark the session as having just produced output, for the
                // activity watcher's quiescence check (NAP-340).
                *last_activity.lock().expect("last_activity poisoned") = Instant::now();
                // Hold the replay lock across both the buffer push and the
                // broadcast send so `subscribe_with_replay` either sees this
                // chunk in its snapshot *or* receives it on the broadcast —
                // never both, never neither.
                let mut s = replay.inner.lock().expect("ReplayBuffer poisoned");
                if s.cap > 0 {
                    let take_from = n.saturating_sub(s.cap);
                    s.buf.extend(&buf[take_from..n]);
                    while s.buf.len() > s.cap {
                        s.buf.pop_front();
                    }
                }
                let _ = tx.send(buf[..n].to_vec());
                drop(s);
            }
            Ok(Err(e)) => {
                tracing::warn!(error = ?e, "fifo read error");
                return;
            }
            Err(_would_block) => {
                // AsyncFd reported readable but read() said WouldBlock.
                // Clear readiness so we wait again. try_io did that for us.
                continue;
            }
        }
    }
}

fn mkfifo(path: &Path) -> io::Result<()> {
    let c =
        CString::new(path.as_os_str().as_bytes()).map_err(|e| io::Error::other(e.to_string()))?;
    // SAFETY: `c` is a valid NUL-terminated path string for the duration of
    // the call. mkfifo only reads it.
    let res = unsafe { libc::mkfifo(c.as_ptr(), 0o600) };
    if res == 0 {
        return Ok(());
    }
    let err = io::Error::last_os_error();
    if err.kind() == io::ErrorKind::AlreadyExists {
        return Ok(());
    }
    Err(err)
}

/// The shell command tmux's `pipe-pane` will run to forward pane output into
/// `fifo_path`. `cat` writes raw bytes through unchanged.
pub fn pipe_pane_command(fifo_path: &Path) -> String {
    format!("cat >> {}", shell_quote(&fifo_path.display().to_string()))
}

fn shell_quote(s: &str) -> String {
    let escaped = s.replace('\'', r#"'\''"#);
    format!("'{escaped}'")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn replay_buffer_keeps_last_cap_bytes() {
        let r = ReplayBuffer::new(4);
        r.push(b"abcdefgh");
        assert_eq!(r.lock_for_snapshot().snapshot(), b"efgh");
    }

    #[test]
    fn replay_buffer_disabled_when_cap_zero() {
        let r = ReplayBuffer::new(0);
        r.push(b"anything");
        assert!(r.lock_for_snapshot().snapshot().is_empty());
    }

    #[test]
    fn replay_buffer_handles_chunk_larger_than_cap() {
        let r = ReplayBuffer::new(3);
        r.push(b"abcdefghij");
        assert_eq!(r.lock_for_snapshot().snapshot(), b"hij");
    }

    #[test]
    fn replay_buffer_accumulates_across_pushes() {
        let r = ReplayBuffer::new(6);
        r.push(b"abc");
        r.push(b"de");
        assert_eq!(r.lock_for_snapshot().snapshot(), b"abcde");
        r.push(b"fghi");
        assert_eq!(r.lock_for_snapshot().snapshot(), b"defghi");
    }
}