vibe-tests 0.0.1

Integration test framework for MCP servers with LLM-powered tool calling.
Documentation
//! Tee writer that duplicates output to file and stdout.
//! Used for logging child process output to both terminal and log file.

use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;
use tracing::Level;

use crate::base::alias::OnLog;
use crate::env::env_log::EnvLog;

/// Writer that duplicates output to file and stdout.
/// Filters terminal output by log level.
pub struct TeeWriter {
    file: Arc<Mutex<std::fs::File>>,
    path: PathBuf,
    on_log: Option<Arc<OnLog>>,
    log_level: Level,
}

impl TeeWriter {
    /// Creates a new TeeWriter that writes to file and forwards to on_log callback.
    pub fn new(path: PathBuf, log_level: Level, on_log: Option<OnLog>) -> std::io::Result<Self> {
        let file = std::fs::File::create(&path)?;
        Ok(Self {
            file: Arc::new(Mutex::new(file)),
            path,
            on_log: on_log.map(Arc::new),
            log_level,
        })
    }

    /// Returns the path to the log file.
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Creates and emits an EnvLog event if the line passes the log level filter.
    fn emit(&self, on_log: &OnLog, line: &str) {
        if self.should_print(line) {
            on_log(EnvLog {
                timestamp: chrono::Utc::now().to_rfc3339(),
                message: line.trim_end_matches('\n').to_string(),
            });
        }
    }

    /// Checks if a line should be printed to terminal based on log level.
    fn should_print(&self, line: &str) -> bool {
        let level = if line.contains("ERROR") && line.contains("[0m") {
            Level::ERROR
        } else if line.contains("WARN") && line.contains("[0m") {
            Level::WARN
        } else if line.contains("INFO") && line.contains("[0m") {
            Level::INFO
        } else if line.contains("DEBUG") && line.contains("[0m") {
            Level::DEBUG
        } else if line.contains("TRACE") && line.contains("[0m") {
            Level::TRACE
        } else {
            // Raw output without ANSI colors (e.g., docker, child processes)
            // Only show at TRACE level
            Level::TRACE
        };

        level <= self.log_level
    }
}

impl Write for TeeWriter {
    /// Writes data to both the log file and stdout (filtered by log level).
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        // Always write to file
        if let Ok(mut file) = self.file.lock() {
            let _ = file.write_all(buf);
            let _ = file.flush();
        }

        // Forward line by line to on_log callback
        if let Some(ref on_log) = self.on_log {
            let mut line_buffer = Vec::new();
            for &byte in buf {
                line_buffer.push(byte);
                if byte == b'\n' {
                    let line = String::from_utf8_lossy(&line_buffer);
                    self.emit(on_log, &line);
                    line_buffer.clear();
                }
            }
            if !line_buffer.is_empty() {
                let line = String::from_utf8_lossy(&line_buffer);
                self.emit(on_log, &line);
            }
        }

        Ok(buf.len())
    }

    /// Flushes the log file.
    fn flush(&mut self) -> std::io::Result<()> {
        if let Ok(mut file) = self.file.lock() {
            file.flush()
        } else {
            Ok(())
        }
    }
}

impl Clone for TeeWriter {
    /// Clones the TeeWriter, sharing the same file handle, path, and on_log callback.
    fn clone(&self) -> Self {
        Self {
            file: Arc::clone(&self.file),
            path: self.path.clone(),
            on_log: self.on_log.clone(),
            log_level: self.log_level,
        }
    }
}

impl From<TeeWriter> for Stdio {
    /// Converts a TeeWriter into Stdio by creating a pipe and spawning a background thread
    /// that copies data from the pipe to the TeeWriter.
    fn from(writer: TeeWriter) -> Self {
        let (mut reader, writer_pipe) =
            os_pipe::pipe().expect("Failed to create pipe for TeeWriter");

        thread::spawn(move || {
            let mut tee = writer;
            // Copy all data from pipe to TeeWriter
            // Ignore errors to avoid panicking in background thread
            let _ = std::io::copy(&mut reader, &mut tee);
        });

        Stdio::from(writer_pipe)
    }
}