newt-core 0.6.5

//! 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)
    }
}

/// 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,
}

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(),
        };

        if cost_part.is_empty() {
            format!("{elapsed} · {token_part}")
        } else {
            format!("{elapsed} · {token_part} · {cost_part}")
        }
    }

    /// 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);
    }
}

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())
}

// ---------------------------------------------------------------------------
// 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://dgx1.home.lab:11434".into(),
        }
    }

    #[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(),
        };
        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 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);
    }
}