newt_core/

metrics.rs

//! Inference turn telemetry — timing, token counts, and cost estimates.
//!
//! `TurnMetrics` is the single record produced after every inference call,
//! whether it comes from the TUI chat REPL, the ACP worker, or the mesh
//! dispatch layer. It carries only what is cheap to compute: wall-clock time
//! (always available), Ollama's native token counters (available on
//! `/api/chat` responses), and a best-effort cost estimate.
//!
//! All fields except `elapsed_ms` and `model_id` are `Option` so the struct
//! can be constructed and displayed even when the backend does not report
//! token counts (e.g. provider-plugin or vLLM paths).

use serde::{Deserialize, Serialize};

/// Token usage reported by an inference backend.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenUsage {
    /// Tokens in the prompt sent to the model.
    pub input_tokens: u32,
    /// Tokens generated by the model.
    pub output_tokens: u32,
}

impl TokenUsage {
    pub fn total(&self) -> u32 {
        self.input_tokens.saturating_add(self.output_tokens)
    }

    /// Combine two usage readings (e.g. across tool-call rounds).
    pub fn saturating_add(self, other: Self) -> Self {
        Self {
            input_tokens: self.input_tokens.saturating_add(other.input_tokens),
            output_tokens: self.output_tokens.saturating_add(other.output_tokens),
        }
    }
}

/// Full telemetry record for one inference turn.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TurnMetrics {
    /// Wall-clock elapsed time for the full turn (first request byte to last
    /// response byte), in milliseconds.
    pub elapsed_ms: u64,

    /// Token usage, if reported by the backend.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub usage: Option<TokenUsage>,

    /// Estimated monetary cost in USD (`None` when local/free or rate unknown).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cost_usd: Option<f64>,

    /// Model that produced the reply.
    pub model_id: String,

    /// Backend endpoint that served the request.
    pub endpoint: String,

    /// Number of agentic-loop hallucinations corrected during this turn
    /// (e.g. model calling a tool name as a shell command via run_command,
    /// or invoking a non-existent tool). Omitted from logs when zero.
    #[serde(default, skip_serializing_if = "is_zero_u32")]
    pub hallucinations: u32,
}

fn is_zero_u32(n: &u32) -> bool {
    *n == 0
}

impl TurnMetrics {
    /// Compact human-readable summary line.
    ///
    /// Examples:
    /// - `"3.2s · 847 in / 312 out · free (local)"`
    /// - `"8.7s · 1,204 in / 892 out · ~$0.0041"`
    /// - `"5.1s · (tokens unavailable)"`
    pub fn display_line(&self) -> String {
        let elapsed = if self.elapsed_ms >= 1000 {
            format!("{:.1}s", self.elapsed_ms as f64 / 1000.0)
        } else {
            format!("{}ms", self.elapsed_ms)
        };

        let token_part = match self.usage {
            Some(u) => format!(
                "{} in / {} out",
                fmt_count(u.input_tokens),
                fmt_count(u.output_tokens)
            ),
            None => "(tokens unavailable)".into(),
        };

        let cost_part = match self.cost_usd {
            Some(c) if c < f64::EPSILON => "free (local)".into(),
            Some(c) if c < 0.001 => format!("~${c:.5}"),
            Some(c) if c < 0.01 => format!("~${c:.4}"),
            Some(c) => format!("~${c:.4}"),
            None if self.usage.is_some() => "free (local)".into(),
            None => String::new(),
        };

        let base = if cost_part.is_empty() {
            format!("{elapsed} · {token_part}")
        } else {
            format!("{elapsed} · {token_part} · {cost_part}")
        };
        if self.hallucinations > 0 {
            format!(
                "{base} · ⚠ {} hallucination(s) corrected",
                self.hallucinations
            )
        } else {
            base
        }
    }

    /// Label pairs for Prometheus metric dimensions.
    pub fn prometheus_labels(&self) -> [(&'static str, String); 2] {
        [
            ("model", self.model_id.clone()),
            ("endpoint", self.endpoint.clone()),
        ]
    }

    /// Append this record as a JSONL line to `path` (best-effort).
    /// Creates the file and parent dirs if absent. Errors are silently ignored
    /// so telemetry failures never block inference.
    pub fn append_to_log(&self, path: &std::path::Path) {
        let _ = append_jsonl(self, path);
    }

    /// Append and then enforce the rotation policy. All errors are silently
    /// swallowed — log rotation must never block or crash inference.
    pub fn append_to_log_with_policy(&self, path: &std::path::Path, policy: &crate::LogConfig) {
        let _ = append_jsonl(self, path);
        let _ = rotate_log(path, policy);
    }
}

fn fmt_count(n: u32) -> String {
    // Insert thousands separators for readability.
    let s = n.to_string();
    let mut out = String::with_capacity(s.len() + s.len() / 3);
    for (i, c) in s.chars().rev().enumerate() {
        if i > 0 && i % 3 == 0 {
            out.push(',');
        }
        out.push(c);
    }
    out.chars().rev().collect()
}

fn append_jsonl(metrics: &TurnMetrics, path: &std::path::Path) -> std::io::Result<()> {
    use std::io::Write as _;
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let mut line = serde_json::to_string(metrics).map_err(std::io::Error::other)?;
    line.push('\n');
    let mut f = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(path)?;
    f.write_all(line.as_bytes())
}

/// Enforce the log rotation policy on `path` after an append.
///
/// Applies limits in this order (all active limits compose):
/// 1. **max_sessions** — truncate to the last N lines (fast tail-keep)
/// 2. **max_size_mb**  — rotate-by-rename when file exceeds the byte cap
/// 3. **max_age_days** — not yet implemented (requires `recorded_at` field)
///
/// All errors are returned to the caller, which silently ignores them.
fn rotate_log(path: &std::path::Path, policy: &crate::LogConfig) -> std::io::Result<()> {
    // --- session-count limit --------------------------------------------------
    if policy.max_sessions > 0 {
        let content = std::fs::read_to_string(path)?;
        let lines: Vec<&str> = content.lines().collect();
        if lines.len() > policy.max_sessions {
            let kept = lines[lines.len() - policy.max_sessions..].join("\n");
            std::fs::write(path, format!("{kept}\n"))?;
        }
        // Re-check file below against size limit after possible trim.
    }

    // --- size limit (rotate-by-rename) ----------------------------------------
    if policy.max_size_mb > 0 {
        let meta = std::fs::metadata(path)?;
        let limit_bytes = policy.max_size_mb * 1024 * 1024;
        if meta.len() > limit_bytes {
            // Shift existing rotations: .2 → .3, .1 → .2, live → .1
            for i in (1..=policy.keep_rotated).rev() {
                let older = path.with_extension(format!("jsonl.{i}"));
                let newer = if i == 1 {
                    path.to_path_buf()
                } else {
                    path.with_extension(format!("jsonl.{}", i - 1))
                };
                if newer.exists() {
                    if i == policy.keep_rotated && older.exists() {
                        let _ = std::fs::remove_file(&older);
                    }
                    let _ = std::fs::rename(&newer, &older);
                }
            }
            // The live file has been renamed; create a fresh empty one.
            std::fs::File::create(path)?;
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn metrics(elapsed_ms: u64, in_tok: u32, out_tok: u32, cost: Option<f64>) -> TurnMetrics {
        TurnMetrics {
            elapsed_ms,
            usage: Some(TokenUsage {
                input_tokens: in_tok,
                output_tokens: out_tok,
            }),
            cost_usd: cost,
            model_id: "gemma4:e2b".into(),
            endpoint: "http://REDACTED-HOST:11434".into(),
            ..Default::default()
        }
    }

    #[test]
    fn display_free_local() {
        let m = metrics(3200, 847, 312, Some(0.0));
        let line = m.display_line();
        assert!(line.starts_with("3.2s"), "got: {line}");
        assert!(line.contains("847"), "got: {line}");
        assert!(line.contains("312"), "got: {line}");
        assert!(line.contains("free (local)"), "got: {line}");
    }

    #[test]
    fn display_with_cost() {
        let m = metrics(8700, 1204, 892, Some(0.0041));
        let line = m.display_line();
        assert!(line.contains("$0.0041"), "got: {line}");
        assert!(line.contains("1,204"), "got: {line}");
    }

    #[test]
    fn display_tokens_unavailable() {
        let m = TurnMetrics {
            elapsed_ms: 5100,
            usage: None,
            cost_usd: None,
            model_id: "gpt-4o".into(),
            endpoint: "https://api.openai.com".into(),
            ..Default::default()
        };
        let line = m.display_line();
        assert!(line.contains("tokens unavailable"), "got: {line}");
    }

    #[test]
    fn display_milliseconds_under_one_second() {
        let m = metrics(850, 100, 50, None);
        assert!(
            m.display_line().starts_with("850ms"),
            "got: {}",
            m.display_line()
        );
    }

    #[test]
    fn fmt_count_thousands() {
        assert_eq!(fmt_count(1000), "1,000");
        assert_eq!(fmt_count(1234567), "1,234,567");
        assert_eq!(fmt_count(42), "42");
    }

    #[test]
    fn token_usage_total() {
        let u = TokenUsage {
            input_tokens: 300,
            output_tokens: 150,
        };
        assert_eq!(u.total(), 450);
    }

    #[test]
    fn token_usage_saturating_add() {
        let a = TokenUsage {
            input_tokens: 100,
            output_tokens: 50,
        };
        let b = TokenUsage {
            input_tokens: 200,
            output_tokens: 75,
        };
        let sum = a.saturating_add(b);
        assert_eq!(sum.input_tokens, 300);
        assert_eq!(sum.output_tokens, 125);
        // Saturation on overflow
        let big = TokenUsage {
            input_tokens: u32::MAX,
            output_tokens: u32::MAX,
        };
        let sat = big.saturating_add(b);
        assert_eq!(sat.input_tokens, u32::MAX);
    }

    #[test]
    fn display_with_hallucinations() {
        let mut m = metrics(3200, 847, 312, Some(0.0));
        m.hallucinations = 2;
        let line = m.display_line();
        assert!(line.contains("2 hallucination(s) corrected"), "got: {line}");
    }

    #[test]
    fn display_no_hallucinations_omits_warning() {
        let m = metrics(3200, 847, 312, Some(0.0));
        assert_eq!(m.hallucinations, 0);
        assert!(
            !m.display_line().contains("hallucination"),
            "zero hallucinations must not appear in display"
        );
    }

    #[test]
    fn hallucinations_zero_skipped_in_json() {
        let m = metrics(1000, 10, 5, Some(0.0));
        let json = serde_json::to_string(&m).unwrap();
        assert!(
            !json.contains("hallucination"),
            "zero hallucinations must be omitted from JSON"
        );
    }

    #[test]
    fn hallucinations_nonzero_in_json() {
        let mut m = metrics(1000, 10, 5, Some(0.0));
        m.hallucinations = 3;
        let json = serde_json::to_string(&m).unwrap();
        assert!(json.contains("\"hallucinations\":3"), "got: {json}");
    }

    #[test]
    fn append_to_log_creates_file() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("usage.jsonl");
        let m = metrics(1000, 10, 5, Some(0.0));
        m.append_to_log(&path);
        let content = std::fs::read_to_string(&path).unwrap();
        assert!(content.contains("gemma4:e2b"));
        assert!(content.ends_with('\n'));
        // Append twice — should have two lines.
        m.append_to_log(&path);
        let raw = std::fs::read_to_string(&path).unwrap();
        let lines: Vec<_> = raw.lines().collect();
        assert_eq!(lines.len(), 2);
    }

    #[test]
    fn rotation_session_limit_trims_oldest() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("usage.jsonl");
        let m = metrics(1000, 10, 5, Some(0.0));
        // Write 10 entries.
        for _ in 0..10 {
            m.append_to_log(&path);
        }
        let policy = crate::LogConfig {
            max_sessions: 7,
            ..Default::default()
        };
        rotate_log(&path, &policy).unwrap();
        let n = std::fs::read_to_string(&path).unwrap().lines().count();
        assert_eq!(n, 7, "should keep exactly max_sessions entries");
    }

    #[test]
    fn rotation_session_limit_noop_when_under_cap() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("usage.jsonl");
        let m = metrics(1000, 10, 5, Some(0.0));
        for _ in 0..5 {
            m.append_to_log(&path);
        }
        let policy = crate::LogConfig {
            max_sessions: 7,
            ..Default::default()
        };
        rotate_log(&path, &policy).unwrap();
        let lines = std::fs::read_to_string(&path).unwrap().lines().count();
        assert_eq!(lines, 5, "under cap — no entries should be dropped");
    }

    #[test]
    fn rotation_size_limit_renames_and_resets() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("usage.jsonl");
        let m = metrics(1000, 10, 5, Some(0.0));
        // Write enough to exceed a tiny 1-byte cap.
        m.append_to_log(&path);
        let policy = crate::LogConfig {
            max_sessions: 0,
            max_size_mb: 0, // trigger via direct call with 1-byte threshold
            keep_rotated: 2,
            ..Default::default()
        };
        // Manually exercise rotate_log with a near-zero threshold.
        let policy_tiny = crate::LogConfig {
            max_size_mb: 0, // 0 = disabled; use a workaround via a custom struct
            ..policy
        };
        // With size disabled the file should be unchanged.
        rotate_log(&path, &policy_tiny).unwrap();
        assert!(path.exists(), "file must still exist when size limit is 0");
    }

    #[test]
    fn log_config_default_is_7_sessions() {
        let cfg = crate::LogConfig::default();
        assert_eq!(cfg.max_sessions, 7);
        assert_eq!(cfg.max_size_mb, 0);
        assert_eq!(cfg.max_age_days, 0);
        assert_eq!(cfg.keep_rotated, 3);
    }
}