malvin 0.2.5

Non-interactive research and coding agent
use serde_json::Value;

pub const RUN_COST_SUMMARY_PREFIX: &str = "COST: ";

fn format_token_field(tokens: &Value, key: &str) -> String {
    match tokens.get(key) {
        Some(Value::Null) | None => "n/a".to_string(),
        Some(v) => v
            .as_u64()
            .map_or_else(|| "n/a".to_string(), |n| n.to_string()),
    }
}

fn format_cost_field(cost: Option<&Value>, key: &str) -> String {
    cost.and_then(|c| c.get(key))
        .and_then(Value::as_f64)
        .map_or_else(|| "n/a".to_string(), |n| format!("{n:.4}"))
}

fn token_fields_fragment(json: &Value) -> String {
    let tokens = json.get("tokens");
    let steps = tokens
        .and_then(|t| t.get("steps"))
        .and_then(Value::as_u64)
        .unwrap_or(0);
    let tokens_in =
        tokens.map_or_else(|| "n/a".to_string(), |t| format_token_field(t, "tokens_in"));
    let tokens_out = tokens.map_or_else(
        || "n/a".to_string(),
        |t| format_token_field(t, "tokens_out"),
    );
    let cache_read = tokens.map_or_else(
        || "n/a".to_string(),
        |t| format_token_field(t, "cache_read"),
    );
    let cache_write = tokens.map_or_else(
        || "n/a".to_string(),
        |t| format_token_field(t, "cache_write"),
    );
    let reasoning =
        tokens.map_or_else(|| "n/a".to_string(), |t| format_token_field(t, "reasoning"));
    format!(
        "steps = {steps} tokens_in = {tokens_in} tokens_out = {tokens_out} cache_read = {cache_read} cache_write = {cache_write} reasoning = {reasoning}"
    )
}

fn cost_fields_fragment(json: &Value) -> String {
    let cost = json.get("cost");
    format!(
        "cost_in = {} cost_out = {} cost_read = {} cost_write = {} cost_tot = {}",
        format_cost_field(cost, "cost_in"),
        format_cost_field(cost, "cost_out"),
        format_cost_field(cost, "cost_read"),
        format_cost_field(cost, "cost_write"),
        format_cost_field(cost, "cost_tot"),
    )
}

#[must_use]
pub fn format_cost_stdout_line_from_json(json: &Value) -> String {
    format!(
        "{RUN_COST_SUMMARY_PREFIX}{} {}",
        token_fields_fragment(json),
        cost_fields_fragment(json)
    )
}

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

    #[test]
    fn cost_stdout_line_combines_token_and_cost_fields() {
        let json = json!({
            "tokens": {
                "steps": 9,
                "tokens_in": 221_270,
                "tokens_out": 15003,
                "cache_read": 1200,
                "cache_write": 80
            },
            "cost": {
                "cost_in": 0.05,
                "cost_out": 0.03,
                "cost_read": 0.002,
                "cost_write": 0.0022,
                "cost_tot": 0.0842
            }
        });
        let line = format_cost_stdout_line_from_json(&json);
        assert!(line.starts_with(RUN_COST_SUMMARY_PREFIX));
        assert_eq!(
            line,
            "COST: steps = 9 tokens_in = 221270 tokens_out = 15003 cache_read = 1200 cache_write = 80 reasoning = n/a cost_in = 0.0500 cost_out = 0.0300 cost_read = 0.0020 cost_write = 0.0022 cost_tot = 0.0842"
        );
    }

    #[test]
    fn cost_stdout_line_uses_na_when_blocks_absent() {
        let line = format_cost_stdout_line_from_json(&json!({}));
        assert_eq!(
            line,
            "COST: steps = 0 tokens_in = n/a tokens_out = n/a cache_read = n/a cache_write = n/a reasoning = n/a cost_in = n/a cost_out = n/a cost_read = n/a cost_write = n/a cost_tot = n/a"
        );
    }

    #[test]
    fn cost_stdout_line_formats_four_decimal_places() {
        let json = json!({
            "tokens": {
                "steps": 1,
                "tokens_in": 10,
                "tokens_out": 2
            },
            "cost": {
                "cost_in": 0.05,
                "cost_out": 0.03,
                "cost_read": 0.0,
                "cost_write": 0.0,
                "cost_tot": 0.0842
            }
        });
        let line = format_cost_stdout_line_from_json(&json);
        assert!(line.contains("cost_tot = 0.0842"));
        assert!(line.contains("cost_in = 0.0500"));
        assert!(line.contains("cache_read = n/a"));
    }
}