kcode-intelligence-router 0.2.3

use serde::{Deserialize, Serialize};

use crate::{Metering, TokenUsage};

/// Published provider prices compiled into this Kennedy build.
pub const PRICING_VERSION: &str = "kennedy-provider-pricing-2026-07-30";

/// Confidence in a locally calculated provider cost.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum CostAccuracy {
    /// Provider usage maps completely to published rates.
    Exact,
    /// Provider usage omitted a billing detail and Kennedy used a documented fallback.
    Estimated,
    /// The estimate deliberately assumes a potentially-free usage item is chargeable.
    Conservative,
}

/// A non-negative estimated cost represented exactly in billionths of one US dollar.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CostEstimate {
    pub usd_nanos: u64,
    pub accuracy: CostAccuracy,
    pub pricing_version: String,
}

impl CostEstimate {
    pub(crate) fn versioned(mut self) -> Self {
        self.pricing_version = PRICING_VERSION.into();
        self
    }

    pub(crate) fn with_surcharge(mut self, usd_nanos: u64, accuracy: CostAccuracy) -> Self {
        self.usd_nanos = self.usd_nanos.saturating_add(usd_nanos);
        if accuracy == CostAccuracy::Conservative
            || accuracy == CostAccuracy::Estimated && self.accuracy == CostAccuracy::Exact
        {
            self.accuracy = accuracy;
        }
        self
    }
}

#[derive(Clone, Copy)]
struct TokenRates {
    input: u64,
    cached_input: u64,
    output: u64,
}

/// Estimates one normalized call using Kennedy's compiled standard-tier price catalog.
///
/// Thinking tokens are billed at the model's output rate. Unknown models and
/// unavailable provider metering intentionally return `None`.
pub fn estimate_cost(model: &str, metering: &Metering) -> Option<CostEstimate> {
    match metering {
        Metering::Tokens(usage) => estimate_token_cost(model, *usage),
        Metering::DurationSeconds { seconds }
            if model == kcode_openai_api::GPT_4O_TRANSCRIBE
                && seconds.is_finite()
                && *seconds >= 0.0 =>
        {
            // OpenAI publishes an estimated $0.006/minute, or $0.0001/second.
            Some(
                CostEstimate {
                    usd_nanos: (*seconds * 100_000.0).round() as u64,
                    accuracy: CostAccuracy::Estimated,
                    pricing_version: String::new(),
                }
                .versioned(),
            )
        }
        _ => None,
    }
}

/// Estimates normalized token usage for a known model.
pub fn estimate_token_cost(model: &str, usage: TokenUsage) -> Option<CostEstimate> {
    let total_input = usage.input_tokens.saturating_add(usage.cached_input_tokens);
    let rates = match model {
        // `gpt-5.6` is the direct-API alias for Sol.
        "gpt-5.6" | "gpt-5.6-sol" if total_input <= 272_000 => TokenRates {
            input: 5_000,
            cached_input: 500,
            output: 30_000,
        },
        "gpt-5.6" | "gpt-5.6-sol" => TokenRates {
            input: 10_000,
            cached_input: 1_000,
            output: 45_000,
        },
        "gpt-5.6-terra" if total_input <= 272_000 => TokenRates {
            input: 2_000,
            cached_input: 200,
            output: 12_000,
        },
        "gpt-5.6-terra" => TokenRates {
            input: 4_000,
            cached_input: 400,
            output: 18_000,
        },
        "gpt-5.6-luna" if total_input <= 272_000 => TokenRates {
            input: 200,
            cached_input: 20,
            output: 1_200,
        },
        "gpt-5.6-luna" => TokenRates {
            input: 400,
            cached_input: 40,
            output: 1_800,
        },
        "gpt-4o-transcribe" => TokenRates {
            input: 2_500,
            cached_input: 2_500,
            output: 10_000,
        },
        "gemini-2.5-flash" => TokenRates {
            input: 300,
            cached_input: 30,
            output: 2_500,
        },
        "gemini-3.1-flash-lite" => TokenRates {
            input: 250,
            cached_input: 25,
            output: 1_500,
        },
        "gemini-3.1-pro-preview" if total_input <= 200_000 => TokenRates {
            input: 2_000,
            cached_input: 200,
            output: 12_000,
        },
        "gemini-3.1-pro-preview" => TokenRates {
            input: 4_000,
            cached_input: 400,
            output: 18_000,
        },
        // Without modality detail, assume native image output for this image model.
        "gemini-3-pro-image" => TokenRates {
            input: 2_000,
            cached_input: 2_000,
            output: 120_000,
        },
        "gpt-image-2" => TokenRates {
            input: 8_000,
            cached_input: 2_000,
            output: 30_000,
        },
        _ => return None,
    };
    Some(
        CostEstimate {
            usd_nanos: usage
                .input_tokens
                .saturating_mul(rates.input)
                .saturating_add(usage.cached_input_tokens.saturating_mul(rates.cached_input))
                .saturating_add(
                    usage
                        .thinking_tokens
                        .saturating_add(usage.output_tokens)
                        .saturating_mul(rates.output),
                ),
            accuracy: if matches!(model, "gpt-image-2" | "gemini-3-pro-image") {
                CostAccuracy::Conservative
            } else {
                CostAccuracy::Exact
            },
            pricing_version: String::new(),
        }
        .versioned(),
    )
}

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

    #[test]
    fn openai_text_uses_distinct_cached_and_output_rates() {
        let cost = estimate_token_cost(
            "gpt-5.6-sol",
            TokenUsage {
                input_tokens: 1_000_000,
                cached_input_tokens: 1_000_000,
                thinking_tokens: 500_000,
                output_tokens: 500_000,
            },
        )
        .unwrap();
        // This crosses the long-context threshold: $10 + $1 + $45.
        assert_eq!(cost.usd_nanos, 56_000_000_000);
    }

    #[test]
    fn thinking_is_billed_at_output_rate() {
        let cost = estimate_token_cost(
            "gemini-3.1-flash-lite",
            TokenUsage {
                thinking_tokens: 10,
                output_tokens: 5,
                ..TokenUsage::default()
            },
        )
        .unwrap();
        assert_eq!(cost.usd_nanos, 22_500);
    }

    #[test]
    fn unknown_models_are_not_silently_free() {
        assert_eq!(
            estimate_token_cost("future-model", TokenUsage::default()),
            None
        );
    }
}