wsb 0.255.6379

//! The sink: process-wide output mode, and the writer every byte passes.

use std::io::{self, BufRead, BufReader, Write};
use std::path::PathBuf;
use std::process::ExitStatus;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};
use std::time::Instant;

use super::record::{Line, Outcome, Record, Stream};

/// Whether this process packages its output.
///
/// A flag rather than a handle: every writer in the crate consults it on every
/// write, and it is set once, before anything writes, by the argument scan at
/// startup. There is no way to turn it off again — a run that changed format
/// half way through would be unparseable from either end.
static JSONL: AtomicBool = AtomicBool::new(false);

/// Monotonic line numbers, so a consumer can see that it lost one.
static SEQ: AtomicU64 = AtomicU64::new(0);

/// Identifies a child process within one run.
static SUBPROCESS: AtomicU64 = AtomicU64::new(0);

/// When the run started, for the duration on `CommandFinished`.
static STARTED: OnceLock<Instant> = OnceLock::new();

/// Serialises writes so two threads cannot interleave half a line.
fn writer() -> &'static Mutex<()> {
    static WRITER: OnceLock<Mutex<()>> = OnceLock::new();
    WRITER.get_or_init(|| Mutex::new(()))
}

/// The file this run is also written to, when it is being recorded.
///
/// Every operation this workspace runs leaves a file behind, whether it was
/// started from the workbench or typed at a terminal. The workbench records
/// the children it starts; this is how a command line records ITSELF, in the
/// same format, in the same directory, under the same limit — so "where is the
/// log of that build" has one answer regardless of where the build was asked
/// for.
static TRANSCRIPT: OnceLock<Option<crate::logs::op_log::OpLogWriter>> = OnceLock::new();

/// Set by whoever already records this process's output, so it does not record
/// itself as well.
///
/// The workbench runs `wsb` as a child and writes every line the child emits
/// to a file of its own. Without this, that run would produce two files with
/// the same content and two answers to where it is.
pub const NO_TRANSCRIPT_ENV: &str = "WSB_NO_TRANSCRIPT";

/// What this run IS, decided once when the transcript is attached.
///
/// Kept because two things need it and only one of them is the file: the log
/// writer, and the report this run makes of itself to the workspace's daemon
/// so the workbench can show it while it happens. Recomputing the second from
/// the first would mean deriving a run's identity twice, in two places, from
/// an argument vector that has already been consumed.
static RUN: OnceLock<Option<super::activity::Run>> = OnceLock::new();

/// Where this run's lines go besides the file, and what stands between the
/// start of the run and a decision about that.
///
/// One state under one lock rather than a flag and a buffer, because the
/// handover has to be atomic: a line emitted between "the reporter exists" and
/// "the buffer has been given to it" would be a line in neither.
enum Reporting {
    /// Nothing has said yet whether this run is reported. Lines are kept, up
    /// to a cap, so an announcement can carry the whole beginning of the run.
    Waiting { held: Vec<Line>, dropped: u64 },
    /// A daemon is being told about this run.
    Live(super::activity::Reporter),
    /// Nothing is watching: no daemon, no transcript, or the run is already
    /// over.
    Off,
}

/// How many lines are kept before anything has decided whether this run is
/// reported.
///
/// The decision is made by the binary within milliseconds of starting, so this
/// is reached only by a program that never makes it — `wsb serve`, a test
/// harness, a library embedding. Capping it is what keeps such a program from
/// holding every line it ever says.
const HELD_BEFORE_DECIDING: usize = 2000;

fn reporting() -> &'static Mutex<Reporting> {
    static REPORTING: OnceLock<Mutex<Reporting>> = OnceLock::new();
    REPORTING.get_or_init(|| {
        Mutex::new(Reporting::Waiting {
            held: Vec::new(),
            dropped: 0,
        })
    })
}

/// Report this run to the workspace's daemon, so the workbench shows it while
/// it happens.
///
/// Called once by the binary, at the moment it decides whether this run goes
/// beside the machine's daemon (`attach_to_the_daemon`).
///
/// A run the daemon STARTED is not reported: it is already in the registry,
/// put there by the daemon that started it, and a second announcement would
/// be one operation appearing twice. That is not a check written here — such a
/// run has no transcript either, for the same reason, and this returns early
/// on that.
pub fn report_to_daemon(client: &crate::daemon::Client) {
    let Some(Some(run)) = RUN.get() else {
        return;
    };
    let mut state = match reporting().lock() {
        Ok(state) => state,
        Err(_) => return,
    };
    let Reporting::Waiting { held, dropped } = &mut *state else {
        return;
    };
    let held = std::mem::take(held);
    let dropped = *dropped;
    match super::activity::Reporter::start(client.clone(), run.clone(), held, dropped) {
        Ok(reporter) => *state = Reporting::Live(reporter),
        // A thread that will not start is not a reason to keep buffering
        // lines for a reporter that will never read them.
        Err(_) => *state = Reporting::Off,
    }
}

/// Hand one line to whatever is reporting this run.
fn report_line(line: &Line) {
    let Ok(mut state) = reporting().lock() else {
        return;
    };
    match &mut *state {
        Reporting::Waiting { held, dropped } => {
            held.push(line.clone());
            while held.len() > HELD_BEFORE_DECIDING {
                held.remove(0);
                *dropped += 1;
            }
        }
        Reporting::Live(reporter) => reporter.line(line),
        Reporting::Off => {}
    }
}

/// Tell the reporter how the run ended, and stop.
fn close_report(exit_code: i32) {
    let Ok(mut state) = reporting().lock() else {
        return;
    };
    if let Reporting::Live(reporter) = std::mem::replace(&mut *state, Reporting::Off) {
        // The lock is released before waiting: `finished` waits for the
        // reporter thread to acknowledge, and that thread is not touching this
        // state — but a lock held across a wait is a deadlock waiting for a
        // reason, and there is nothing to gain by holding it.
        drop(state);
        reporter.finished(exit_code);
    }
}

/// Begin recording this run to a file, and answer where it is.
///
/// Called once, from `main`, as soon as the workspace root is known and before
/// anything is written — the path is announced as the first thing a run says,
/// so somebody watching a long build knows where to look before it has said
/// anything else, and somebody reading a finished terminal can still find it.
///
/// Returns None when the run is not being recorded: no workspace, a parent
/// that is already recording it, or a directory that could not be opened. A
/// run that cannot be written down still runs.
pub fn attach_transcript(workspace_root: &std::path::Path, argv: &[String]) -> Option<PathBuf> {
    if std::env::var_os(NO_TRANSCRIPT_ENV).is_some() {
        let _ = TRANSCRIPT.set(None);
        // Already recorded and already registered by whoever started this
        // process. Saying so a second time would put one operation in the
        // activity panel twice.
        let _ = RUN.set(None);
        return None;
    }
    let logs_root = crate::logs::root(workspace_root);
    // The command as the kind, so a listing of the directory reads as a list
    // of what was run: `142233-update-1f2e….jsonl`.
    let kind = argv
        .iter()
        .skip(1)
        .find(|arg| !arg.starts_with('-'))
        .cloned()
        .unwrap_or_else(|| "wsb".to_string());
    // The command as a person would say it. argv[0] is whatever path the
    // binary was invoked through — an absolute build path under a disposable
    // directory, more often than not — and a label is for reading.
    let label = std::iter::once(
        std::path::Path::new(argv.first().map(String::as_str).unwrap_or("wsb"))
            .file_name()
            .map(|name| name.to_string_lossy().into_owned())
            .unwrap_or_else(|| "wsb".to_string()),
    )
    .chain(argv.iter().skip(1).cloned())
    .collect::<Vec<_>>()
    .join(" ");
    // Bound rather than passed inline: the run reports itself to the daemon
    // under the SAME id and the same start time the file is written under, so
    // the row in the activity panel and the file on disk are one run rather
    // than two records of it.
    let id = uuid::Uuid::new_v4();
    let started_at = chrono::Utc::now();
    let writer = crate::logs::op_log::OpLogWriter::create(
        &logs_root,
        id,
        &kind,
        &label,
        argv,
        started_at,
    );
    match writer {
        Ok(writer) => {
            let path = writer.path().to_path_buf();
            let _ = RUN.set(Some(super::activity::Run {
                id,
                kind,
                label,
                argv: argv.to_vec(),
                started_at,
                log_path: path.display().to_string(),
            }));
            let _ = TRANSCRIPT.set(Some(writer));
            Some(path)
        }
        Err(error) => {
            // Reported, not swallowed: a run with no file is a run whose
            // output ends with the terminal, and that is worth knowing at the
            // start rather than discovering afterwards.
            let _ = writeln!(io::stderr(), "[wsb] this run is not being recorded: {error}");
            let _ = TRANSCRIPT.set(None);
            let _ = RUN.set(None);
            None
        }
    }
}

/// Close the transcript with how the run ended.
fn close_transcript(exit_code: i32) {
    close_report(exit_code);
    let Some(Some(writer)) = TRANSCRIPT.get() else {
        return;
    };
    let status = if exit_code == 0 {
        crate::logs::op_log::Status::Completed
    } else {
        crate::logs::op_log::Status::Failed
    };
    let error = (exit_code != 0).then(|| format!("exited with status {exit_code}"));
    writer.finish(status, error);
}

/// Write one line to the transcript, and to whatever is watching this run.
///
/// One function for both because they are one act: what the file records and
/// what the activity panel shows must be the same lines, and two call sites
/// would eventually be two sets of them.
fn record_line(line: &Line) {
    if let Some(Some(writer)) = TRANSCRIPT.get() {
        writer.write_line(line);
    }
    report_line(line);
}

/// Whether output is being packaged as JSONL.
pub fn is_jsonl() -> bool {
    JSONL.load(Ordering::Relaxed)
}

/// Package this process's output from here on.
///
/// Called once, from the argument scan, before any command runs.
pub fn jsonl_mode() {
    JSONL.store(true, Ordering::Relaxed);
    // Colour is for a terminal. In JSONL mode it would arrive as escape
    // sequences inside JSON strings, which is not text a consumer can put on a
    // screen — and the same run would produce different bytes depending on
    // where it was started from.
    colored::control::set_override(false);
}

/// Take `--jsonl` out of an argument vector, reporting whether it was there.
///
/// The flag is removed rather than parsed by the argument parser because wsb
/// forwards whole argument vectors to subcommands that read their own
/// arguments — a flag only the top-level parser knew about would be an
/// unknown argument by the time those saw it. Removing it here makes it
/// available to every command, including the ones that parse nothing.
///
/// One name across every wsb tool: they all take `--jsonl` and all answer in
/// the same envelope, so a consumer that can read one can read any of them.
pub fn take_jsonl_flag(argv: &mut Vec<String>) -> bool {
    let before = argv.len();
    argv.retain(|arg| arg != "--jsonl");
    before != argv.len()
}

/// Announce the run. The first line of a JSONL run; nothing in
/// human mode, where the command line is already on screen.
pub fn started(argv: &[String]) {
    let _ = STARTED.set(Instant::now());
    if !is_jsonl() {
        return;
    }
    emit(Record::CommandStarted {
        argv: argv.to_vec(),
        cwd: std::env::current_dir()
            .map(|p| p.display().to_string())
            .unwrap_or_default(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        pid: std::process::id(),
    });
}

/// Close the run. The last line of a JSONL run; nothing in human
/// mode. Returns the exit code it was given, so a caller can
/// `exit(output::finish(code))`.
pub fn finish(exit_code: i32) -> i32 {
    if is_jsonl() {
        let outcome = if exit_code == 0 {
            Outcome::Ok
        } else {
            Outcome::Failed
        };
        emit(Record::CommandFinished {
            outcome,
            exit_code,
            duration_ms: elapsed_ms(),
        });
    }
    let _ = io::stdout().flush();
    close_transcript(exit_code);
    exit_code
}

fn elapsed_ms() -> u64 {
    STARTED
        .get()
        .map(|start| start.elapsed().as_millis() as u64)
        .unwrap_or(0)
}

/// The id the next child process is reported under.
pub(super) fn next_subprocess_id() -> u64 {
    SUBPROCESS.fetch_add(1, Ordering::Relaxed) + 1
}

/// Emit one record.
///
/// Does nothing in human mode for the structured variants: they exist to say
/// in data what the human rendering already says in prose, and saying it twice
/// on a terminal is noise. `Stdout`/`Stderr` are the exception — they ARE the
/// prose, and [`write_out`]/[`write_err`] route them.
pub fn emit(record: Record) {
    let line = Line {
        seq: SEQ.fetch_add(1, Ordering::Relaxed) + 1,
        ts: crate::logging::rfc3339_now(),
        record,
    };
    // The FILE gets every line, in both modes. A run on a terminal says the
    // same things a packaged run says — it just renders them as prose — and a
    // recording that held only the packaged runs would be missing exactly the
    // runs somebody was watching when they went wrong.
    record_line(&line);
    if !is_jsonl() {
        return;
    }
    // A record that cannot be serialised is a defect in this crate, not
    // operator input: report it on stderr rather than dropping the line
    // silently, and keep going — losing one line must not end the run.
    let encoded = match serde_json::to_string(&line) {
        Ok(encoded) => encoded,
        Err(error) => {
            let _ = writeln!(io::stderr(), "[wsb] could not encode an output record: {error}");
            return;
        }
    };
    let _guard = writer().lock();
    let mut out = io::stdout().lock();
    let _ = writeln!(out, "{encoded}");
    // Flushed per line: a consumer reading this stream is watching a run in
    // progress, and a buffered line is one it cannot act on yet.
    let _ = out.flush();
}

/// Write text the command produced for a person.
pub fn write_out(text: &str) {
    if is_jsonl() {
        for line in split_lines(text) {
            emit(Record::Stdout { text: line });
        }
        return;
    }
    // On a terminal the prose goes to the terminal AND into the file, as the
    // same Stdout records a packaged run would have produced. One format on
    // disk whichever mode the run was in, so what reads a log does not have to
    // ask how the run was invoked.
    if TRANSCRIPT.get().is_some_and(Option::is_some) {
        for line in split_lines(text) {
            record_line(&Line {
                seq: SEQ.fetch_add(1, Ordering::Relaxed) + 1,
                ts: crate::logging::rfc3339_now(),
                record: Record::Stdout { text: line },
            });
        }
    }
    let _guard = writer().lock();
    let mut out = io::stdout().lock();
    let _ = out.write_all(text.as_bytes());
    let _ = out.flush();
}

/// Write text the command produced to report trouble.
pub fn write_err(text: &str) {
    if is_jsonl() {
        for line in split_lines(text) {
            emit(Record::Stderr { text: line });
        }
        return;
    }
    if TRANSCRIPT.get().is_some_and(Option::is_some) {
        for line in split_lines(text) {
            record_line(&Line {
                seq: SEQ.fetch_add(1, Ordering::Relaxed) + 1,
                ts: crate::logging::rfc3339_now(),
                record: Record::Stderr { text: line },
            });
        }
    }
    let _guard = writer().lock();
    let mut err = io::stderr().lock();
    let _ = err.write_all(text.as_bytes());
    let _ = err.flush();
}

/// Record text that reached the terminal by another road.
///
/// A child that draws its own live display is given the terminal itself, and
/// what it wrote was in nobody's log: a run that was refused said why on the
/// screen and the recording of that release held an exit status. The caller
/// passes the bytes through to the terminal and hands the lines that SETTLED
/// here, so the file says what the screen said without the redraws.
pub fn record_err(text: &str) {
    if is_jsonl() {
        for line in split_lines(text) {
            emit(Record::Stderr { text: line });
        }
        return;
    }
    if TRANSCRIPT.get().is_some_and(Option::is_some) {
        for line in split_lines(text) {
            record_line(&Line {
                seq: SEQ.fetch_add(1, Ordering::Relaxed) + 1,
                ts: crate::logging::rfc3339_now(),
                record: Record::Stderr { text: line },
            });
        }
    }
}

/// The lines a live display leaves behind, out of the bytes it draws with.
///
/// A live display redraws one line in place — a carriage return, the line
/// again — and ends it with a newline once it will not change. What a log
/// wants is each line as it was left: the text after the last carriage return,
/// with the terminal's control sequences removed, when the newline arrives.
#[derive(Default)]
pub struct SettledLines {
    pending: Vec<u8>,
}

impl SettledLines {
    /// Take more of the stream; return the lines it completed.
    pub fn feed(&mut self, bytes: &[u8]) -> Vec<String> {
        let mut settled = Vec::new();
        for &byte in bytes {
            if byte == b'\n' {
                if let Some(line) = self.take() {
                    settled.push(line);
                }
            } else {
                self.pending.push(byte);
            }
        }
        settled
    }

    /// The stream ended: what was drawn last and never ended with a newline.
    pub fn finish(mut self) -> Option<String> {
        self.take()
    }

    fn take(&mut self) -> Option<String> {
        let text = String::from_utf8_lossy(&self.pending).into_owned();
        self.pending.clear();
        let last = text.rsplit('\r').next().unwrap_or_default();
        let last = strip_ansi(last);
        let last = last.trim_end();
        (!last.trim().is_empty()).then(|| last.to_string())
    }
}

/// Pass a live display through to the terminal unchanged, recording the lines
/// it settles on.
///
/// Bytes are written as they arrive and flushed: a display that redraws a line
/// in place is a stream of partial lines, and holding one back for its newline
/// would freeze the bar it is drawing.
pub fn pass_live_display(mut stream: impl io::Read, terminal: &mut impl io::Write) {
    let mut lines = SettledLines::default();
    let mut buffer = [0u8; 4096];
    loop {
        let read = match stream.read(&mut buffer) {
            // The end, or a broken pipe: the child's exit status says whether
            // that mattered, and it is read next.
            Ok(0) | Err(_) => break,
            Ok(read) => read,
        };
        let _ = terminal.write_all(&buffer[..read]);
        let _ = terminal.flush();
        for line in lines.feed(&buffer[..read]) {
            record_err(&line);
        }
    }
    if let Some(line) = lines.finish() {
        record_err(&line);
    }
}

/// Text with terminal control sequences removed.
///
/// A record's text is text, not a picture of text. Colour reaches this format
/// three ways — the colour crate (turned off with the mode), escape sequences
/// written by hand, and children that colour their own output — and only the
/// first of those can be told not to. So it is stripped here, once, for all
/// three, rather than left for every consumer to do again.
fn strip_ansi(text: &str) -> String {
    if !text.contains('\x1b') {
        return text.to_string();
    }
    let mut out = String::with_capacity(text.len());
    let mut chars = text.chars().peekable();
    while let Some(c) = chars.next() {
        if c != '\x1b' {
            out.push(c);
            continue;
        }
        match chars.next() {
            // CSI: parameters, then one final byte in @..~ that ends it.
            Some('[') => {
                for c in chars.by_ref() {
                    if ('\x40'..='\x7e').contains(&c) {
                        break;
                    }
                }
            }
            // OSC: a string, ended by BEL or by ESC \.
            Some(']') => {
                while let Some(c) = chars.next() {
                    if c == '\x07' {
                        break;
                    }
                    if c == '\x1b' && chars.peek() == Some(&'\\') {
                        chars.next();
                        break;
                    }
                }
            }
            // Anything else is a two-character sequence, already consumed.
            Some(_) | None => {}
        }
    }
    out
}

/// One record per line of text, with the trailing newline removed.
///
/// A write of `"a\nb\n"` is two lines, not one string containing a newline: a
/// consumer that has to split strings itself is parsing again, which is what
/// this format exists to end. A write with no newline at all is still one
/// line — the caller said something, and holding it back until a later write
/// happened to end it would reorder the run.
fn split_lines(text: &str) -> Vec<String> {
    if text.is_empty() {
        return Vec::new();
    }
    let trimmed = text.strip_suffix('\n').unwrap_or(text);
    if trimmed.is_empty() {
        // A lone newline: a deliberate blank line.
        return vec![String::new()];
    }
    trimmed.split('\n').map(strip_ansi).collect()
}

/// Report a finished child, naming a signal death as one.
pub(super) fn finished_child(id: u64, status: &ExitStatus, started: Instant) {
    emit(Record::SubprocessFinished {
        id,
        exit_code: status.code(),
        signal: exit_signal(status),
        duration_ms: started.elapsed().as_millis() as u64,
    });
}

/// Whether this run is writing a log a child's output can be recorded into.
pub(super) fn transcribing() -> bool {
    TRANSCRIPT.get().is_some_and(Option::is_some)
}

/// Read one of a child's pipes to its end, one record per line — and, with
/// `echo`, write each line to the terminal as it arrives, as the child would
/// have, so recording it takes nothing away from the person watching.
pub(super) fn pump<R: io::Read>(id: u64, stream: Stream, pipe: R, echo: bool) {
    let mut reader = BufReader::new(pipe);
    let mut buffer = Vec::new();
    loop {
        buffer.clear();
        match reader.read_until(b'\n', &mut buffer) {
            Ok(0) => return,
            Ok(_) => {
                if echo {
                    let _guard = writer().lock();
                    let _ = match stream {
                        Stream::Stdout => io::stdout().lock().write_all(&buffer),
                        Stream::Stderr => io::stderr().lock().write_all(&buffer),
                    };
                }
                while matches!(buffer.last(), Some(b'\n') | Some(b'\r')) {
                    buffer.pop();
                }
                // A child is free to write bytes that are not UTF-8; the run
                // does not end over it, and the replacement character says
                // where it happened.
                let text = strip_ansi(&String::from_utf8_lossy(&buffer));
                emit(Record::SubprocessOutput { id, stream, text });
            }
            Err(error) => {
                emit(Record::SubprocessOutput {
                    id,
                    stream,
                    text: format!("[wsb: could not read this process's output: {error}]"),
                });
                return;
            }
        }
    }
}

fn exit_signal(status: &ExitStatus) -> Option<i32> {
    #[cfg(unix)]
    {
        use std::os::unix::process::ExitStatusExt;
        return status.signal();
    }
    #[cfg(not(unix))]
    {
        let _ = status;
        None
    }
}

/// Write for a person, or package it. Use exactly as `println!`.
#[macro_export]
macro_rules! outln {
    () => { $crate::output::write_out("\n") };
    ($($arg:tt)*) => { $crate::output::write_out(&format!("{}\n", format_args!($($arg)*))) };
}

/// Write for a person with no trailing newline. Use exactly as `print!`.
#[macro_export]
macro_rules! out {
    ($($arg:tt)*) => { $crate::output::write_out(&format!("{}", format_args!($($arg)*))) };
}

/// Report trouble. Use exactly as `eprintln!`.
#[macro_export]
macro_rules! errln {
    () => { $crate::output::write_err("\n") };
    ($($arg:tt)*) => { $crate::output::write_err(&format!("{}\n", format_args!($($arg)*))) };
}

/// Report trouble with no trailing newline. Use exactly as `eprint!`.
#[macro_export]
macro_rules! err {
    ($($arg:tt)*) => { $crate::output::write_err(&format!("{}", format_args!($($arg)*))) };
}

#[cfg(test)]
#[path = "sink_tests.rs"]
mod tests;