agentwerk 0.1.13

//! Run-time stats. One [`Stats`] struct of atomic counters records every
//! observable event and exposes inherent read accessors. Each domain
//! interacts with its own write-only protocol — [`LoopStats`] for the
//! agent loop, `TicketStats` for the ticket system — so a domain
//! cannot reach another domain's events. The wiring is internal: the
//! caller never sees `Stats` at construction time, only afterwards
//! through `TicketSystem::stats()`.
//!
//! Lock-free for counter increments; readers do one atomic load per
//! call.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// Recorder protocol for the agent loop. Each agent holds an
/// `Arc<dyn LoopStats + Send + Sync>` and reports loop events through
/// it. Write-only; reads happen on `Stats` directly.
pub trait LoopStats: Send + Sync {
    fn record_turn(&self);
    fn record_request(&self, input_tokens: u64, output_tokens: u64);
    fn record_tool_call(&self);
    fn record_error(&self);
}

/// Recorder protocol for the ticket system. The ticket system holds an
/// `Arc<Stats>` directly but only exercises these methods, so by
/// convention the ticket-domain surface is narrow.
pub(crate) trait TicketStats: Send + Sync {
    fn record_created(&self);
    /// First call wins (CAS into `started_at`); later calls are no-ops.
    fn record_started(&self, when: u64);
    /// Adds `ticket_duration.as_secs()` and `work_duration.as_secs()` to
    /// the corresponding atomic sums.
    fn record_done(&self, ticket_duration: Duration, work_duration: Duration);
    fn record_failed(&self, ticket_duration: Duration, work_duration: Duration);
}

/// Run-wide counters. Implements every recorder protocol; exposes
/// inherent read methods for the caller to consume after a run.
pub struct Stats {
    turns: AtomicU64,
    requests: AtomicU64,
    tool_calls: AtomicU64,
    errors: AtomicU64,
    input_tokens: AtomicU64,
    output_tokens: AtomicU64,
    tickets_created: AtomicU64,
    tickets_done: AtomicU64,
    tickets_failed: AtomicU64,

    /// Run-start wall clock (millis since epoch). 0 = unset; first
    /// `record_started` wins via CAS.
    started_at: AtomicU64,
    /// Run-end wall clock (millis since epoch). 0 = still running;
    /// `mark_finished` stamps it when the watcher fires.
    finished_at: AtomicU64,
    /// Sum of finished tickets' creation→terminal durations, seconds.
    total_ticket_duration: AtomicU64,
    /// Sum of finished tickets' started→terminal durations, seconds.
    /// With concurrent agents this can exceed the run's wall-clock
    /// duration.
    total_work_duration: AtomicU64,
    /// Lazy-init map of nested counter slices keyed by ticket label.
    /// Always empty on a slice itself; populated only on the run-wide
    /// `Stats` owned by `TicketSystem`.
    label_stats: Mutex<HashMap<String, Arc<Stats>>>,
}

impl Stats {
    pub fn new() -> Self {
        Self {
            turns: AtomicU64::new(0),
            requests: AtomicU64::new(0),
            tool_calls: AtomicU64::new(0),
            errors: AtomicU64::new(0),
            input_tokens: AtomicU64::new(0),
            output_tokens: AtomicU64::new(0),
            tickets_created: AtomicU64::new(0),
            tickets_done: AtomicU64::new(0),
            tickets_failed: AtomicU64::new(0),
            started_at: AtomicU64::new(0),
            finished_at: AtomicU64::new(0),
            total_ticket_duration: AtomicU64::new(0),
            total_work_duration: AtomicU64::new(0),
            label_stats: Mutex::new(HashMap::new()),
        }
    }

    /// Live counters scoped to one ticket label. Lazy-init on first
    /// access; subsequent calls return the same `Arc`. Reads use the
    /// same accessors as the run-wide `Stats`; `run_duration()` is
    /// always `None` on a slice (run timing stays global).
    pub fn stats_for_label(&self, label: &str) -> Arc<Stats> {
        let mut map = self.label_stats.lock().unwrap();
        map.entry(label.to_string())
            .or_insert_with(|| Arc::new(Stats::new()))
            .clone()
    }

    pub fn turns(&self) -> u64 {
        self.turns.load(Ordering::Relaxed)
    }

    pub fn requests(&self) -> u64 {
        self.requests.load(Ordering::Relaxed)
    }

    pub fn tool_calls(&self) -> u64 {
        self.tool_calls.load(Ordering::Relaxed)
    }

    pub fn errors(&self) -> u64 {
        self.errors.load(Ordering::Relaxed)
    }

    pub fn input_tokens(&self) -> u64 {
        self.input_tokens.load(Ordering::Relaxed)
    }

    pub fn output_tokens(&self) -> u64 {
        self.output_tokens.load(Ordering::Relaxed)
    }

    pub fn tickets_created(&self) -> u64 {
        self.tickets_created.load(Ordering::Relaxed)
    }

    pub fn tickets_done(&self) -> u64 {
        self.tickets_done.load(Ordering::Relaxed)
    }

    pub fn tickets_failed(&self) -> u64 {
        self.tickets_failed.load(Ordering::Relaxed)
    }

    /// Wall time of the run, measured from the first `record_started`
    /// call. Live while the run is in progress; freezes at
    /// `finished_at - started_at` once `mark_finished` has fired.
    /// `None` until the run has started.
    pub fn run_duration(&self) -> Option<Duration> {
        let s = self.started_at.load(Ordering::Relaxed);
        if s == 0 {
            return None;
        }
        let f = self.finished_at.load(Ordering::Relaxed);
        if f != 0 && f >= s {
            return Some(Duration::from_millis(f - s));
        }
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .ok()?
            .as_millis() as u64;
        Some(Duration::from_millis(now.saturating_sub(s)))
    }

    /// `tickets_done / (tickets_done + tickets_failed)`. `None` when
    /// no ticket has reached a terminal state.
    pub fn tickets_success_rate(&self) -> Option<f64> {
        let done = self.tickets_done.load(Ordering::Relaxed);
        let failed = self.tickets_failed.load(Ordering::Relaxed);
        let total = done + failed;
        if total == 0 {
            None
        } else {
            Some(done as f64 / total as f64)
        }
    }

    /// Sum of finished tickets' creation→terminal spans.
    pub fn ticket_duration(&self) -> Duration {
        Duration::from_secs(self.total_ticket_duration.load(Ordering::Relaxed))
    }

    /// Mean of finished tickets' creation→terminal spans. `None`
    /// while no ticket has finished.
    pub fn avg_ticket_duration(&self) -> Option<Duration> {
        let n =
            self.tickets_done.load(Ordering::Relaxed) + self.tickets_failed.load(Ordering::Relaxed);
        if n == 0 {
            None
        } else {
            let secs = self.total_ticket_duration.load(Ordering::Relaxed);
            Some(Duration::from_secs(secs / n))
        }
    }

    /// Sum of finished tickets' started→terminal spans. With
    /// concurrent agents this aggregates work across all of them, so
    /// it can exceed `run_duration`.
    pub fn work_duration(&self) -> Duration {
        Duration::from_secs(self.total_work_duration.load(Ordering::Relaxed))
    }

    /// Mean of finished tickets' started→terminal spans. `None` while
    /// no ticket has finished.
    pub fn avg_work_duration(&self) -> Option<Duration> {
        let n =
            self.tickets_done.load(Ordering::Relaxed) + self.tickets_failed.load(Ordering::Relaxed);
        if n == 0 {
            None
        } else {
            let secs = self.total_work_duration.load(Ordering::Relaxed);
            Some(Duration::from_secs(secs / n))
        }
    }

    /// Stamp the run's finish wall-clock. Idempotent in practice
    /// (the watcher fires once); successive calls overwrite, which
    /// is fine.
    pub(crate) fn mark_finished(&self, when: u64) {
        self.finished_at.store(when, Ordering::Relaxed);
    }

    /// Rebuild from already-loaded tickets. Used as the fallback when
    /// the stats file is absent or unreadable; otherwise `Stats::load`
    /// is preferred to keep observability continuous across restarts.
    pub(crate) fn derive(tickets: &HashMap<String, crate::agents::tickets::Ticket>) -> Self {
        let stats = Stats::new();
        for t in tickets.values() {
            TicketStats::record_created(&stats);
            for label in t.labels.iter() {
                TicketStats::record_created(&*stats.stats_for_label(label));
            }
            let ticket_dur = ticket_duration(t).unwrap_or_default();
            let work_dur = work_duration(t).unwrap_or_default();
            match t.status() {
                "done" => {
                    TicketStats::record_done(&stats, ticket_dur, work_dur);
                    for label in t.labels.iter() {
                        TicketStats::record_done(
                            &*stats.stats_for_label(label),
                            ticket_dur,
                            work_dur,
                        );
                    }
                }
                "failed" => {
                    TicketStats::record_failed(&stats, ticket_dur, work_dur);
                    for label in t.labels.iter() {
                        TicketStats::record_failed(
                            &*stats.stats_for_label(label),
                            ticket_dur,
                            work_dur,
                        );
                    }
                }
                _ => {}
            }
        }
        stats
    }

    /// Sugar over the `Persist` impl for the keyless case.
    pub(crate) fn load(dir: &std::path::Path) -> std::io::Result<Self> {
        <Self as crate::persistence::Persist>::load(dir, &())
    }
}

impl Stats {
    pub(crate) const FILE: &'static str = "stats.json";
}

impl crate::persistence::Persist for Stats {
    type Key = ();

    fn save(&self, dir: &std::path::Path) -> std::io::Result<()> {
        let body = serde_json::to_vec_pretty(&self.as_json()).map_err(std::io::Error::other)?;
        crate::persistence::write_atomic(&dir.join(Self::FILE), &body)
    }

    fn load(dir: &std::path::Path, _: &Self::Key) -> std::io::Result<Self> {
        let bytes = std::fs::read(dir.join(Self::FILE))?;
        let value: serde_json::Value = serde_json::from_slice(&bytes).map_err(std::io::Error::other)?;
        let stats = Stats::new();
        stats.load_fields(&value);
        if let Some(labels) = value.get("labels").and_then(|v| v.as_object()) {
            for (name, body) in labels {
                let slice = stats.stats_for_label(name);
                slice.load_fields(body);
            }
        }
        Ok(stats)
    }
}

impl Stats {
    fn as_json(&self) -> serde_json::Value {
        let labels = self.label_stats.lock().unwrap();
        let label_obj: serde_json::Map<String, serde_json::Value> = labels
            .iter()
            .map(|(name, slice)| (name.clone(), slice.fields_as_json()))
            .collect();
        let mut body = self.fields_as_json();
        if let serde_json::Value::Object(map) = &mut body {
            map.insert("labels".into(), serde_json::Value::Object(label_obj));
        }
        body
    }

    fn fields_as_json(&self) -> serde_json::Value {
        serde_json::json!({
            "turns": self.turns.load(Ordering::Relaxed),
            "requests": self.requests.load(Ordering::Relaxed),
            "tool_calls": self.tool_calls.load(Ordering::Relaxed),
            "errors": self.errors.load(Ordering::Relaxed),
            "input_tokens": self.input_tokens.load(Ordering::Relaxed),
            "output_tokens": self.output_tokens.load(Ordering::Relaxed),
            "tickets_created": self.tickets_created.load(Ordering::Relaxed),
            "tickets_done": self.tickets_done.load(Ordering::Relaxed),
            "tickets_failed": self.tickets_failed.load(Ordering::Relaxed),
            "total_ticket_duration_secs": self.total_ticket_duration.load(Ordering::Relaxed),
            "total_work_duration_secs": self.total_work_duration.load(Ordering::Relaxed),
        })
    }

    fn load_fields(&self, value: &serde_json::Value) {
        let get = |key: &str| value.get(key).and_then(|v| v.as_u64()).unwrap_or(0);
        self.turns.store(get("turns"), Ordering::Relaxed);
        self.requests.store(get("requests"), Ordering::Relaxed);
        self.tool_calls.store(get("tool_calls"), Ordering::Relaxed);
        self.errors.store(get("errors"), Ordering::Relaxed);
        self.input_tokens
            .store(get("input_tokens"), Ordering::Relaxed);
        self.output_tokens
            .store(get("output_tokens"), Ordering::Relaxed);
        self.tickets_created
            .store(get("tickets_created"), Ordering::Relaxed);
        self.tickets_done
            .store(get("tickets_done"), Ordering::Relaxed);
        self.tickets_failed
            .store(get("tickets_failed"), Ordering::Relaxed);
        self.total_ticket_duration
            .store(get("total_ticket_duration_secs"), Ordering::Relaxed);
        self.total_work_duration
            .store(get("total_work_duration_secs"), Ordering::Relaxed);
    }
}

fn ticket_duration(t: &crate::agents::tickets::Ticket) -> Option<Duration> {
    let end = t.finished_at().or_else(|| t.failed_at())?;
    Some(Duration::from_millis(end.saturating_sub(t.created_at())))
}

fn work_duration(t: &crate::agents::tickets::Ticket) -> Option<Duration> {
    let end = t.finished_at().or_else(|| t.failed_at())?;
    let start = t.started_at()?;
    Some(Duration::from_millis(end.saturating_sub(start)))
}

impl Default for Stats {
    fn default() -> Self {
        Self::new()
    }
}

impl LoopStats for Stats {
    fn record_turn(&self) {
        self.turns.fetch_add(1, Ordering::Relaxed);
    }

    fn record_request(&self, input_tokens: u64, output_tokens: u64) {
        self.requests.fetch_add(1, Ordering::Relaxed);
        self.input_tokens.fetch_add(input_tokens, Ordering::Relaxed);
        self.output_tokens
            .fetch_add(output_tokens, Ordering::Relaxed);
    }

    fn record_tool_call(&self) {
        self.tool_calls.fetch_add(1, Ordering::Relaxed);
    }

    fn record_error(&self) {
        self.errors.fetch_add(1, Ordering::Relaxed);
    }
}

impl TicketStats for Stats {
    fn record_created(&self) {
        self.tickets_created.fetch_add(1, Ordering::Relaxed);
    }

    fn record_started(&self, when: u64) {
        // First call wins. Subsequent claims (Path A reclaim, late
        // bind) leave the original run-start untouched.
        let _ = self
            .started_at
            .compare_exchange(0, when, Ordering::Relaxed, Ordering::Relaxed);
    }

    fn record_done(&self, ticket_duration: Duration, work_duration: Duration) {
        self.tickets_done.fetch_add(1, Ordering::Relaxed);
        self.total_ticket_duration
            .fetch_add(ticket_duration.as_secs(), Ordering::Relaxed);
        self.total_work_duration
            .fetch_add(work_duration.as_secs(), Ordering::Relaxed);
    }

    fn record_failed(&self, ticket_duration: Duration, work_duration: Duration) {
        self.tickets_failed.fetch_add(1, Ordering::Relaxed);
        self.total_ticket_duration
            .fetch_add(ticket_duration.as_secs(), Ordering::Relaxed);
        self.total_work_duration
            .fetch_add(work_duration.as_secs(), Ordering::Relaxed);
    }
}

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

    #[test]
    fn fresh_stats_are_zero() {
        let s = Stats::new();
        assert_eq!(s.turns(), 0);
        assert_eq!(s.requests(), 0);
        assert_eq!(s.tool_calls(), 0);
        assert_eq!(s.errors(), 0);
        assert_eq!(s.input_tokens(), 0);
        assert_eq!(s.output_tokens(), 0);
        assert_eq!(s.tickets_created(), 0);
        assert_eq!(s.tickets_done(), 0);
        assert_eq!(s.tickets_failed(), 0);
        assert_eq!(s.ticket_duration(), Duration::ZERO);
        assert_eq!(s.work_duration(), Duration::ZERO);
        assert!(s.run_duration().is_none());
        assert!(s.avg_ticket_duration().is_none());
        assert!(s.avg_work_duration().is_none());
        assert!(s.tickets_success_rate().is_none());
    }

    #[test]
    fn loop_stats_writes_show_up_in_reads() {
        let s = Stats::new();
        s.record_turn();
        s.record_turn();
        s.record_request(10, 5);
        s.record_request(2, 1);
        s.record_tool_call();
        s.record_error();

        assert_eq!(s.turns(), 2);
        assert_eq!(s.requests(), 2);
        assert_eq!(s.tool_calls(), 1);
        assert_eq!(s.errors(), 1);
        assert_eq!(s.input_tokens(), 12);
        assert_eq!(s.output_tokens(), 6);
    }

    #[test]
    fn ticket_stats_writes_show_up_in_reads() {
        let s = Stats::new();
        s.record_created();
        s.record_created();
        s.record_done(Duration::from_secs(3), Duration::from_secs(2));
        s.record_failed(Duration::from_secs(5), Duration::from_secs(4));

        assert_eq!(s.tickets_created(), 2);
        assert_eq!(s.tickets_done(), 1);
        assert_eq!(s.tickets_failed(), 1);
        assert_eq!(s.ticket_duration(), Duration::from_secs(8));
        assert_eq!(s.work_duration(), Duration::from_secs(6));
    }

    #[test]
    fn record_started_first_call_wins() {
        let s = Stats::new();
        s.record_started(1_000);
        s.record_started(2_000);
        s.record_started(3_000);
        s.mark_finished(4_500);
        assert_eq!(s.run_duration(), Some(Duration::from_millis(3500)));
    }

    #[test]
    fn run_duration_freezes_at_finish() {
        let s = Stats::new();
        assert!(s.run_duration().is_none());
        s.record_started(1_000);
        // Live before finish: anchored at started_at = 1_000ms epoch, so the
        // delta to "now" is enormous. We just check it's some duration.
        assert!(s.run_duration().is_some());
        s.mark_finished(2_500);
        assert_eq!(s.run_duration(), Some(Duration::from_millis(1500)));
        // Stays frozen on a subsequent call.
        assert_eq!(s.run_duration(), Some(Duration::from_millis(1500)));
    }

    #[test]
    fn tickets_success_rate_done_failed_mix() {
        let s = Stats::new();
        s.record_done(Duration::from_secs(1), Duration::from_secs(1));
        s.record_done(Duration::from_secs(2), Duration::from_secs(2));
        s.record_failed(Duration::from_secs(3), Duration::from_secs(3));
        let rate = s.tickets_success_rate().unwrap();
        assert!((rate - 2.0 / 3.0).abs() < 1e-9, "rate = {rate}");
    }

    #[test]
    fn tickets_success_rate_none_when_nothing_finished() {
        let s = Stats::new();
        assert!(s.tickets_success_rate().is_none());
    }

    #[test]
    fn avg_ticket_duration_is_arithmetic_mean() {
        let s = Stats::new();
        s.record_done(Duration::from_secs(2), Duration::from_secs(2));
        s.record_failed(Duration::from_secs(4), Duration::from_secs(4));
        assert_eq!(s.avg_ticket_duration(), Some(Duration::from_secs(3)));
    }

    #[test]
    fn avg_work_duration_is_arithmetic_mean() {
        let s = Stats::new();
        s.record_done(Duration::from_secs(3), Duration::from_secs(2));
        s.record_failed(Duration::from_secs(5), Duration::from_secs(4));
        assert_eq!(s.avg_work_duration(), Some(Duration::from_secs(3)));
    }

    #[test]
    fn stats_for_label_returns_same_slice_on_repeat_access() {
        let s = Stats::new();
        let a = s.stats_for_label("scan");
        let b = s.stats_for_label("scan");
        assert!(Arc::ptr_eq(&a, &b));
    }

    #[test]
    fn stats_for_label_slice_records_independently() {
        let s = Stats::new();
        let slice = s.stats_for_label("scan");
        slice.record_turn();
        slice.record_request(10, 5);
        assert_eq!(slice.turns(), 1);
        assert_eq!(slice.input_tokens(), 10);
        assert_eq!(slice.output_tokens(), 5);
        assert_eq!(s.turns(), 0);
        assert_eq!(s.input_tokens(), 0);
    }

    #[test]
    fn stats_for_label_slice_run_duration_is_none() {
        let s = Stats::new();
        let slice = s.stats_for_label("scan");
        slice.record_done(Duration::from_secs(2), Duration::from_secs(1));
        assert!(slice.run_duration().is_none());
        assert_eq!(slice.tickets_done(), 1);
    }

    #[test]
    fn work_duration_can_exceed_run_duration_with_concurrency() {
        // Two tickets, each 5s of work, finished in a 6s window —
        // models 2 agents working in parallel.
        let s = Stats::new();
        s.record_started(1_000);
        s.record_done(Duration::from_secs(5), Duration::from_secs(5));
        s.record_done(Duration::from_secs(5), Duration::from_secs(5));
        s.mark_finished(7_000);
        assert_eq!(s.run_duration(), Some(Duration::from_secs(6)));
        assert_eq!(s.work_duration(), Duration::from_secs(10));
    }

    #[test]
    fn stats_round_trips_through_save_load() {
        let dir = crate::test_util::TempDir::new().unwrap();

        let s = Stats::new();
        s.record_turn();
        s.record_turn();
        s.record_request(100, 50);
        s.record_tool_call();
        s.record_error();
        s.record_created();
        s.record_done(Duration::from_secs(7), Duration::from_secs(5));
        s.record_failed(Duration::from_secs(3), Duration::from_secs(2));

        let slice = s.stats_for_label("scan");
        slice.record_turn();
        slice.record_request(40, 20);
        slice.record_created();
        slice.record_done(Duration::from_secs(4), Duration::from_secs(3));

        use crate::persistence::Persist;
        s.save(dir.path()).unwrap();
        let restored = Stats::load(dir.path()).unwrap();
        assert_eq!(restored.turns(), 2);
        assert_eq!(restored.requests(), 1);
        assert_eq!(restored.tool_calls(), 1);
        assert_eq!(restored.errors(), 1);
        assert_eq!(restored.input_tokens(), 100);
        assert_eq!(restored.output_tokens(), 50);
        assert_eq!(restored.tickets_created(), 1);
        assert_eq!(restored.tickets_done(), 1);
        assert_eq!(restored.tickets_failed(), 1);
        assert_eq!(restored.ticket_duration(), Duration::from_secs(10));
        assert_eq!(restored.work_duration(), Duration::from_secs(7));

        let restored_slice = restored.stats_for_label("scan");
        assert_eq!(restored_slice.turns(), 1);
        assert_eq!(restored_slice.input_tokens(), 40);
        assert_eq!(restored_slice.tickets_done(), 1);
        assert_eq!(restored_slice.ticket_duration(), Duration::from_secs(4));
    }
}