bevy_brp_mcp 0.22.2

MCP server for Bevy Remote Protocol (BRP) integration
//! Optimized watch logging with buffering and batching

use std::fmt::Write;
use std::path::PathBuf;
use std::time::UNIX_EPOCH;

use chrono::DateTime;
use chrono::Local;
use serde_json::Value;
use tokio::fs::File;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::Instant;
use tracing::debug;
use tracing::error;

use super::constants::BUFFER_FLUSH_SIZE;
use super::constants::WATCH_LOG_BUFFER_CAPACITY;
use super::constants::WATCH_LOG_BUFFER_SIZE;
use super::constants::WATCH_LOG_FLUSH_INTERVAL;
use crate::log_tools::TracingLevel;

/// Log entry to be written
#[derive(Debug)]
pub(super) struct LogEntry {
    pub(super) update_type: String,
    pub(super) data:        Value,
    pub(super) timestamp:   DateTime<Local>,
}

/// Buffered logger for watch updates
pub(super) struct BufferedWatchLogger {
    tx:          mpsc::Sender<LogEntry>,
    shutdown_tx: Option<oneshot::Sender<()>>,
}

impl BufferedWatchLogger {
    /// Create a new buffered logger and spawn the writer task
    pub(super) fn new(log_path: PathBuf) -> Self {
        let (tx, rx) = mpsc::channel(WATCH_LOG_BUFFER_SIZE);
        let (shutdown_tx, shutdown_rx) = oneshot::channel();

        // Spawn the writer task
        tokio::spawn(async move {
            if let Err(e) = write_task(log_path, rx, shutdown_rx).await {
                error!("Watch logger write task failed: {e}");
            }
        });

        Self {
            tx,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// Queue a log entry for writing (non-blocking)
    pub(super) async fn write_update(&self, update_type: &str, data: Value) -> Result<(), String> {
        let entry = LogEntry {
            update_type: update_type.to_string(),
            data,
            timestamp: chrono::Local::now(),
        };

        self.tx
            .send(entry)
            .await
            .map_err(|_| "Logger channel closed".to_string())
    }

    /// Queue a debug log entry for writing only if debug mode is enabled
    pub(super) async fn write_debug_update(
        &self,
        update_type: &str,
        data: Value,
    ) -> Result<(), String> {
        if matches!(
            TracingLevel::get_current_tracing_level(),
            TracingLevel::Debug | TracingLevel::Trace
        ) {
            self.write_update(update_type, data).await
        } else {
            Ok(())
        }
    }

    /// Get the log file path for a watch (same as before)
    pub(super) fn get_watch_log_path(watch_id: u32, entity_id: u64, watch_type: &str) -> PathBuf {
        let timestamp = std::time::SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_or(0, |d| d.as_secs());

        let filename =
            format!("bevy_brp_mcp_watch_{watch_id}_{watch_type}_{entity_id}_{timestamp}.log");

        std::env::temp_dir().join(filename)
    }
}

impl Drop for BufferedWatchLogger {
    fn drop(&mut self) {
        // Signal shutdown to the write task
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            let _ = shutdown_tx.send(());
            debug!("Sent shutdown signal to watch logger");
        }
    }
}

/// Helper function to flush the buffer to the file
async fn flush_buffer(
    file: &mut File,
    buffer: &mut String,
    last_flush: &mut Instant,
) -> std::io::Result<()> {
    if !buffer.is_empty() {
        file.write_all(buffer.as_bytes()).await?;
        file.flush().await?;
        buffer.clear();
        *last_flush = tokio::time::Instant::now();
        debug!("Flushed watch log buffer");
    }
    Ok(())
}

/// Background task that batches and writes log entries
async fn write_task(
    log_path: PathBuf,
    mut rx: mpsc::Receiver<LogEntry>,
    mut shutdown_rx: oneshot::Receiver<()>,
) -> std::io::Result<()> {
    // Open file once and keep it open
    let mut file = OpenOptions::new()
        .create(true)
        .append(true)
        .open(&log_path)
        .await?;

    // Buffer for batching writes
    let mut buffer = String::with_capacity(WATCH_LOG_BUFFER_CAPACITY);
    let mut last_flush = tokio::time::Instant::now();
    let flush_interval = WATCH_LOG_FLUSH_INTERVAL;

    loop {
        // Try to receive with timeout, but also check for shutdown signal
        tokio::select! {
            // Check for shutdown signal
            _ = &mut shutdown_rx => {
                debug!("Watch logger received shutdown signal");
                break;
            }

            // Try to receive log entry with timeout
            timeout_result = tokio::time::timeout(flush_interval, rx.recv()) => {
                match timeout_result {
                    Ok(Some(entry)) => {
                        // Format entry into buffer
                        let timestamp = entry.timestamp.format("%Y-%m-%d %H:%M:%S%.3f");
                        if let Ok(json) = serde_json::to_string(&entry.data) {
                            let _ = writeln!(
                                &mut buffer,
                                "[{}] {}: {}",
                                timestamp, entry.update_type, json
                            );
                        }

                        // Check if we should flush (buffer size or time)
                        if buffer.len() > BUFFER_FLUSH_SIZE
                            || last_flush.elapsed() > flush_interval
                        {
                            flush_buffer(&mut file, &mut buffer, &mut last_flush).await?;
                        }
                    }
                    Ok(None) => {
                        // Channel closed, flush remaining buffer and exit
                        debug!("Watch logger message channel closed");
                        break;
                    }
                    Err(_) => {
                        // Timeout - flush if buffer has content
                        flush_buffer(&mut file, &mut buffer, &mut last_flush).await?;
                    }
                }
            }
        }
    }

    // Final flush before shutdown
    flush_buffer(&mut file, &mut buffer, &mut last_flush).await?;
    debug!("Watch logger write task shutting down cleanly");

    Ok(())
}