loopflow 0.10.0

//! Output streaming hub with file-backed persistence.
//!
//! Output lines are broadcast in-memory AND appended to per-run log files
//! at `~/.lf/output/<run_id>.log`. The logs endpoint replays from
//! the file, then follows live — no race conditions.

use std::collections::HashMap;
use std::fs::{self, File, OpenOptions};
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::{Duration, SystemTime};

use tokio::sync::broadcast;
use tracing::warn;

use crate::lfd::security::{path_within_root_existing, path_within_root_planned, validate_safe_id};

#[derive(Debug, Clone)]
pub struct OutputEvent {
    pub wave_id: String,
    pub run_id: String,
    pub agent_id: String,
    pub text: String,
}

/// File writer cache — one writer per run_id.
struct Writers {
    dir: PathBuf,
    files: HashMap<String, File>,
}

impl Writers {
    fn new(dir: PathBuf) -> Self {
        fs::create_dir_all(&dir).ok();
        Self {
            dir,
            files: HashMap::new(),
        }
    }

    fn append(&mut self, run_id: &str, line: &str) {
        let Some(relative) = relative_log_path(run_id) else {
            warn!(run_id = %run_id, "rejecting unsafe run_id for output log");
            return;
        };

        let file = if let Some(file) = self.files.get_mut(run_id) {
            file
        } else {
            let path = match path_within_root_planned(&self.dir, &relative) {
                Ok(path) => path,
                Err(err) => {
                    warn!(run_id = %run_id, error = %err, "failed to resolve output log path");
                    return;
                }
            };
            let file = match OpenOptions::new().create(true).append(true).open(path) {
                Ok(file) => file,
                Err(err) => {
                    warn!(run_id = %run_id, error = %err, "failed to open output log file");
                    return;
                }
            };
            self.files.entry(run_id.to_string()).or_insert(file)
        };
        let _ = writeln!(file, "{line}");
    }

    fn close(&mut self, run_id: &str) {
        self.files.remove(run_id);
    }
}

/// Running token totals for one run, summed from `StreamEvent::Usage` events.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct UsageTotals {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
}

impl UsageTotals {
    pub fn is_empty(&self) -> bool {
        self.input_tokens == 0 && self.output_tokens == 0 && self.cache_read_tokens == 0
    }
}

#[derive(Clone)]
pub struct OutputHub {
    sender: broadcast::Sender<OutputEvent>,
    writers: std::sync::Arc<Mutex<Writers>>,
    usage: std::sync::Arc<Mutex<HashMap<String, UsageTotals>>>,
    output_dir: PathBuf,
}

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

impl OutputHub {
    pub fn new(buffer: usize, output_dir: PathBuf) -> Self {
        let (sender, _) = broadcast::channel(buffer);
        let writers = std::sync::Arc::new(Mutex::new(Writers::new(output_dir.clone())));
        Self {
            sender,
            writers,
            usage: std::sync::Arc::new(Mutex::new(HashMap::new())),
            output_dir,
        }
    }

    /// Add token counts to a run's running usage total. Called as
    /// `StreamEvent::Usage` events are parsed from the agent's stream.
    pub fn add_usage(
        &self,
        run_id: &str,
        input_tokens: Option<u64>,
        output_tokens: Option<u64>,
        cache_read_tokens: Option<u64>,
    ) {
        if let Ok(mut usage) = self.usage.lock() {
            let entry = usage.entry(run_id.to_string()).or_default();
            entry.input_tokens += input_tokens.unwrap_or(0);
            entry.output_tokens += output_tokens.unwrap_or(0);
            entry.cache_read_tokens += cache_read_tokens.unwrap_or(0);
        }
    }

    /// Remove and return a run's accumulated usage. Returns `None` when no
    /// usage was recorded (e.g. tmux-backed runs that never stream in-process).
    pub fn take_usage(&self, run_id: &str) -> Option<UsageTotals> {
        self.usage
            .lock()
            .ok()
            .and_then(|mut usage| usage.remove(run_id))
    }

    /// Append to log file, then broadcast. File write happens first so
    /// that replay-then-follow has a clean dedup boundary.
    pub fn send(&self, event: OutputEvent) {
        if let Ok(mut w) = self.writers.lock() {
            w.append(&event.run_id, &event.text);
        }
        let _ = self.sender.send(event);
    }

    pub fn subscribe(&self) -> broadcast::Receiver<OutputEvent> {
        self.sender.subscribe()
    }

    /// Read the log file for a wave run. Returns lines and the byte offset
    /// at end of read (for dedup with the broadcast stream).
    pub fn read_log(&self, run_id: &str) -> Option<(Vec<String>, u64)> {
        let relative = relative_log_path(run_id)?;
        let path = path_within_root_existing(&self.output_dir, &relative).ok()?;
        let file = File::open(&path).ok()?;
        let metadata = file.metadata().ok()?;
        let size = metadata.len();
        let reader = BufReader::new(file);
        let lines: Vec<String> = reader.lines().map_while(Result::ok).collect();
        Some((lines, size))
    }

    /// Drop the file handle for a completed/failed wave run.
    pub fn close_writer(&self, run_id: &str) {
        if let Ok(mut w) = self.writers.lock() {
            w.close(run_id);
        }
    }

    /// Output directory path.
    pub fn output_dir(&self) -> &Path {
        &self.output_dir
    }
}

/// Delete output log files older than `max_age` based on filesystem mtime.
pub fn prune_output_logs(dir: &Path, max_age: Duration) {
    let entries = match fs::read_dir(dir) {
        Ok(entries) => entries,
        Err(err) => {
            tracing::debug!(dir = %dir.display(), error = %err, "skipping output log pruning");
            return;
        }
    };
    let now = SystemTime::now();
    let mut pruned = 0u32;
    for entry in entries.flatten() {
        let path = entry.path();
        if path.extension().and_then(|e| e.to_str()) != Some("log") {
            continue;
        }
        let mtime = match entry.metadata().and_then(|m| m.modified()) {
            Ok(t) => t,
            Err(_) => continue,
        };
        if now.duration_since(mtime).unwrap_or(Duration::ZERO) > max_age
            && fs::remove_file(&path).is_ok()
        {
            pruned += 1;
        }
    }
    if pruned > 0 {
        tracing::info!(count = pruned, dir = %dir.display(), "pruned old output logs");
    }
}

fn relative_log_path(run_id: &str) -> Option<PathBuf> {
    validate_safe_id(run_id).ok()?;
    Some(PathBuf::from(format!("{run_id}.log")))
}

#[cfg(test)]
mod tests {
    use super::{OutputEvent, OutputHub};
    use tempfile::tempdir;

    #[test]
    fn read_log_returns_none_for_unsafe_run_id() {
        let tmp = tempdir().expect("tempdir");
        let hub = OutputHub::new(8, tmp.path().to_path_buf());
        assert!(hub.read_log("../escape").is_none());
    }

    #[test]
    fn send_ignores_unsafe_run_id_without_creating_file() {
        let tmp = tempdir().expect("tempdir");
        let hub = OutputHub::new(8, tmp.path().to_path_buf());
        let outside = tmp.path().parent().expect("parent").join("escape.log");

        hub.send(OutputEvent {
            wave_id: "wave-1".to_string(),
            run_id: "../escape".to_string(),
            agent_id: "agent-1".to_string(),
            text: "nope".to_string(),
        });

        assert!(!outside.exists());
    }
}