aimux_server/

attach.rs

use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use tokio::io::{self, AsyncRead, AsyncWrite};
use tokio::sync::Mutex;
use tracing::{debug, info, warn};
use aimux_protocol::{
    read_frame, write_frame, ClientInput, LayoutInfo, PaneId, PaneLayout, Response,
    ScreenCell, ScreenCellAttrs, ScreenColor, ServerPush, WindowBarInfo, WindowId,
};

use crate::screen;
use crate::session::{Session, SessionManager, Window};

// ---------------------------------------------------------------------------
// From conversions: server screen types → protocol wire types
// ---------------------------------------------------------------------------

impl From<&screen::Color> for ScreenColor {
    fn from(c: &screen::Color) -> Self {
        match c {
            screen::Color::Default => ScreenColor::Default,
            screen::Color::Indexed(n) => ScreenColor::Indexed(*n),
            screen::Color::Rgb(r, g, b) => ScreenColor::Rgb(*r, *g, *b),
        }
    }
}

impl From<&screen::CellAttrs> for ScreenCellAttrs {
    fn from(a: &screen::CellAttrs) -> Self {
        ScreenCellAttrs {
            bold: a.bold,
            dim: a.dim,
            italic: a.italic,
            underline: a.underline,
            fg: ScreenColor::from(&a.fg),
            bg: ScreenColor::from(&a.bg),
        }
    }
}

impl From<&screen::Cell> for ScreenCell {
    fn from(c: &screen::Cell) -> Self {
        ScreenCell {
            ch: c.ch,
            attrs: ScreenCellAttrs::from(&c.attrs),
        }
    }
}

// ---------------------------------------------------------------------------
// Layout helpers
// ---------------------------------------------------------------------------

fn build_layout_info(
    session: &Session,
    window: &Window,
    terminal_cols: u16,
    terminal_rows: u16,
) -> LayoutInfo {
    let positions = window.compute_pane_positions(terminal_cols, terminal_rows);
    let active_pane_id = window
        .panes
        .get(window.active_pane)
        .map(|p| p.id.clone())
        .unwrap_or_default();

    let panes: Vec<PaneLayout> = positions
        .into_iter()
        .map(|(pane_id, pos)| PaneLayout {
            is_active: pane_id == active_pane_id,
            is_zoomed: false, // Phase 5 adds zoom
            pane_id,
            x: pos.x,
            y: pos.y,
            width: pos.width,
            height: pos.height,
        })
        .collect();

    let windows: Vec<WindowBarInfo> = session
        .windows
        .iter()
        .enumerate()
        .map(|(idx, w)| WindowBarInfo {
            id: w.id,
            title: format!("{}:{}", w.id, w.panes.first().map(|p| p.title.as_str()).unwrap_or("")),
            is_active: idx == session.active_window,
        })
        .collect();

    LayoutInfo {
        session: session.name.clone(),
        window_id: window.id,
        panes,
        terminal_cols,
        terminal_rows,
        windows,
    }
}

/// Collect the current pane IDs and active window ID for mutation detection.
fn snapshot_pane_set(_session: &Session, window: &Window) -> (WindowId, Vec<PaneId>) {
    let pane_ids: Vec<PaneId> = window.panes.iter().map(|p| p.id.clone()).collect();
    (window.id, pane_ids)
}

// ---------------------------------------------------------------------------
// Snapshot serialization
// ---------------------------------------------------------------------------

async fn build_snapshot(
    pane_id: &str,
    screen_arc: &Arc<tokio::sync::RwLock<screen::Screen>>,
) -> ServerPush {
    let screen = screen_arc.read().await;
    let cells: Vec<Vec<ScreenCell>> = screen
        .grid()
        .iter()
        .map(|row| row.iter().map(ScreenCell::from).collect())
        .collect();
    let cursor = screen.cursor();
    let size = screen.size();
    ServerPush::ScreenSnapshot {
        pane_id: pane_id.to_string(),
        cells,
        cursor: aimux_protocol::CursorPos {
            row: cursor.row,
            col: cursor.col,
        },
        size,
        cursor_visible: screen.cursor_visible(),
        title: screen.title().to_string(),
    }
}

// ---------------------------------------------------------------------------
// Main attach handler
// ---------------------------------------------------------------------------

pub async fn handle_attach<S>(
    pipe: S,
    manager: Arc<Mutex<SessionManager>>,
    session_name: &str,
) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    info!("attach: client connecting to session '{}'", session_name);
    let default_cols: u16 = 80;
    let default_rows: u16 = 24;

    // 1. Validate session, build initial layout
    let layout = {
        let mgr = manager.lock().await;
        let session = mgr
            .get_session(session_name)
            .context("session not found")?;
        let window = session
            .windows
            .get(session.active_window)
            .context("no active window")?;
        build_layout_info(session, window, default_cols, default_rows)
    };

    // 2. Send AttachAccepted
    let (mut reader, mut writer) = io::split(pipe);
    let response = Response::AttachAccepted {
        layout: layout.clone(),
    };
    let bytes = rmp_serde::to_vec(&response).context("serialize AttachAccepted")?;
    write_frame(&mut writer, &bytes).await.context("send AttachAccepted")?;

    info!("attach: accepted for session '{}' ({} panes)", session_name, layout.panes.len());

    // 3. Set up coalescing state
    let mut dirty: HashSet<PaneId> = HashSet::new();
    let mut interval = tokio::time::interval(Duration::from_millis(16));
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    // Track last-known pane set for mutation detection
    let (mut last_window_id, mut last_pane_ids) = {
        let mgr = manager.lock().await;
        let session = match mgr.get_session(session_name) {
            Some(s) => s,
            None => {
                warn!("attach: session '{}' disappeared before streaming started", session_name);
                let err = ServerPush::Error {
                    message: "session killed".into(),
                };
                let _ = write_frame(&mut writer, &rmp_serde::to_vec(&err)?).await;
                return Ok(());
            }
        };
        let window = &session.windows[session.active_window];
        snapshot_pane_set(session, window)
    };

    // Subscribe to update channels for all visible panes (tagged with pane ID)
    let mut update_rxs = subscribe_pane_updates(&manager, session_name).await;

    // Mark all panes dirty initially so client gets first full frame
    for pane_id in &last_pane_ids {
        dirty.insert(pane_id.clone());
    }

    let mut terminal_cols = default_cols;
    let mut terminal_rows = default_rows;

    // Spawn a reader task to make frame reading cancel-safe in select!.
    // read_frame does multiple sequential reads (version, length, payload) and
    // is NOT cancel-safe — if dropped mid-read, already-consumed bytes are lost.
    let (frame_tx, mut frame_rx) = tokio::sync::mpsc::channel::<Result<Vec<u8>>>(4);
    tokio::spawn(async move {
        loop {
            let result = read_frame(&mut reader).await;
            let is_err = result.is_err();
            if frame_tx.send(result).await.is_err() {
                break; // Receiver dropped
            }
            if is_err {
                break; // Stop reading after error
            }
        }
    });

    // 4. Enter select! loop
    loop {
        tokio::select! {
            _ = interval.tick() => {
                // Drain update notifications, marking specific panes dirty
                for (pane_id, rx) in &mut update_rxs {
                    if rx.try_recv().is_ok() {
                        dirty.insert(pane_id.clone());
                        // Drain any remaining
                        while rx.try_recv().is_ok() {}
                    }
                }

                // Check for session mutation (pane set changes)
                let mutation_result = {
                    let mgr = manager.lock().await;
                    match mgr.get_session(session_name) {
                        Some(session) => {
                            if session.windows.is_empty() {
                                None
                            } else {
                                let window = &session.windows[session.active_window];
                                let (wid, pids) = snapshot_pane_set(session, window);
                                Some((wid, pids, session.active_window))
                            }
                        }
                        None => None,
                    }
                };

                match mutation_result {
                    None => {
                        // Session killed
                        warn!("attach: session '{}' was killed, disconnecting client", session_name);
                        let msg = ServerPush::Error { message: "session killed".into() };
                        let _ = write_frame(&mut writer, &rmp_serde::to_vec(&msg)?).await;
                        return Ok(());
                    }
                    Some((wid, pids, _active_window_idx)) => {
                        if wid != last_window_id || pids != last_pane_ids {
                            // Pane set changed — rebuild subscriptions, push new layout
                            debug!("attach: pane set changed, rebuilding subscriptions");
                            update_rxs = subscribe_pane_updates(&manager, session_name).await;
                            last_window_id = wid;
                            last_pane_ids = pids.clone();

                            // Mark all new panes dirty
                            for pid in &pids {
                                dirty.insert(pid.clone());
                            }

                            // Push LayoutChanged
                            let layout = {
                                let mgr = manager.lock().await;
                                let session = mgr.get_session(session_name).context("session gone")?;
                                let window = &session.windows[session.active_window];
                                build_layout_info(session, window, terminal_cols, terminal_rows)
                            };
                            let msg = ServerPush::LayoutChanged(layout);
                            let bytes = rmp_serde::to_vec(&msg)?;
                            if write_frame(&mut writer, &bytes).await.is_err() {
                                info!("attach: pipe broken while sending layout change for '{}'", session_name);
                                return Ok(());
                            }
                        }
                    }
                }

                // Send snapshots for dirty panes
                if !dirty.is_empty() {
                    let pane_ids: Vec<PaneId> = dirty.drain().collect();
                    for pane_id in &pane_ids {
                        let screen_arc = {
                            let mgr = manager.lock().await;
                            match mgr.find_pane(pane_id) {
                                Some((_, _, pane)) => pane.screen.clone(),
                                None => continue,
                            }
                        };
                        let snapshot = build_snapshot(pane_id, &screen_arc).await;
                        let bytes = rmp_serde::to_vec(&snapshot)?;
                        if write_frame(&mut writer, &bytes).await.is_err() {
                            info!("attach: pipe broken while sending snapshot for '{}'", session_name);
                            return Ok(());
                        }
                    }
                }
            }

            frame_result = frame_rx.recv() => {
                let payload = match frame_result {
                    Some(Ok(p)) => p,
                    Some(Err(e)) => {
                        info!("attach: client disconnected from '{}': {}", session_name, e);
                        return Ok(());
                    }
                    None => {
                        info!("attach: client reader closed for '{}'", session_name);
                        return Ok(());
                    }
                };
                let input: ClientInput = rmp_serde::from_slice(&payload)
                    .context("deserialize ClientInput")?;

                match input {
                    ClientInput::Keys { data } => {
                        let mgr = manager.lock().await;
                        let session = match mgr.get_session(session_name) {
                            Some(s) => s,
                            None => return Ok(()),
                        };
                        let window = &session.windows[session.active_window];
                        if let Some(pane) = window.panes.get(window.active_pane) {
                            let _ = pane.pty.write(&data);
                        }
                    }

                    ClientInput::Resize { cols, rows } => {
                        debug!("attach: client resize {}x{} for '{}'", cols, rows, session_name);
                        terminal_cols = cols;
                        terminal_rows = rows;

                        // Recompute layout and resize all panes
                        let resize_targets = {
                            let mgr = manager.lock().await;
                            let session = match mgr.get_session(session_name) {
                                Some(s) => s,
                                None => return Ok(()),
                            };
                            let window = &session.windows[session.active_window];
                            let positions = window.compute_pane_positions(cols, rows);
                            positions
                                .into_iter()
                                .map(|(pane_id, pos)| (pane_id, pos.width, pos.height))
                                .collect::<Vec<_>>()
                        };

                        // Resize each pane's PTY and screen
                        for (pane_id, w, h) in &resize_targets {
                            let screen_arc = {
                                let mut mgr = manager.lock().await;
                                match mgr.resize_pane_pty(pane_id, *w, *h) {
                                    Ok(s) => s,
                                    Err(_) => continue,
                                }
                            };
                            let mut screen = screen_arc.write().await;
                            screen.resize(*w, *h);
                            dirty.insert(pane_id.clone());
                        }

                        // Push LayoutChanged
                        let layout = {
                            let mgr = manager.lock().await;
                            let session = match mgr.get_session(session_name) {
                                Some(s) => s,
                                None => return Ok(()),
                            };
                            let window = &session.windows[session.active_window];
                            build_layout_info(session, window, cols, rows)
                        };
                        let msg = ServerPush::LayoutChanged(layout);
                        let bytes = rmp_serde::to_vec(&msg)?;
                        if write_frame(&mut writer, &bytes).await.is_err() {
                            info!("attach: pipe broken during resize for '{}'", session_name);
                            return Ok(());
                        }

                        // Rebuild pane set tracking after resize
                        let mgr = manager.lock().await;
                        if let Some(session) = mgr.get_session(session_name) {
                            let window = &session.windows[session.active_window];
                            let (wid, pids) = snapshot_pane_set(session, window);
                            last_window_id = wid;
                            last_pane_ids = pids;
                        }
                    }

                    ClientInput::Detach => {
                        info!("attach: client sent detach for session '{}'", session_name);
                        return Ok(());
                    }

                    ClientInput::Command(request) => {
                        debug!("attach: command from client in '{}': {:?}", session_name, request);
                        let response = crate::ipc_dispatch::dispatch(request, &manager).await;
                        let msg = rmp_serde::to_vec(&response)?;
                        if write_frame(&mut writer, &msg).await.is_err() {
                            info!("attach: pipe broken sending command response for '{}'", session_name);
                            return Ok(());
                        }
                        // After a command, mark all panes dirty to pick up changes
                        for pane_id in &last_pane_ids {
                            dirty.insert(pane_id.clone());
                        }
                    }
                }
            }
        }
    }
}

/// Subscribe to update_tx for all visible panes in the session's active window.
/// Returns receivers tagged with their pane IDs for per-pane dirty tracking.
async fn subscribe_pane_updates(
    manager: &Arc<Mutex<SessionManager>>,
    session_name: &str,
) -> Vec<(PaneId, tokio::sync::broadcast::Receiver<()>)> {
    let mgr = manager.lock().await;
    let Some(session) = mgr.get_session(session_name) else {
        return vec![];
    };
    let window = &session.windows[session.active_window];
    window
        .panes
        .iter()
        .map(|pane| (pane.id.clone(), pane.update_tx.subscribe()))
        .collect()
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::screen::{Cell, CellAttrs, Color, Screen};

    #[test]
    fn color_default_converts() {
        let c = Color::Default;
        assert_eq!(ScreenColor::from(&c), ScreenColor::Default);
    }

    #[test]
    fn color_indexed_converts() {
        let c = Color::Indexed(42);
        assert_eq!(ScreenColor::from(&c), ScreenColor::Indexed(42));
    }

    #[test]
    fn color_rgb_converts() {
        let c = Color::Rgb(10, 20, 30);
        assert_eq!(ScreenColor::from(&c), ScreenColor::Rgb(10, 20, 30));
    }

    #[test]
    fn cell_attrs_converts() {
        let attrs = CellAttrs {
            bold: true,
            dim: false,
            italic: true,
            underline: false,
            fg: Color::Indexed(1),
            bg: Color::Rgb(255, 0, 128),
        };
        let proto = ScreenCellAttrs::from(&attrs);
        assert!(proto.bold);
        assert!(!proto.dim);
        assert!(proto.italic);
        assert!(!proto.underline);
        assert_eq!(proto.fg, ScreenColor::Indexed(1));
        assert_eq!(proto.bg, ScreenColor::Rgb(255, 0, 128));
    }

    #[test]
    fn cell_converts() {
        let cell = Cell {
            ch: 'X',
            attrs: CellAttrs {
                bold: true,
                ..CellAttrs::default()
            },
        };
        let proto = ScreenCell::from(&cell);
        assert_eq!(proto.ch, 'X');
        assert!(proto.attrs.bold);
    }

    #[test]
    fn default_cell_converts() {
        let cell = Cell::default();
        let proto = ScreenCell::from(&cell);
        assert_eq!(proto.ch, '\0');
        assert_eq!(proto.attrs, ScreenCellAttrs::default());
    }

    #[tokio::test]
    async fn snapshot_from_screen() {
        let mut screen = Screen::new(10, 3, 100);
        screen.feed(b"Hello\r\nWorld");

        let screen_arc = Arc::new(tokio::sync::RwLock::new(screen));
        let snapshot = build_snapshot("%0", &screen_arc).await;

        match snapshot {
            ServerPush::ScreenSnapshot {
                pane_id,
                cells,
                cursor,
                size,
                cursor_visible,
                ..
            } => {
                assert_eq!(pane_id, "%0");
                assert_eq!(size, (10, 3));
                assert!(cursor_visible);
                // Cursor should be after "World" — row 1, col 5
                assert_eq!(cursor.row, 1);
                assert_eq!(cursor.col, 5);
                // Check first row has "Hello"
                assert_eq!(cells.len(), 3);
                let first_row_text: String = cells[0].iter().map(|c| c.ch).collect();
                assert!(first_row_text.starts_with("Hello"));
                let second_row_text: String = cells[1].iter().map(|c| c.ch).collect();
                assert!(second_row_text.starts_with("World"));
            }
            _ => panic!("expected ScreenSnapshot"),
        }
    }

    #[test]
    fn screen_cursor_visible_getter() {
        let screen = Screen::new(80, 24, 100);
        assert!(screen.cursor_visible());
    }

    #[test]
    fn screen_grid_getter() {
        let screen = Screen::new(10, 5, 100);
        let grid = screen.grid();
        assert_eq!(grid.len(), 5);
        assert_eq!(grid[0].len(), 10);
    }

    #[test]
    fn screen_grid_after_feed() {
        let mut screen = Screen::new(10, 3, 100);
        screen.feed(b"AB");
        let grid = screen.grid();
        assert_eq!(grid[0][0].ch, 'A');
        assert_eq!(grid[0][1].ch, 'B');
    }
}