maestro/crates/maestro-tui/src/session_task.rs

//! `SessionTask` — long-running Tokio task that owns one session's WS to the
//! daemon's `WS /api/attach/:id` and feeds bytes into a `vt100::Parser`.
//!
//! The parser's screen lives in `Arc<Mutex<vt100::Parser>>` so the render loop
//! can read it from any thread. The handle returned by [`SessionTask::spawn`]
//! is cheap to clone and is what the rest of the TUI carries around.
//!
//! Background liveness:
//! - The task stays alive even when no pane currently renders the session.
//!   This is what makes "switch tabs and the other session's output keeps
//!   flowing" work — the screen is up to date the moment focus returns.
//! - On WS drop the task transitions to `Reconnecting` and retries with the
//!   same exponential backoff as the catalog (1s → 30s, ±25% jitter). The
//!   parser is **not** reset on reconnect — the last-known frame stays
//!   visible until new bytes arrive, so the user doesn't see a blank.

use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use futures_util::{SinkExt, StreamExt};
use maestro_protocol::ClientControl;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;

/// Match the browser terminal's retained history without letting long-running
/// sessions grow the client without bound.
pub const SCROLLBACK_ROWS: usize = 5_000;

/// Input the dispatcher (or the resize tracker) pushes to a session task.
#[derive(Debug, Clone)]
pub enum SessionInput {
    /// Raw bytes to send on the attach WS as a binary frame (forwarded
    /// keystrokes, paste content, etc.).
    Bytes(Vec<u8>),
    /// Per-session resize. The daemon resizes its PTY in response.
    Resize { cols: u16, rows: u16 },
}

/// Lifecycle state of the underlying WS connection. The renderer reads this
/// to draw a `reconnecting…` overlay or freeze the screen on `Ended`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum SessionState {
    Connecting = 0,
    Live = 1,
    Reconnecting = 2,
    Ended = 3,
}

/// Cheap-to-clone handle returned by [`SessionTask::spawn`]. Holds:
///
/// - `tx` — `mpsc::Sender<SessionInput>`: the dispatcher pushes keystrokes
///   and resizes here.
/// - `screen` — `Arc<Mutex<vt100::Parser>>`: the renderer reads cells from
///   here.
/// - `state` — `Arc<AtomicU8>`: live status of the WS, queryable without
///   holding the screen lock.
/// - `revision` — `Arc<AtomicU64>`: bumped every time the parser's screen is
///   mutated. The render loop compares it against the last drawn frame's so an
///   idle pane costs a `load` instead of a repaint.
#[derive(Clone)]
pub struct SessionHandle {
    pub tx: mpsc::Sender<SessionInput>,
    pub screen: Arc<Mutex<vt100::Parser>>,
    pub(crate) state: Arc<AtomicU8>,
    pub(crate) revision: Arc<AtomicU64>,
}

impl std::fmt::Debug for SessionHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SessionHandle")
            .field("state", &self.state())
            .finish_non_exhaustive()
    }
}

impl SessionHandle {
    pub fn state(&self) -> SessionState {
        match self.state.load(Ordering::Acquire) {
            0 => SessionState::Connecting,
            1 => SessionState::Live,
            2 => SessionState::Reconnecting,
            _ => SessionState::Ended,
        }
    }

    /// Monotonic counter of screen mutations. Only equality across two reads
    /// is meaningful — the loop uses it to decide whether a frame is stale.
    pub fn revision(&self) -> u64 {
        self.revision.load(Ordering::Acquire)
    }

    /// Stop the underlying task. Used when the pane closes and the user has
    /// no other tabs pointing at this session.
    pub fn end(&self) {
        self.state
            .store(SessionState::Ended as u8, Ordering::Release);
    }
}

#[derive(Debug, thiserror::Error)]
pub enum SessionTaskError {
    #[error("ws connect failed: {0}")]
    ConnectFailed(#[from] tokio_tungstenite::tungstenite::Error),
}

#[derive(Debug)]
pub struct SessionTask;

impl SessionTask {
    /// Spawn the background task. Returns immediately with a handle; the WS
    /// connect happens asynchronously. The handle's `state()` reports
    /// `Connecting` until the WS handshake succeeds.
    ///
    /// `rows`/`cols` are the **initial** size hint — the renderer will follow
    /// up with `SessionInput::Resize` on every frame whose allocation changes.
    pub fn spawn(ws_url: String, rows: u16, cols: u16) -> SessionHandle {
        let (tx, mut rx) = mpsc::channel::<SessionInput>(64);
        let parser = Arc::new(Mutex::new(vt100::Parser::new(rows, cols, SCROLLBACK_ROWS)));
        let state = Arc::new(AtomicU8::new(SessionState::Connecting as u8));
        let revision = Arc::new(AtomicU64::new(0));

        let parser_bg = parser.clone();
        let state_bg = state.clone();
        let revision_bg = revision.clone();

        tokio::spawn(async move {
            let url = ws_url;
            let mut backoff = Duration::from_secs(1);

            loop {
                if state_bg.load(Ordering::Acquire) == SessionState::Ended as u8 {
                    return;
                }
                state_bg.store(SessionState::Connecting as u8, Ordering::Release);

                let (ws, _resp) = match tokio_tungstenite::connect_async(&url).await {
                    Ok(x) => x,
                    Err(e) => {
                        tracing::debug!(error = ?e, "session: connect failed");
                        state_bg.store(SessionState::Reconnecting as u8, Ordering::Release);
                        tokio::time::sleep(jitter(backoff)).await;
                        backoff = std::cmp::min(backoff.saturating_mul(2), Duration::from_secs(30));
                        continue;
                    }
                };
                backoff = Duration::from_secs(1);
                state_bg.store(SessionState::Live as u8, Ordering::Release);

                let (mut sink, mut stream) = ws.split();

                // Pump until the WS closes or the sender side errors.
                loop {
                    tokio::select! {
                        biased;
                        Some(inp) = rx.recv() => {
                            match inp {
                                SessionInput::Bytes(b) => {
                                    if sink.send(Message::Binary(b)).await.is_err() {
                                        break;
                                    }
                                }
                                SessionInput::Resize { cols, rows } => {
                                    parser_bg.lock().expect("vt100 parser mutex").set_size(rows, cols);
                                    revision_bg.fetch_add(1, Ordering::Release);
                                    let frame = serde_json::to_string(&ClientControl::Resize {
                                        cols,
                                        rows,
                                    })
                                    .expect("serialize resize");
                                    if sink.send(Message::Text(frame)).await.is_err() {
                                        break;
                                    }
                                }
                            }
                        }
                        frame = stream.next() => {
                            match frame {
                                Some(Ok(Message::Binary(bytes))) => {
                                    parser_bg.lock().expect("vt100 parser mutex").process(&bytes);
                                    revision_bg.fetch_add(1, Ordering::Release);
                                }
                                Some(Ok(Message::Text(text))) => {
                                    // Ping/Pong control frames travel as text JSON. Reply to Ping.
                                    if let Ok(ClientControl::Ping) =
                                        serde_json::from_str::<ClientControl>(&text)
                                    {
                                        let pong = serde_json::to_string(&ClientControl::Pong)
                                            .expect("serialize pong");
                                        if sink.send(Message::Text(pong)).await.is_err() {
                                            break;
                                        }
                                    }
                                }
                                Some(Ok(Message::Close(_))) | None => break,
                                Some(Err(e)) => {
                                    tracing::debug!(error = ?e, "session: ws error");
                                    break;
                                }
                                _ => {}
                            }
                        }
                    }
                }

                // Disconnected. Keep the last-known parser screen intact;
                // reconnect after a backoff.
                state_bg.store(SessionState::Reconnecting as u8, Ordering::Release);
                tokio::time::sleep(jitter(backoff)).await;
                backoff = std::cmp::min(backoff.saturating_mul(2), Duration::from_secs(30));
            }
        });

        SessionHandle {
            tx,
            screen: parser,
            state,
            revision,
        }
    }
}

pub(crate) fn jitter(base: Duration) -> Duration {
    let ms = base.as_millis() as u64;
    let now_nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.subsec_nanos())
        .unwrap_or(0) as u64;
    let delta = (now_nanos % 500) as i64 - 250;
    let scaled = (ms as i64).saturating_mul(delta).div_euclid(1000);
    let total = ms.saturating_add_signed(scaled);
    Duration::from_millis(total)
}

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

    #[test]
    fn state_constructs_as_connecting() {
        let h = SessionHandle {
            tx: mpsc::channel(1).0,
            screen: Arc::new(Mutex::new(vt100::Parser::new(24, 80, 0))),
            state: Arc::new(AtomicU8::new(SessionState::Connecting as u8)),
            revision: Arc::new(AtomicU64::new(0)),
        };
        assert_eq!(h.state(), SessionState::Connecting);
    }

    #[test]
    fn end_transitions_to_ended() {
        let h = SessionHandle {
            tx: mpsc::channel(1).0,
            screen: Arc::new(Mutex::new(vt100::Parser::new(24, 80, 0))),
            state: Arc::new(AtomicU8::new(SessionState::Live as u8)),
            revision: Arc::new(AtomicU64::new(0)),
        };
        h.end();
        assert_eq!(h.state(), SessionState::Ended);
    }

    #[test]
    fn revision_starts_at_zero_and_tracks_screen_mutations() {
        let revision = Arc::new(AtomicU64::new(0));
        let h = SessionHandle {
            tx: mpsc::channel(1).0,
            screen: Arc::new(Mutex::new(vt100::Parser::new(24, 80, 0))),
            state: Arc::new(AtomicU8::new(SessionState::Live as u8)),
            revision: revision.clone(),
        };
        assert_eq!(h.revision(), 0);

        // Stands in for the WS pump: process bytes, then bump.
        h.screen.lock().unwrap().process(b"hello");
        revision.fetch_add(1, Ordering::Release);
        assert_eq!(h.revision(), 1);
    }
}