klieo-runlog 3.5.0

use super::model_alias::family_alias;
use super::*;

#[test]
fn cost_for_cached_prices_each_tier_separately() {
    // prompt $1/1k, completion $2/1k, cache_read $0.1/1k, cache_creation $1.25/1k.
    let r = Rates::new(1.0, 2.0).with_cache_rates(0.1, 1.25);
    // 1000 prompt, 1000 completion, 2000 cache_read, 1000 cache_creation.
    let c = r.cost_for_cached(1000, 1000, 2000, 1000);
    assert!((c.prompt_usd - 1.0).abs() < 1e-9);
    assert!((c.completion_usd - 2.0).abs() < 1e-9);
    assert!(
        (c.cache_usd - (0.2 + 1.25)).abs() < 1e-9,
        "cache_usd = read + creation"
    );
    assert!(
        (c.total_usd - 4.45).abs() < 1e-9,
        "total sums all four tiers"
    );
}

#[test]
fn cost_for_is_cache_free_and_backward_compatible() {
    let r = Rates::new(1.0, 2.0).with_cache_rates(0.1, 1.25);
    let c = r.cost_for(1000, 1000);
    assert!(c.cache_usd.abs() < 1e-12, "no cache tokens → no cache cost");
    assert!((c.total_usd - 3.0).abs() < 1e-9);
}

#[test]
fn anthropic_default_rates_carry_cache_tiers() {
    let t = PriceTable::with_default_rates();
    let r = t
        .lookup("anthropic", "claude-sonnet-4-6")
        .expect("sonnet family");
    // Anthropic: cache read ~0.1x input, cache creation ~1.25x input.
    assert!((r.cache_read_usd_per_1k() - 0.0003).abs() < 1e-9);
    assert!((r.cache_creation_usd_per_1k() - 0.00375).abs() < 1e-9);
}

#[test]
fn exact_lookup_hit() {
    let t = PriceTable::with_default_rates();
    let r = t.lookup("openai", "gpt-4o").expect("exact hit");
    assert!((r.prompt_usd_per_1k - 0.0025).abs() < 1e-9);
    assert!((r.completion_usd_per_1k - 0.01).abs() < 1e-9);
}

#[test]
fn family_alias_fallback_hit() {
    let t = PriceTable::with_default_rates();
    // Versioned Anthropic ID should collapse to the -x family alias.
    let r = t
        .lookup("anthropic", "claude-sonnet-4-6")
        .expect("family alias hit");
    assert!((r.prompt_usd_per_1k - 0.003).abs() < 1e-9);
    let r2 = t
        .lookup("anthropic", "claude-opus-4-7")
        .expect("opus family alias hit");
    assert!((r2.prompt_usd_per_1k - 0.015).abs() < 1e-9);
}

#[test]
fn lookup_miss_returns_none() {
    let t = PriceTable::with_default_rates();
    assert!(t.lookup("openai", "no-such-model").is_none());
    // gpt-4o-mini-bogus has alpha tail; no exact, no family alias hit
    // (no openai wildcard registered).
    assert!(t.lookup("openai", "gpt-4o-mini-bogus").is_none());
    assert!(t.lookup("unknown-provider", "gpt-4o").is_none());
}

#[test]
fn builtin_table_contains_all_six_providers_models() {
    let t = PriceTable::with_default_rates();
    assert!(t.lookup("anthropic", "claude-sonnet-4-x").is_some());
    assert!(t.lookup("anthropic", "claude-opus-4-x").is_some());
    assert!(t.lookup("anthropic", "claude-haiku-4-x").is_some());
    assert!(t.lookup("openai", "gpt-4o").is_some());
    assert!(t.lookup("openai", "gpt-4o-mini").is_some());
    assert!(t.lookup("google", "gemini-2.0-flash").is_some());
    assert!(t.lookup("ollama", "*").is_some());
}

#[test]
fn ollama_zero_rate_yields_zero_cost() {
    let t = PriceTable::with_default_rates();
    let r = t.lookup("ollama", "*").unwrap();
    let c = r.cost_for(10_000, 5_000);
    assert!(c.prompt_usd.abs() < 1e-12);
    assert!(c.completion_usd.abs() < 1e-12);
    assert!(c.total_usd.abs() < 1e-12);
}

#[test]
fn cost_arithmetic_matches_hand_compute() {
    // 1000 prompt @ $0.003/1k = $0.003
    // 500 completion @ $0.015/1k = $0.0075
    // total = $0.0105
    let r = Rates::new(0.003, 0.015);
    let c = r.cost_for(1000, 500);
    assert!((c.prompt_usd - 0.003).abs() <= 1e-9);
    assert!((c.completion_usd - 0.0075).abs() <= 1e-9);
    assert!((c.total_usd - 0.0105).abs() <= 1e-9);
}

#[test]
fn family_alias_only_rewrites_trailing_digits() {
    // `gpt-4o-mini` has alpha tail → no rewrite, no fallback.
    assert!(family_alias("gpt-4o-mini").is_none());
    // gpt-4o has a `-` separator but its tail "4o" is alphanumeric (not
    // pure digits), so family_alias returns None.
    assert!(family_alias("gpt-4o").is_none());
    // `claude-sonnet-4-6` rewrites to `claude-sonnet-4-x`.
    assert_eq!(
        family_alias("claude-sonnet-4-6").as_deref(),
        Some("claude-sonnet-4-x")
    );
}

#[test]
fn set_overrides_existing_rate() {
    let mut t = PriceTable::with_default_rates();
    t.set("openai", "gpt-4o", Rates::new(0.001, 0.002)).unwrap();
    let r = t.lookup("openai", "gpt-4o").unwrap();
    assert!((r.prompt_usd_per_1k - 0.001).abs() < 1e-9);
    assert!((r.completion_usd_per_1k - 0.002).abs() < 1e-9);
}

#[test]
fn rates_round_trip_through_json() {
    let r = Rates::new(0.003, 0.015);
    let s = serde_json::to_string(&r).unwrap();
    let back: Rates = serde_json::from_str(&s).unwrap();
    assert!((back.prompt_usd_per_1k - 0.003).abs() < 1e-12);
    assert!((back.completion_usd_per_1k - 0.015).abs() < 1e-12);
}

#[test]
fn try_new_rejects_nan() {
    let err = Rates::try_new(f64::NAN, 0.001).unwrap_err();
    assert!(matches!(err, RatesError::NotFinite { .. }));
    let err = Rates::try_new(0.001, f64::NAN).unwrap_err();
    assert!(matches!(err, RatesError::NotFinite { .. }));
}

#[test]
fn try_new_rejects_negative() {
    let err = Rates::try_new(-0.001, 0.001).unwrap_err();
    assert!(matches!(err, RatesError::Negative { .. }));
    let err = Rates::try_new(0.001, -0.001).unwrap_err();
    assert!(matches!(err, RatesError::Negative { .. }));
}

#[test]
fn try_new_rejects_infinity() {
    let err = Rates::try_new(f64::INFINITY, 0.001).unwrap_err();
    assert!(matches!(err, RatesError::NotFinite { .. }));
    let err = Rates::try_new(0.001, f64::NEG_INFINITY).unwrap_err();
    assert!(matches!(err, RatesError::NotFinite { .. }));
}

#[test]
fn family_alias_strips_anthropic_iso_date() {
    let mut t = PriceTable::new();
    t.set("anthropic", "claude-3-5-sonnet-x", Rates::new(0.003, 0.015))
        .unwrap();
    let r = t
        .lookup("anthropic", "claude-3-5-sonnet-2024-11-20")
        .expect("ISO-date strip then family alias");
    assert!((r.prompt_usd_per_1k - 0.003).abs() < 1e-9);
}

#[test]
fn family_alias_strips_anthropic_compact_date() {
    let mut t = PriceTable::new();
    t.set("anthropic", "claude-3-5-sonnet-x", Rates::new(0.003, 0.015))
        .unwrap();
    let r = t
        .lookup("anthropic", "claude-3-5-sonnet-20241022")
        .expect("compact-date strip then family alias");
    assert!((r.prompt_usd_per_1k - 0.003).abs() < 1e-9);
}

#[test]
fn family_alias_strips_ollama_colon_tag() {
    let mut t = PriceTable::new();
    t.set("ollama", "qwen2.5", Rates::new(0.0, 0.0)).unwrap();
    let r = t
        .lookup("ollama", "qwen2.5:14b")
        .expect("ollama colon-tag head");
    assert!(r.prompt_usd_per_1k.abs() < 1e-12);
}

#[test]
fn rates_deserialize_rejects_negative_via_serde_try_from() {
    // Negative input must fail closed at the deserialise boundary,
    // not land as a poisoned Rates value.
    let err = serde_json::from_str::<Rates>(
        r#"{"prompt_usd_per_1k": -0.001, "completion_usd_per_1k": 0.001}"#,
    )
    .unwrap_err();
    let msg = err.to_string();
    assert!(
        msg.contains("non-negative") || msg.contains("Negative"),
        "expected negative-rate message, got {msg}"
    );
}

#[test]
fn rates_deserialize_rejects_infinite_via_serde_try_from() {
    // serde_json represents floats as JSON numbers; we encode +inf via
    // an out-of-range exponent in JSON syntax. Most serde_json builds
    // accept this as a finite f64::MAX-ish value, so we round-trip an
    // explicit f64::INFINITY through serde_json::Value to bypass the
    // string-parser. The point is that even when an upstream feeder
    // produces a non-finite Number, the try_from gate rejects it.
    let v = serde_json::json!({
        "prompt_usd_per_1k": f64::MAX,
        "completion_usd_per_1k": f64::MAX,
    });
    // Use direct float arithmetic to manufacture an Infinity through
    // the gate: serde_json::Number::from_f64 returns None for non-finite,
    // so we cannot encode infinity as a JSON Value. Instead, drive the
    // TryFrom impl directly to confirm the gate's behaviour.
    let _ = v;
    let unchecked = RatesUnchecked {
        prompt_usd_per_1k: f64::INFINITY,
        completion_usd_per_1k: 0.001,
        cache_read_usd_per_1k: 0.0,
        cache_creation_usd_per_1k: 0.0,
    };
    let err = Rates::try_from(unchecked).unwrap_err();
    assert!(matches!(err, RatesError::NotFinite { .. }));
}

#[test]
fn negative_cache_rate_is_rejected_by_gate() {
    let unchecked = RatesUnchecked {
        prompt_usd_per_1k: 0.003,
        completion_usd_per_1k: 0.015,
        cache_read_usd_per_1k: -0.1,
        cache_creation_usd_per_1k: 0.0,
    };
    let err = Rates::try_from(unchecked).unwrap_err();
    assert!(matches!(err, RatesError::CacheRateInvalid { .. }));
}

#[test]
fn family_alias_falls_through_to_provider_wildcard() {
    let t = PriceTable::with_default_rates();
    // Built-in ollama "*" should hit any unrecognised ollama model.
    let r = t
        .lookup("ollama", "qwen2.5:14b")
        .expect("ollama wildcard fallback");
    assert!(r.prompt_usd_per_1k.abs() < 1e-12);
    let r2 = t
        .lookup("ollama", "llama3.2")
        .expect("ollama wildcard fallback no tag");
    assert!(r2.prompt_usd_per_1k.abs() < 1e-12);
}

#[test]
fn gemini_provider_alias_resolves_to_google_keys() {
    let t = PriceTable::with_default_rates();
    // Clients report "gemini"; the table keys Google models under "google".
    let via_gemini = t
        .lookup("gemini", "gemini-2.0-flash")
        .expect("gemini aliases to the google price key");
    let via_google = t
        .lookup("google", "gemini-2.0-flash")
        .expect("google key exists");
    assert_eq!(
        via_gemini.cost_for(1000, 1000).total_usd,
        via_google.cost_for(1000, 1000).total_usd,
    );
}