wraith/src-tauri/src/mcp/scrollback.rs

//! Per-session scrollback buffer for MCP terminal_read.
//!
//! A thread-safe circular buffer that stores the last N bytes of terminal
//! output. Both SSH and PTY output loops write to it. The MCP tools read
//! from it without touching xterm.js or the frontend.

use std::sync::Mutex;

const DEFAULT_CAPACITY: usize = 64 * 1024; // 64KB per session

/// Thread-safe circular buffer for terminal output.
pub struct ScrollbackBuffer {
    inner: Mutex<RingBuffer>,
}

struct RingBuffer {
    data: Vec<u8>,
    capacity: usize,
    /// Write position (wraps around)
    write_pos: usize,
    /// Total bytes written (for detecting wrap)
    total_written: usize,
}

impl ScrollbackBuffer {
    pub fn new() -> Self {
        Self::with_capacity(DEFAULT_CAPACITY)
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            inner: Mutex::new(RingBuffer {
                data: vec![0u8; capacity],
                capacity,
                write_pos: 0,
                total_written: 0,
            }),
        }
    }

    /// Append bytes to the buffer. Old data is overwritten when full.
    pub fn push(&self, bytes: &[u8]) {
        if bytes.is_empty() {
            return;
        }
        let mut buf = self.inner.lock().unwrap();
        let cap = buf.capacity;
        // If input exceeds capacity, only keep the last `cap` bytes
        let data = if bytes.len() > cap {
            &bytes[bytes.len() - cap..]
        } else {
            bytes
        };
        let write_pos = buf.write_pos;
        let first_len = (cap - write_pos).min(data.len());
        buf.data[write_pos..write_pos + first_len].copy_from_slice(&data[..first_len]);
        if first_len < data.len() {
            buf.data[..data.len() - first_len].copy_from_slice(&data[first_len..]);
        }
        buf.write_pos = (write_pos + data.len()) % cap;
        buf.total_written += bytes.len();
    }

    /// Read the last `n` lines from the buffer, with ANSI escape codes stripped.
    pub fn read_lines(&self, n: usize) -> String {
        let raw = self.read_raw();
        let text = strip_ansi(&raw);
        let lines: Vec<&str> = text.lines().collect();
        let start = lines.len().saturating_sub(n);
        lines[start..].join("\n")
    }

    /// Read all buffered content as raw bytes (ordered oldest→newest).
    pub fn read_raw(&self) -> String {
        let buf = self.inner.lock().unwrap();
        let bytes = if buf.total_written >= buf.capacity {
            // Buffer has wrapped — read from write_pos to end, then start to write_pos
            let mut out = Vec::with_capacity(buf.capacity);
            out.extend_from_slice(&buf.data[buf.write_pos..]);
            out.extend_from_slice(&buf.data[..buf.write_pos]);
            out
        } else {
            // Buffer hasn't wrapped yet
            buf.data[..buf.write_pos].to_vec()
        };
        String::from_utf8_lossy(&bytes).to_string()
    }

    /// Total bytes written since creation.
    pub fn total_written(&self) -> usize {
        self.inner.lock().unwrap().total_written
    }
}

/// Strip ANSI escape sequences from text.
fn strip_ansi(input: &str) -> String {
    let mut output = String::with_capacity(input.len());
    let mut chars = input.chars().peekable();

    while let Some(ch) = chars.next() {
        if ch == '\x1b' {
            // ESC sequence — consume until terminator
            if let Some(&next) = chars.peek() {
                if next == '[' {
                    chars.next(); // consume '['
                    // CSI sequence — consume until letter
                    while let Some(&c) = chars.peek() {
                        chars.next();
                        if c.is_ascii_alphabetic() || c == '~' || c == '@' {
                            break;
                        }
                    }
                } else if next == ']' {
                    chars.next(); // consume ']'
                    // OSC sequence — consume until BEL or ST
                    while let Some(&c) = chars.peek() {
                        chars.next();
                        if c == '\x07' {
                            break;
                        }
                        if c == '\x1b' {
                            if chars.peek() == Some(&'\\') {
                                chars.next();
                            }
                            break;
                        }
                    }
                } else {
                    chars.next(); // consume single-char escape
                }
            }
        } else if ch == '\r' {
            // Skip carriage returns for cleaner output
            continue;
        } else {
            output.push(ch);
        }
    }

    output
}

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

    #[test]
    fn push_and_read_basic() {
        let buf = ScrollbackBuffer::new();
        buf.push(b"hello world\n");
        let lines = buf.read_lines(10);
        assert!(lines.contains("hello world"));
    }

    #[test]
    fn read_lines_limits_output() {
        let buf = ScrollbackBuffer::new();
        buf.push(b"line1\nline2\nline3\nline4\nline5\n");
        let lines = buf.read_lines(2);
        assert!(!lines.contains("line3"));
        assert!(lines.contains("line4"));
        assert!(lines.contains("line5"));
    }

    #[test]
    fn circular_buffer_wraps() {
        let buf = ScrollbackBuffer::with_capacity(16);
        buf.push(b"AAAAAAAAAAAAAAAA"); // fill 16 bytes
        buf.push(b"BBBB"); // overwrite first 4
        let raw = buf.read_raw();
        assert!(raw.starts_with("AAAAAAAAAAAA")); // 12 A's remain
        assert!(raw.ends_with("BBBB"));
    }

    #[test]
    fn strip_ansi_removes_csi() {
        let input = "\x1b[32mgreen\x1b[0m normal";
        assert_eq!(strip_ansi(input), "green normal");
    }

    #[test]
    fn strip_ansi_removes_osc() {
        let input = "\x1b]0;title\x07text";
        assert_eq!(strip_ansi(input), "text");
    }

    #[test]
    fn strip_ansi_preserves_plain_text() {
        let input = "no escapes here\njust text";
        assert_eq!(strip_ansi(input), "no escapes here\njust text");
    }

    #[test]
    fn empty_buffer_returns_empty() {
        let buf = ScrollbackBuffer::new();
        assert_eq!(buf.read_lines(10), "");
        assert_eq!(buf.total_written(), 0);
    }

    #[test]
    fn total_written_tracks_all_bytes() {
        let buf = ScrollbackBuffer::with_capacity(8);
        buf.push(b"12345678"); // 8 bytes
        buf.push(b"ABCD"); // 4 more, wraps
        assert_eq!(buf.total_written(), 12);
    }

    #[test]
    fn push_empty_is_noop() {
        let buf = ScrollbackBuffer::with_capacity(8);
        buf.push(b"hello");
        buf.push(b"");
        assert_eq!(buf.total_written(), 5);
        assert!(buf.read_raw().contains("hello"));
    }

    #[test]
    fn push_larger_than_capacity() {
        let buf = ScrollbackBuffer::with_capacity(4);
        buf.push(b"ABCDEFGH"); // 8 bytes into 4-byte buffer
        let raw = buf.read_raw();
        assert_eq!(raw, "EFGH"); // only last 4 bytes kept
        assert_eq!(buf.total_written(), 8);
    }

    #[test]
    fn push_exact_capacity() {
        let buf = ScrollbackBuffer::with_capacity(8);
        buf.push(b"12345678");
        let raw = buf.read_raw();
        assert_eq!(raw, "12345678");
        assert_eq!(buf.total_written(), 8);
    }

    #[test]
    fn push_wrap_around_boundary() {
        let buf = ScrollbackBuffer::with_capacity(8);
        buf.push(b"123456"); // write_pos = 6
        buf.push(b"ABCD");  // wraps: 2 at end, 2 at start
        let raw = buf.read_raw();
        // Buffer: [C, D, 3, 4, 5, 6, A, B], write_pos=2
        // Read from pos 2: "3456AB" + wrap: no, read from write_pos to end then start
        assert_eq!(raw, "3456ABCD");
    }
}