keep_running/

client.rs

use crate::protocol::{decode_message, encode_message, ClientMessage, DaemonMessage};
use crate::session::{self, SessionInfo};
use crate::terminal::{self, status, status_dim, RawModeGuard};
use anyhow::{Context, Result};
use crossterm::tty::IsTty;
use crossterm::{
    cursor::{Hide, MoveTo, Show},
    execute,
    terminal::{Clear, ClearType},
};
use std::io::{self, Read, Write};
use std::os::fd::AsRawFd;
use std::os::unix::net::UnixStream;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

/// Detach sequence: Ctrl+a followed by 'd'
const CTRL_A: u8 = 0x01;

/// Drop raw mode (restoring ONLCR etc.) and emit a clean newline so the
/// next status line lands at column 0. Centralises the previous
/// `drop(raw_guard); eprint!("\r\n");` repetition — the explicit `\r` was
/// redundant once cooked mode was restored.
fn exit_raw(guard: RawModeGuard) {
    drop(guard);
    eprintln!();
}

/// Connect to a session and run the interactive client
pub fn attach(session: &SessionInfo) -> Result<()> {
    // Check if we have a terminal
    if !io::stdin().is_tty() {
        anyhow::bail!("stdin is not a terminal - cannot attach to session");
    }

    let socket_path = session::socket_path(&session.name)?;
    let mut stream =
        UnixStream::connect(&socket_path).context("Failed to connect to session daemon")?;

    // Set read timeout for handshake
    stream.set_read_timeout(Some(Duration::from_secs(5)))?;
    stream.set_nonblocking(false)?;

    // Get terminal size
    let (cols, rows) = terminal::get_size()?;

    // Send attach message
    let msg = ClientMessage::Attach { cols, rows };
    let encoded = encode_message(&msg)?;
    stream.write_all(&encoded)?;

    // Wait for attached confirmation (blocking with timeout)
    let mut buf = [0u8; 8192];
    let n = stream
        .read(&mut buf)
        .context("Failed to read attach confirmation")?;
    if n == 0 {
        anyhow::bail!("Connection closed while waiting for attach confirmation");
    }
    let mut msg_buf = buf[..n].to_vec();

    loop {
        if let Some((msg, consumed)) = decode_message::<DaemonMessage>(&msg_buf)? {
            msg_buf.drain(0..consumed);
            match msg {
                DaemonMessage::Attached => break,
                DaemonMessage::Error(e) => anyhow::bail!("Daemon error: {}", e),
                _ => {}
            }
        } else {
            let n = stream
                .read(&mut buf)
                .context("Failed to read from daemon")?;
            if n == 0 {
                anyhow::bail!("Connection closed while waiting for attach confirmation");
            }
            msg_buf.extend_from_slice(&buf[..n]);
        }
    }

    // Clear timeout and switch to non-blocking for main loop
    stream.set_read_timeout(None)?;
    stream.set_nonblocking(true)?;

    // Clear screen and move cursor to top-left before showing session content
    execute!(io::stdout(), Clear(ClearType::All), MoveTo(0, 0))?;

    // Welcome banner (cooked mode — survives at the top of the cleared screen
    // until program output scrolls over it).
    status(&format!(
        "attached to '{}' · pid {}",
        session.name, session.pid
    ));
    status_dim("detach with Ctrl+a d  ·  kill with Ctrl+a k");

    // Enter raw mode
    let raw_guard = RawModeGuard::enter()?;

    // Set up signal handling for SIGWINCH (terminal resize)
    let resize_flag = Arc::new(AtomicBool::new(false));
    let resize_flag_clone = resize_flag.clone();

    // Install SIGWINCH handler
    unsafe {
        signal_hook::low_level::register(signal_hook::consts::SIGWINCH, move || {
            resize_flag_clone.store(true, Ordering::SeqCst);
        })?;
    }

    // Main I/O loop
    let mut input_buf = [0u8; 1024];
    let mut daemon_buf = [0u8; 8192];
    let mut daemon_msg_buf = msg_buf; // May have leftover data

    // State for detecting Ctrl+a d
    let mut saw_ctrl_a = false;

    // Holds a protocol error message captured during decoding so we can
    // surface it cleanly after dropping raw mode.
    let mut protocol_error: Option<String> = None;

    // Get file descriptors for polling
    let stdin_fd = 0i32;
    let socket_fd = stream.as_raw_fd();

    loop {
        // Check for resize
        if resize_flag.swap(false, Ordering::SeqCst) {
            if let Ok((cols, rows)) = terminal::get_size() {
                let msg = ClientMessage::Resize { cols, rows };
                if let Ok(encoded) = encode_message(&msg) {
                    let _ = stream.write_all(&encoded);
                }
            }
        }

        // Use poll to check for data on stdin and socket
        let mut poll_fds = [
            libc::pollfd {
                fd: stdin_fd,
                events: libc::POLLIN,
                revents: 0,
            },
            libc::pollfd {
                fd: socket_fd,
                events: libc::POLLIN,
                revents: 0,
            },
        ];

        let poll_result = unsafe { libc::poll(poll_fds.as_mut_ptr(), 2, 10) }; // 10ms timeout

        if poll_result < 0 {
            let err = io::Error::last_os_error();
            if err.kind() != io::ErrorKind::Interrupted {
                return Err(err).context("poll failed");
            }
            continue;
        }

        // Check stdin
        if poll_fds[0].revents & libc::POLLIN != 0 {
            let n = unsafe {
                libc::read(
                    stdin_fd,
                    input_buf.as_mut_ptr() as *mut libc::c_void,
                    input_buf.len(),
                )
            };

            if n == 0 {
                // EOF on stdin
                break;
            } else if n > 0 {
                let data = &input_buf[..n as usize];

                // Check for detach/kill sequence (Ctrl+a then d/k)
                // Batch regular bytes together for efficiency (critical for paste)
                let mut i = 0;
                while i < data.len() {
                    if saw_ctrl_a {
                        saw_ctrl_a = false;
                        let byte = data[i];
                        i += 1;

                        match byte {
                            b'd' | b'D' => {
                                // Detach!
                                let msg = ClientMessage::Detach;
                                if let Ok(encoded) = encode_message(&msg) {
                                    let _ = stream.write_all(&encoded);
                                }
                                // Drop raw mode before printing so the banner reflows cleanly.
                                exit_raw(raw_guard);
                                status(&format!("detached from '{}'", session.name));
                                status_dim(&format!("reattach: keep-running {}", session.name));
                                return Ok(());
                            }
                            b'k' | b'K' => {
                                // Kill session!
                                unsafe {
                                    libc::kill(session.pid as i32, libc::SIGHUP);
                                }
                                exit_raw(raw_guard);
                                status(&format!("killed '{}'", session.name));
                                return Ok(());
                            }
                            CTRL_A => {
                                // Double Ctrl+a - send a literal Ctrl+a
                                let msg = ClientMessage::Input(vec![CTRL_A]);
                                if let Ok(encoded) = encode_message(&msg) {
                                    let _ = stream.write_all(&encoded);
                                }
                            }
                            _ => {
                                // Not a command, send the Ctrl+a we held back, plus this byte
                                let msg = ClientMessage::Input(vec![CTRL_A, byte]);
                                if let Ok(encoded) = encode_message(&msg) {
                                    let _ = stream.write_all(&encoded);
                                }
                            }
                        }
                    } else {
                        // Find the next Ctrl+A or end of data
                        let start = i;
                        while i < data.len() && data[i] != CTRL_A {
                            i += 1;
                        }

                        // Send batch of regular bytes as a single message
                        if i > start {
                            let msg = ClientMessage::Input(data[start..i].to_vec());
                            if let Ok(encoded) = encode_message(&msg) {
                                let _ = stream.write_all(&encoded);
                            }
                        }

                        // If we stopped at Ctrl+A, consume it and set flag
                        if i < data.len() && data[i] == CTRL_A {
                            saw_ctrl_a = true;
                            i += 1;
                        }
                    }
                }
            }
        }

        // Check socket
        let mut socket_eof = false;
        if poll_fds[1].revents & libc::POLLIN != 0 {
            match stream.read(&mut daemon_buf) {
                Ok(0) => {
                    // Daemon disconnected - but process buffered messages first
                    socket_eof = true;
                }
                Ok(n) => {
                    daemon_msg_buf.extend_from_slice(&daemon_buf[..n]);
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {}
                Err(e) => {
                    return Err(e).context("Error reading from daemon");
                }
            }
        } else if poll_fds[1].revents & (libc::POLLHUP | libc::POLLERR) != 0 {
            // Only treat hangup as fatal when there's no data to read (POLLIN not set).
            // On macOS, POLLHUP can be returned alongside POLLIN during normal
            // socket state transitions; we should drain data before exiting.
            socket_eof = true;
        }

        // Process any messages in the buffer
        loop {
            match decode_message::<DaemonMessage>(&daemon_msg_buf) {
                Ok(Some((msg, consumed))) => {
                    daemon_msg_buf.drain(0..consumed);

                    match msg {
                        DaemonMessage::Output(data) => {
                            terminal::write_stdout(&data)?;
                        }
                        DaemonMessage::ReplayStart => {
                            // Hide cursor during replay to reduce visual noise
                            let _ = execute!(io::stdout(), Hide);
                        }
                        DaemonMessage::ReplayEnd => {
                            // Restore cursor after replay
                            let _ = execute!(io::stdout(), Show);
                        }
                        DaemonMessage::ChildExited { code } => {
                            exit_raw(raw_guard);
                            match code {
                                Some(c) => status(&format!("process exited with code {}", c)),
                                None => status("process terminated by signal"),
                            }
                            return Ok(());
                        }
                        DaemonMessage::Error(e) => {
                            exit_raw(raw_guard);
                            status(&format!("daemon error: {}", e));
                            return Ok(());
                        }
                        DaemonMessage::Attached => {
                            // Already handled
                        }
                    }
                }
                Ok(None) => break,
                Err(e) => {
                    protocol_error = Some(e.to_string());
                    break;
                }
            }
        }

        // Now that we've processed all buffered messages, handle EOF/hangup
        if let Some(e) = protocol_error.take() {
            exit_raw(raw_guard);
            status(&format!("protocol error: {}", e));
            return Ok(());
        }
        if socket_eof {
            exit_raw(raw_guard);
            status("session ended");
            break;
        }
    }

    Ok(())
}

/// Start a new session and immediately attach to it
pub fn run_and_attach(name: &str, command: &[String]) -> Result<()> {
    // Start the daemon
    crate::daemon::start_daemon(name.to_string(), command.to_vec())?;

    // Poll for the daemon to register itself + bind its socket. Replaces a
    // fixed 100ms sleep that could race on slower machines.
    let deadline = std::time::Instant::now() + Duration::from_secs(2);
    let session = loop {
        if let Some(s) = session::load_session(name)? {
            if std::path::Path::new(&s.socket_path).exists() {
                break s;
            }
        }
        if std::time::Instant::now() >= deadline {
            anyhow::bail!("session daemon failed to start within 2s");
        }
        std::thread::sleep(Duration::from_millis(20));
    };

    // Don't print a banner here — `attach()` clears the screen before printing
    // its own welcome banner, so anything we print would be wallpapered over.
    attach(&session)
}