From eeddb8462a498d0b78e3339fa606cae09b07463c Mon Sep 17 00:00:00 2001 From: jimyag Date: Wed, 8 Apr 2026 13:54:34 +0800 Subject: [PATCH 1/4] feat: support external model and pricing configuration from TOML --- src/config.rs | 62 ++ src/main.rs | 3 + src/models.rs | 1605 +++++++++++++++++++++++++++---------------------- 3 files changed, 959 insertions(+), 711 deletions(-) diff --git a/src/config.rs b/src/config.rs index c2a9528..c399ea3 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,8 @@ +use crate::models::ModelInfo; use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; use std::cell::RefCell; +use std::collections::HashMap; use std::fs; use std::path::PathBuf; @@ -9,6 +11,10 @@ pub struct Config { pub server: ServerConfig, pub upload: UploadConfig, pub formatting: FormattingConfig, + #[serde(default)] + pub models: HashMap, + #[serde(default)] + pub aliases: HashMap, } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -52,6 +58,8 @@ impl Default for Config { locale: "en".to_string(), decimal_places: 2, }, + models: HashMap::new(), + aliases: HashMap::new(), } } } @@ -171,6 +179,12 @@ pub fn show_config() -> Result<()> { println!(" Number Human: {}", config.formatting.number_human); println!(" Locale: {}", config.formatting.locale); println!(" Decimal Places: {}", config.formatting.decimal_places); + if !config.models.is_empty() { + println!(" Custom Models: {}", config.models.len()); + } + if !config.aliases.is_empty() { + println!(" Custom Aliases: {}", config.aliases.len()); + } } None => { println!("❌ No configuration file found."); @@ -235,6 +249,54 @@ mod tests { (dir, config_path) } + #[test] + fn test_config_with_custom_models() { + let toml_str = r#" +[server] +url = "https://custom.example.com" +api_token = "test-token" + +[upload] +auto_upload = true +upload_today_only = false +retry_attempts = 5 +last_date_uploaded = 0 + +[formatting] +number_comma = true +number_human = false +locale = "zh" +decimal_places = 4 + +[models."custom-model"] +pricing = { Flat = { input_per_1m = 10.0, output_per_1m = 20.0 } } +caching = "None" +is_estimated = true + +[aliases] +"my-alias" = "custom-model" +"#; + + let config: Config = toml::from_str(toml_str).unwrap(); + + assert_eq!(config.server.url, "https://custom.example.com"); + assert!(config.models.contains_key("custom-model")); + + let custom_model = config.models.get("custom-model").unwrap(); + match &custom_model.pricing { + PricingStructure::Flat { + input_per_1m, + output_per_1m, + } => { + assert_eq!(*input_per_1m, 10.0); + assert_eq!(*output_per_1m, 20.0); + } + _ => panic!("Expected flat pricing"), + } + + assert_eq!(config.aliases.get("my-alias").unwrap(), "custom-model"); + } + #[test] fn default_config_round_trip() { let (_dir, _path) = setup_test_config(); diff --git a/src/main.rs b/src/main.rs index 71e2b06..c25057a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -131,6 +131,9 @@ async fn main() { // Load config file to get defaults let config = config::Config::load().unwrap_or(None).unwrap_or_default(); + // Initialize external models from config + models::init_external_models(config.models.clone(), config.aliases.clone()); + // Create format options merging config defaults with CLI overrides let format_options = utils::NumberFormatOptions { use_comma: cli.number_comma || config.formatting.number_comma, diff --git a/src/models.rs b/src/models.rs index fe0a9aa..cee3c83 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1,5 +1,7 @@ -use phf::phf_map; +use parking_lot::RwLock; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::OnceLock; use crate::utils::warn_once; @@ -14,16 +16,16 @@ pub struct PricingTier { pub output_per_1m: f64, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct TieredPricing { /// Pricing tiers ordered from lowest threshold to highest. - pub tiers: &'static [PricingTier], + pub tiers: Vec, /// If true, bill the entire token count at the single matching tier's rate. pub bracket_pricing: bool, } /// Different pricing structures supported by various model providers -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum PricingStructure { /// Flat rate pricing (same cost regardless of token count) Flat { @@ -43,16 +45,16 @@ pub struct CachingTier { pub cached_input_per_1m: f64, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct TieredCaching { /// Cache tiers ordered from lowest threshold to highest. - pub tiers: &'static [CachingTier], + pub tiers: Vec, /// If true, bill the entire token count at the single matching tier's rate. pub bracket_pricing: bool, } /// Different caching support models -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub enum CachingSupport { /// Model does not support caching None, @@ -68,7 +70,7 @@ pub enum CachingSupport { } /// Complete model information with all pricing details -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelInfo { /// Pricing structure (flat or tiered) pub pricing: PricingStructure, @@ -78,1063 +80,1204 @@ pub struct ModelInfo { pub is_estimated: bool, } -static MODEL_INDEX: phf::Map<&'static str, ModelInfo> = phf_map! { +/// Global registry for models and aliases +struct Registry { + index: HashMap, + aliases: HashMap, +} + +impl Registry { + fn new_with_defaults() -> Self { + let mut index = HashMap::new(); + let mut aliases = HashMap::new(); + populate_defaults(&mut index, &mut aliases); + Self { index, aliases } + } + + fn merge( + &mut self, + external_models: HashMap, + external_aliases: HashMap, + ) { + for (name, info) in external_models { + self.index.insert(name, info); + } + for (alias, canonical) in external_aliases { + self.aliases.insert(alias, canonical); + } + } +} + +static REGISTRY: OnceLock> = OnceLock::new(); +static EXTERNAL_MODELS_INITIALIZED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +/// Initialize the model registry with external configuration. +/// This should be called once at startup. +pub fn init_external_models( + external_models: HashMap, + external_aliases: HashMap, +) { + if EXTERNAL_MODELS_INITIALIZED.swap(true, std::sync::atomic::Ordering::SeqCst) { + warn_once( + "WARNING: init_external_models called multiple times. Skipping subsequent loads.", + ); + return; + } + + let rwlock = REGISTRY.get_or_init(|| RwLock::new(Registry::new_with_defaults())); + let mut registry = rwlock.write(); + registry.merge(external_models, external_aliases); +} + +fn get_registry_lock() -> &'static RwLock { + REGISTRY.get_or_init(|| RwLock::new(Registry::new_with_defaults())) +} + +fn populate_defaults( + index: &mut HashMap, + aliases: &mut HashMap, +) { + macro_rules! add_model { + ($name:expr, $pricing:expr, $caching:expr, $est:expr) => { + index.insert( + $name.to_string(), + ModelInfo { + pricing: $pricing, + caching: $caching, + is_estimated: $est, + }, + ); + }; + } + // OpenAI Models - "o4-mini" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "o4-mini", + PricingStructure::Flat { input_per_1m: 1.1, - output_per_1m: 4.4, + output_per_1m: 4.4 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.275, + CachingSupport::OpenAI { + cached_input_per_1m: 0.275 }, - is_estimated: false, - }, - "o3" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "o3", + PricingStructure::Flat { input_per_1m: 2.0, - output_per_1m: 8.0, + output_per_1m: 8.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.5, + CachingSupport::OpenAI { + cached_input_per_1m: 0.5 }, - is_estimated: false, - }, - "o3-pro" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "o3-pro", + PricingStructure::Flat { input_per_1m: 20.0, - output_per_1m: 80.0, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "o3-mini" => ModelInfo { - pricing: PricingStructure::Flat { + output_per_1m: 80.0 + }, + CachingSupport::None, + false + ); + add_model!( + "o3-mini", + PricingStructure::Flat { input_per_1m: 1.1, - output_per_1m: 4.4, + output_per_1m: 4.4 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.55, + CachingSupport::OpenAI { + cached_input_per_1m: 0.55 }, - is_estimated: false, - }, - "o1" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "o1", + PricingStructure::Flat { input_per_1m: 15.0, - output_per_1m: 60.0, + output_per_1m: 60.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 7.5, + CachingSupport::OpenAI { + cached_input_per_1m: 7.5 }, - is_estimated: false, - }, - "o1-preview" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "o1-preview", + PricingStructure::Flat { input_per_1m: 15.0, - output_per_1m: 60.0, + output_per_1m: 60.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 7.5, + CachingSupport::OpenAI { + cached_input_per_1m: 7.5 }, - is_estimated: false, - }, - "o1-mini" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "o1-mini", + PricingStructure::Flat { input_per_1m: 1.1, - output_per_1m: 4.4, + output_per_1m: 4.4 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.55, + CachingSupport::OpenAI { + cached_input_per_1m: 0.55 }, - is_estimated: false, - }, - "o1-pro" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "o1-pro", + PricingStructure::Flat { input_per_1m: 150.0, - output_per_1m: 600.0, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "gpt-4.1" => ModelInfo { - pricing: PricingStructure::Flat { + output_per_1m: 600.0 + }, + CachingSupport::None, + false + ); + add_model!( + "gpt-4.1", + PricingStructure::Flat { input_per_1m: 2.0, - output_per_1m: 8.0, + output_per_1m: 8.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.5, + CachingSupport::OpenAI { + cached_input_per_1m: 0.5 }, - is_estimated: false, - }, - "gpt-4o" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-4o", + PricingStructure::Flat { input_per_1m: 2.5, - output_per_1m: 10.0, + output_per_1m: 10.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 1.25, + CachingSupport::OpenAI { + cached_input_per_1m: 1.25 }, - is_estimated: false, - }, - "gpt-4o-2024-05-13" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-4o-2024-05-13", + PricingStructure::Flat { input_per_1m: 5.0, - output_per_1m: 10.0, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "gpt-4.1-mini" => ModelInfo { - pricing: PricingStructure::Flat { + output_per_1m: 10.0 + }, + CachingSupport::None, + false + ); + add_model!( + "gpt-4.1-mini", + PricingStructure::Flat { input_per_1m: 0.4, - output_per_1m: 1.6, + output_per_1m: 1.6 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.1, + CachingSupport::OpenAI { + cached_input_per_1m: 0.1 }, - is_estimated: false, - }, - "gpt-4.1-nano" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-4.1-nano", + PricingStructure::Flat { input_per_1m: 0.1, - output_per_1m: 0.4, + output_per_1m: 0.4 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.025, + CachingSupport::OpenAI { + cached_input_per_1m: 0.025 }, - is_estimated: false, - }, - "gpt-4o-mini" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-4o-mini", + PricingStructure::Flat { input_per_1m: 0.15, - output_per_1m: 0.6, + output_per_1m: 0.6 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.075, + CachingSupport::OpenAI { + cached_input_per_1m: 0.075 }, - is_estimated: false, - }, - "codex-mini-latest" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "codex-mini-latest", + PricingStructure::Flat { input_per_1m: 1.5, - output_per_1m: 6.0, + output_per_1m: 6.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.375, + CachingSupport::OpenAI { + cached_input_per_1m: 0.375 }, - is_estimated: false, - }, - "gpt-4-turbo" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-4-turbo", + PricingStructure::Flat { input_per_1m: 10.0, - output_per_1m: 30.0, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "gpt-5" => ModelInfo { - pricing: PricingStructure::Flat { + output_per_1m: 30.0 + }, + CachingSupport::None, + false + ); + add_model!( + "gpt-5", + PricingStructure::Flat { input_per_1m: 1.25, - output_per_1m: 10.0, + output_per_1m: 10.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.125, + CachingSupport::OpenAI { + cached_input_per_1m: 0.125 }, - is_estimated: false, - }, - "gpt-5.1" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.1", + PricingStructure::Flat { input_per_1m: 1.25, - output_per_1m: 10.0, + output_per_1m: 10.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.125, + CachingSupport::OpenAI { + cached_input_per_1m: 0.125 }, - is_estimated: false, - }, - "gpt-5-mini" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5-mini", + PricingStructure::Flat { input_per_1m: 0.25, - output_per_1m: 2.0, + output_per_1m: 2.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.025, + CachingSupport::OpenAI { + cached_input_per_1m: 0.025 }, - is_estimated: false, - }, - "gpt-5-nano" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5-nano", + PricingStructure::Flat { input_per_1m: 0.05, - output_per_1m: 0.4, + output_per_1m: 0.4 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.005, + CachingSupport::OpenAI { + cached_input_per_1m: 0.005 }, - is_estimated: false, - }, - "gpt-5-codex-mini" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5-codex-mini", + PricingStructure::Flat { input_per_1m: 0.25, - output_per_1m: 2.0, + output_per_1m: 2.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.025, + CachingSupport::OpenAI { + cached_input_per_1m: 0.025 }, - is_estimated: false, - }, - // GPT-5.1 Codex models - "gpt-5.1-codex" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.1-codex", + PricingStructure::Flat { input_per_1m: 1.25, - output_per_1m: 10.0, + output_per_1m: 10.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.125, + CachingSupport::OpenAI { + cached_input_per_1m: 0.125 }, - is_estimated: false, - }, - "gpt-5.1-codex-mini" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.1-codex-mini", + PricingStructure::Flat { input_per_1m: 0.25, - output_per_1m: 2.0, + output_per_1m: 2.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.025, + CachingSupport::OpenAI { + cached_input_per_1m: 0.025 }, - is_estimated: false, - }, - "gpt-5.1-codex-max" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.1-codex-max", + PricingStructure::Flat { input_per_1m: 1.25, - output_per_1m: 10.0, + output_per_1m: 10.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.125, + CachingSupport::OpenAI { + cached_input_per_1m: 0.125 }, - is_estimated: false, - }, - "gpt-5.2" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.2", + PricingStructure::Flat { input_per_1m: 1.75, - output_per_1m: 14.0, + output_per_1m: 14.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.175, + CachingSupport::OpenAI { + cached_input_per_1m: 0.175 }, - is_estimated: false, - }, - "gpt-5.2-pro" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.2-pro", + PricingStructure::Flat { input_per_1m: 21.0, - output_per_1m: 168.0, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "gpt-5.2-codex" => ModelInfo { - pricing: PricingStructure::Flat { + output_per_1m: 168.0 + }, + CachingSupport::None, + false + ); + add_model!( + "gpt-5.2-codex", + PricingStructure::Flat { input_per_1m: 1.75, - output_per_1m: 14.0, + output_per_1m: 14.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.175, + CachingSupport::OpenAI { + cached_input_per_1m: 0.175 }, - is_estimated: false, - }, - // GPT-5.3 Codex - "gpt-5.3-codex" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5.3-codex", + PricingStructure::Flat { input_per_1m: 1.75, - output_per_1m: 14.0, + output_per_1m: 14.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.175, + CachingSupport::OpenAI { + cached_input_per_1m: 0.175 }, - is_estimated: false, - }, - "gpt-5-pro" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gpt-5-pro", + PricingStructure::Flat { input_per_1m: 15.0, - output_per_1m: 120.0, + output_per_1m: 120.0 }, - caching: CachingSupport::None, - is_estimated: false, - }, - "gpt-5.4" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + CachingSupport::None, + false + ); + + add_model!( + "gpt-5.4", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(272_000), input_per_1m: 2.50, - output_per_1m: 15.0, + output_per_1m: 15.0 }, PricingTier { max_tokens: None, input_per_1m: 5.0, - output_per_1m: 22.5, + output_per_1m: 22.5 }, ], bracket_pricing: false, }), - caching: CachingSupport::Google(TieredCaching { - tiers: &[ + CachingSupport::Google(TieredCaching { + tiers: vec![ CachingTier { max_tokens: Some(272_000), - cached_input_per_1m: 0.25, + cached_input_per_1m: 0.25 }, CachingTier { max_tokens: None, - cached_input_per_1m: 0.50, + cached_input_per_1m: 0.50 }, ], bracket_pricing: false, }), - is_estimated: false, - }, - "gpt-5.4-mini" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + + add_model!( + "gpt-5.4-mini", + PricingStructure::Flat { input_per_1m: 0.75, - output_per_1m: 4.5, + output_per_1m: 4.5 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.075, + CachingSupport::OpenAI { + cached_input_per_1m: 0.075 }, - is_estimated: false, - }, + false + ); // Anthropic Models - "claude-opus-4-6" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "claude-opus-4-6", + PricingStructure::Flat { input_per_1m: 5.0, - output_per_1m: 25.0, + output_per_1m: 25.0 }, - caching: CachingSupport::Anthropic { - cache_write_per_1m: 6.25, // 1.25x base input - cache_read_per_1m: 0.5, // 0.1x base input + CachingSupport::Anthropic { + cache_write_per_1m: 6.25, + cache_read_per_1m: 0.5 }, - is_estimated: false, - }, - "claude-opus-4-5" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-opus-4-5", + PricingStructure::Flat { input_per_1m: 5.0, - output_per_1m: 25.0, + output_per_1m: 25.0 }, - caching: CachingSupport::Anthropic { - cache_write_per_1m: 6.25, // 1.25x base input - cache_read_per_1m: 0.5, // 0.1x base input + CachingSupport::Anthropic { + cache_write_per_1m: 6.25, + cache_read_per_1m: 0.5 }, - is_estimated: false, - }, - "claude-opus-4-1" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-opus-4-1", + PricingStructure::Flat { input_per_1m: 15.0, - output_per_1m: 75.0, + output_per_1m: 75.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 18.75, - cache_read_per_1m: 1.5, + cache_read_per_1m: 1.5 }, - is_estimated: false, - }, - "claude-opus-4" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-opus-4", + PricingStructure::Flat { input_per_1m: 15.0, - output_per_1m: 75.0, + output_per_1m: 75.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 18.75, - cache_read_per_1m: 1.5, + cache_read_per_1m: 1.5 }, - is_estimated: false, - }, - "claude-sonnet-4" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-sonnet-4", + PricingStructure::Flat { input_per_1m: 3.0, - output_per_1m: 15.0, + output_per_1m: 15.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 3.75, - cache_read_per_1m: 0.3, + cache_read_per_1m: 0.3 }, - is_estimated: false, - }, - "claude-sonnet-4-6" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-sonnet-4-6", + PricingStructure::Flat { input_per_1m: 3.0, - output_per_1m: 15.0, + output_per_1m: 15.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 3.75, - cache_read_per_1m: 0.3, + cache_read_per_1m: 0.3 }, - is_estimated: false, - }, - "claude-sonnet-4-5" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-sonnet-4-5", + PricingStructure::Flat { input_per_1m: 3.0, - output_per_1m: 15.0, + output_per_1m: 15.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 3.75, - cache_read_per_1m: 0.3, + cache_read_per_1m: 0.3 }, - is_estimated: false, - }, - "claude-3-7-sonnet" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-3-7-sonnet", + PricingStructure::Flat { input_per_1m: 3.0, - output_per_1m: 15.0, + output_per_1m: 15.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 3.75, - cache_read_per_1m: 0.3, + cache_read_per_1m: 0.3 }, - is_estimated: false, - }, - "claude-3-5-sonnet" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-3-5-sonnet", + PricingStructure::Flat { input_per_1m: 3.0, - output_per_1m: 15.0, + output_per_1m: 15.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 3.75, - cache_read_per_1m: 0.3, + cache_read_per_1m: 0.3 }, - is_estimated: false, - }, - "claude-3-5-haiku" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-3-5-haiku", + PricingStructure::Flat { input_per_1m: 0.8, - output_per_1m: 4.0, + output_per_1m: 4.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 1.0, - cache_read_per_1m: 0.08, + cache_read_per_1m: 0.08 }, - is_estimated: false, - }, - "claude-haiku-4-5" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-haiku-4-5", + PricingStructure::Flat { input_per_1m: 1.0, - output_per_1m: 5.0, + output_per_1m: 5.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 1.25, - cache_read_per_1m: 0.10, + cache_read_per_1m: 0.10 }, - is_estimated: false, - }, - "claude-3-opus" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-3-opus", + PricingStructure::Flat { input_per_1m: 15.0, - output_per_1m: 75.0, + output_per_1m: 75.0 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 18.75, - cache_read_per_1m: 1.5, + cache_read_per_1m: 1.5 }, - is_estimated: false, - }, - "claude-3-haiku" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "claude-3-haiku", + PricingStructure::Flat { input_per_1m: 0.25, - output_per_1m: 1.25, + output_per_1m: 1.25 }, - caching: CachingSupport::Anthropic { + CachingSupport::Anthropic { cache_write_per_1m: 0.3, - cache_read_per_1m: 0.03, + cache_read_per_1m: 0.03 }, - is_estimated: false, - }, + false + ); // Google Models - "gemini-3-flash-preview" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "gemini-3-flash-preview", + PricingStructure::Flat { input_per_1m: 0.5, - output_per_1m: 3.0, + output_per_1m: 3.0 }, - caching: CachingSupport::Google(TieredCaching { - tiers: &[ - CachingTier { - max_tokens: None, - cached_input_per_1m: 0.05, - }, - ], + CachingSupport::Google(TieredCaching { + tiers: vec![CachingTier { + max_tokens: None, + cached_input_per_1m: 0.05 + }], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-3.1-pro-preview" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + false + ); + add_model!( + "gemini-3.1-pro-preview", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(200_000), input_per_1m: 2.0, - output_per_1m: 12.0, + output_per_1m: 12.0 }, PricingTier { max_tokens: None, input_per_1m: 4.0, - output_per_1m: 18.0, + output_per_1m: 18.0 }, ], bracket_pricing: true, }), - caching: CachingSupport::Google(TieredCaching { - tiers: &[ + CachingSupport::Google(TieredCaching { + tiers: vec![ CachingTier { max_tokens: Some(200_000), - cached_input_per_1m: 0.20, + cached_input_per_1m: 0.20 }, CachingTier { max_tokens: None, - cached_input_per_1m: 0.40, + cached_input_per_1m: 0.40 }, ], bracket_pricing: true, }), - is_estimated: false, - }, - "gemini-3-pro-preview-11-2025" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + false + ); + add_model!( + "gemini-3-pro-preview-11-2025", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(200_000), input_per_1m: 2.0, - output_per_1m: 12.0, + output_per_1m: 12.0 }, PricingTier { max_tokens: None, input_per_1m: 4.0, - output_per_1m: 18.0, + output_per_1m: 18.0 }, ], bracket_pricing: false, }), - caching: CachingSupport::None, - is_estimated: false, - }, - "gemini-2.5-pro" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + CachingSupport::None, + false + ); + add_model!( + "gemini-2.5-pro", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(200_000), input_per_1m: 1.25, - output_per_1m: 10.0, + output_per_1m: 10.0 }, PricingTier { max_tokens: None, input_per_1m: 2.5, - output_per_1m: 15.0, + output_per_1m: 15.0 }, ], bracket_pricing: false, }), - caching: CachingSupport::Google(TieredCaching { - tiers: &[ + CachingSupport::Google(TieredCaching { + tiers: vec![ CachingTier { max_tokens: Some(200_000), - cached_input_per_1m: 0.31, + cached_input_per_1m: 0.31 }, CachingTier { max_tokens: None, - cached_input_per_1m: 0.625, + cached_input_per_1m: 0.625 }, ], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-2.5-flash" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gemini-2.5-flash", + PricingStructure::Flat { input_per_1m: 0.3, - output_per_1m: 2.5, + output_per_1m: 2.5 }, - caching: CachingSupport::Google(TieredCaching { - tiers: &[ - CachingTier { - max_tokens: None, - cached_input_per_1m: 0.075, - }, - ], + CachingSupport::Google(TieredCaching { + tiers: vec![CachingTier { + max_tokens: None, + cached_input_per_1m: 0.075 + }], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-2.5-flash-lite" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gemini-2.5-flash-lite", + PricingStructure::Flat { input_per_1m: 0.1, - output_per_1m: 0.4, + output_per_1m: 0.4 }, - caching: CachingSupport::Google(TieredCaching { - tiers: &[ - CachingTier { - max_tokens: None, - cached_input_per_1m: 0.025, - }, - ], + CachingSupport::Google(TieredCaching { + tiers: vec![CachingTier { + max_tokens: None, + cached_input_per_1m: 0.025 + }], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-2.0-pro-exp-02-05" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gemini-2.0-pro-exp-02-05", + PricingStructure::Flat { input_per_1m: 0.0, - output_per_1m: 0.0, + output_per_1m: 0.0 }, - caching: CachingSupport::Google(TieredCaching { - tiers: &[ - CachingTier { - max_tokens: None, - cached_input_per_1m: 0.0, - }, - ], + CachingSupport::Google(TieredCaching { + tiers: vec![CachingTier { + max_tokens: None, + cached_input_per_1m: 0.0 + }], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-2.0-flash" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gemini-2.0-flash", + PricingStructure::Flat { input_per_1m: 0.1, - output_per_1m: 0.4, + output_per_1m: 0.4 }, - caching: CachingSupport::Google(TieredCaching { - tiers: &[ - CachingTier { - max_tokens: None, - cached_input_per_1m: 0.025, - }, - ], + CachingSupport::Google(TieredCaching { + tiers: vec![CachingTier { + max_tokens: None, + cached_input_per_1m: 0.025 + }], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-2.0-flash-lite" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "gemini-2.0-flash-lite", + PricingStructure::Flat { input_per_1m: 0.075, - output_per_1m: 0.3, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "gemini-1.5-flash" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + output_per_1m: 0.3 + }, + CachingSupport::None, + false + ); + add_model!( + "gemini-1.5-flash", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(128_000), input_per_1m: 0.075, - output_per_1m: 0.3, + output_per_1m: 0.3 }, PricingTier { max_tokens: None, input_per_1m: 0.15, - output_per_1m: 0.6, + output_per_1m: 0.6 }, ], bracket_pricing: false, }), - caching: CachingSupport::Google(TieredCaching { - tiers: &[ + CachingSupport::Google(TieredCaching { + tiers: vec![ CachingTier { max_tokens: Some(128_000), - cached_input_per_1m: 0.01875, + cached_input_per_1m: 0.01875 }, CachingTier { max_tokens: None, - cached_input_per_1m: 0.0375, + cached_input_per_1m: 0.0375 }, ], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-1.5-flash-8b" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + false + ); + add_model!( + "gemini-1.5-flash-8b", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(128_000), input_per_1m: 0.0375, - output_per_1m: 0.15, + output_per_1m: 0.15 }, PricingTier { max_tokens: None, input_per_1m: 0.075, - output_per_1m: 0.3, + output_per_1m: 0.3 }, ], bracket_pricing: false, }), - caching: CachingSupport::Google(TieredCaching { - tiers: &[ + CachingSupport::Google(TieredCaching { + tiers: vec![ CachingTier { max_tokens: Some(128_000), - cached_input_per_1m: 0.01, + cached_input_per_1m: 0.01 }, CachingTier { max_tokens: None, - cached_input_per_1m: 0.02, + cached_input_per_1m: 0.02 }, ], bracket_pricing: false, }), - is_estimated: false, - }, - "gemini-1.5-pro" => ModelInfo { - pricing: PricingStructure::Tiered(TieredPricing { - tiers: &[ + false + ); + add_model!( + "gemini-1.5-pro", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ PricingTier { max_tokens: Some(128_000), input_per_1m: 1.25, - output_per_1m: 5.0, + output_per_1m: 5.0 }, PricingTier { max_tokens: None, input_per_1m: 2.5, - output_per_1m: 10.0, + output_per_1m: 10.0 }, ], bracket_pricing: false, }), - caching: CachingSupport::Google(TieredCaching { - tiers: &[ + CachingSupport::Google(TieredCaching { + tiers: vec![ CachingTier { max_tokens: Some(128_000), - cached_input_per_1m: 0.3125, + cached_input_per_1m: 0.3125 }, CachingTier { max_tokens: None, - cached_input_per_1m: 0.625, + cached_input_per_1m: 0.625 }, ], bracket_pricing: false, }), - is_estimated: false, - }, + false + ); // Z.AI (Zhipu AI) Models - "glm-4.6" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "glm-4.6", + PricingStructure::Flat { input_per_1m: 0.60, - output_per_1m: 2.20, + output_per_1m: 2.20 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.11, + CachingSupport::OpenAI { + cached_input_per_1m: 0.11 }, - is_estimated: false, - }, - "glm-4.7" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "glm-4.7", + PricingStructure::Flat { input_per_1m: 0.60, - output_per_1m: 2.20, + output_per_1m: 2.20 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.11, + CachingSupport::OpenAI { + cached_input_per_1m: 0.11 }, - is_estimated: false, - }, - "glm-4.7-flash" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "glm-4.7-flash", + PricingStructure::Flat { input_per_1m: 0.0, - output_per_1m: 0.0, + output_per_1m: 0.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.0, + CachingSupport::OpenAI { + cached_input_per_1m: 0.0 }, - is_estimated: false, - }, - "glm-4.6v" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "glm-4.6v", + PricingStructure::Flat { input_per_1m: 0.30, - output_per_1m: 0.90, + output_per_1m: 0.90 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.05, + CachingSupport::OpenAI { + cached_input_per_1m: 0.05 }, - is_estimated: false, - }, + false + ); // xAI Models - "grok-code-fast-1" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "grok-code-fast-1", + PricingStructure::Flat { input_per_1m: 0.20, - output_per_1m: 1.50, + output_per_1m: 1.50 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.02, + CachingSupport::OpenAI { + cached_input_per_1m: 0.02 }, - is_estimated: false, - }, + false + ); // Synthetic.new Models - "hf:zai-org/GLM-4.6" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "hf:zai-org/GLM-4.6", + PricingStructure::Flat { input_per_1m: 0.55, - output_per_1m: 2.19, - }, - caching: CachingSupport::None, - is_estimated: false, - }, - "hf:MiniMaxAI/MiniMax-M2" => ModelInfo { - pricing: PricingStructure::Flat { + output_per_1m: 2.19 + }, + CachingSupport::None, + false + ); + add_model!( + "hf:MiniMaxAI/MiniMax-M2", + PricingStructure::Flat { input_per_1m: 0.55, - output_per_1m: 2.19, + output_per_1m: 2.19 }, - caching: CachingSupport::None, - is_estimated: false, - }, + CachingSupport::None, + false + ); // Z.AI (Zhipu AI) - Additional Models - "glm-5" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "glm-5", + PricingStructure::Flat { input_per_1m: 1.0, - output_per_1m: 3.2, + output_per_1m: 3.2 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.2, + CachingSupport::OpenAI { + cached_input_per_1m: 0.2 }, - is_estimated: false, - }, - "glm-5-code" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "glm-5-code", + PricingStructure::Flat { input_per_1m: 1.2, - output_per_1m: 5.0, + output_per_1m: 5.0 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.3, + CachingSupport::OpenAI { + cached_input_per_1m: 0.3 }, - is_estimated: false, - }, - "glm-4.5-air" => ModelInfo { - pricing: PricingStructure::Flat { + false + ); + add_model!( + "glm-4.5-air", + PricingStructure::Flat { input_per_1m: 0.2, - output_per_1m: 1.1, + output_per_1m: 1.1 }, - caching: CachingSupport::OpenAI { - cached_input_per_1m: 0.03, + CachingSupport::OpenAI { + cached_input_per_1m: 0.03 }, - is_estimated: false, - }, + false + ); // MiniMax Models - "minimax-m2.5" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "minimax-m2.5", + PricingStructure::Flat { input_per_1m: 0.30, - output_per_1m: 1.10, + output_per_1m: 1.10 }, - caching: CachingSupport::None, - is_estimated: false, - }, + CachingSupport::None, + false + ); // StepFun Models - "step-3.5-flash" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "step-3.5-flash", + PricingStructure::Flat { input_per_1m: 0.10, - output_per_1m: 0.30, + output_per_1m: 0.30 }, - caching: CachingSupport::None, - is_estimated: false, - }, + CachingSupport::None, + false + ); // Upstage Models - "solar-pro-3" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "solar-pro-3", + PricingStructure::Flat { input_per_1m: 0.15, - output_per_1m: 0.60, + output_per_1m: 0.60 }, - caching: CachingSupport::None, - is_estimated: false, - }, + CachingSupport::None, + false + ); // OpenRouter Models - "aurora-alpha" => ModelInfo { - pricing: PricingStructure::Flat { + add_model!( + "aurora-alpha", + PricingStructure::Flat { input_per_1m: 0.0, - output_per_1m: 0.0, + output_per_1m: 0.0 }, - caching: CachingSupport::None, - is_estimated: false, - }, -}; + CachingSupport::None, + false + ); + + // Populate Aliases + macro_rules! add_alias { + ($alias:expr, $canonical:expr) => { + if $alias != $canonical { + aliases.insert($alias.to_string(), $canonical.to_string()); + } + }; + } -static MODEL_ALIASES: phf::Map<&'static str, &'static str> = phf_map! { // OpenAI aliases - "o4-mini" => "o4-mini", - "o4-mini-2025-04-16" => "o4-mini", - "o3" => "o3", - "o3-2025-04-16" => "o3", - "o3-pro" => "o3-pro", - "o3-pro-2025-06-10" => "o3-pro", - "o3-mini" => "o3-mini", - "o3-mini-2025-01-31" => "o3-mini", - "o1" => "o1", - "o1-2024-12-17" => "o1", - "o1-preview" => "o1-preview", - "o1-preview-2024-09-12" => "o1-preview", - "o1-mini" => "o1-mini", - "o1-mini-2024-09-12" => "o1-mini", - "o1-pro" => "o1-pro", - "o1-pro-2025-03-19" => "o1-pro", - "gpt-4.1" => "gpt-4.1", - "gpt-4.1-2025-04-14" => "gpt-4.1", - "gpt-4o" => "gpt-4o", - "gpt-4o-2024-11-20" => "gpt-4o", - "gpt-4o-2024-08-06" => "gpt-4o", - "gpt-4o-2024-05-13" => "gpt-4o-2024-05-13", - "gpt-4.1-mini" => "gpt-4.1-mini", - "gpt-4.1-mini-2025-04-14" => "gpt-4.1-mini", - "gpt-4.1-nano" => "gpt-4.1-nano", - "gpt-4.1-nano-2025-04-14" => "gpt-4.1-nano", - "gpt-4o-mini" => "gpt-4o-mini", - "gpt-4o-mini-2024-07-18" => "gpt-4o-mini", - "codex-mini-latest" => "codex-mini-latest", - "gpt-4-turbo" => "gpt-4-turbo", - "gpt-4-turbo-2024-04-09" => "gpt-4-turbo", - "gpt-5" => "gpt-5", - "gpt-5-codex" => "gpt-5", - "gpt-5-2025-08-07" => "gpt-5", - "gpt-5.1" => "gpt-5.1", - "gpt-5.1-2025-08-07" => "gpt-5.1", - "gpt-5-mini" => "gpt-5-mini", - "gpt-5-mini-2025-08-07" => "gpt-5-mini", - "gpt-5-nano" => "gpt-5-nano", - "gpt-5-nano-2025-08-07" => "gpt-5-nano", - "gpt-5-codex-mini" => "gpt-5-codex-mini", - "gpt-5.1-codex" => "gpt-5.1-codex", - "gpt-5.1-codex-mini" => "gpt-5.1-codex-mini", - "gpt-5.1-codex-max" => "gpt-5.1-codex-max", - "gpt-5.2" => "gpt-5.2", - "gpt-5.2-2025-12-11" => "gpt-5.2", - "gpt-5.2-pro" => "gpt-5.2-pro", - "gpt-5.2-codex" => "gpt-5.2-codex", - "gpt-5.3-codex" => "gpt-5.3-codex", - "gpt-5-pro" => "gpt-5-pro", + add_alias!("o4-mini", "o4-mini"); + add_alias!("o4-mini-2025-04-16", "o4-mini"); + add_alias!("o3", "o3"); + add_alias!("o3-2025-04-16", "o3"); + add_alias!("o3-pro", "o3-pro"); + add_alias!("o3-pro-2025-06-10", "o3-pro"); + add_alias!("o3-mini", "o3-mini"); + add_alias!("o3-mini-2025-01-31", "o3-mini"); + add_alias!("o1", "o1"); + add_alias!("o1-2024-12-17", "o1"); + add_alias!("o1-preview", "o1-preview"); + add_alias!("o1-preview-2024-09-12", "o1-preview"); + add_alias!("o1-mini", "o1-mini"); + add_alias!("o1-mini-2024-09-12", "o1-mini"); + add_alias!("o1-pro", "o1-pro"); + add_alias!("o1-pro-2025-03-19", "o1-pro"); + add_alias!("gpt-4.1", "gpt-4.1"); + add_alias!("gpt-4.1-2025-04-14", "gpt-4.1"); + add_alias!("gpt-4o", "gpt-4o"); + add_alias!("gpt-4o-2024-11-20", "gpt-4o"); + add_alias!("gpt-4o-2024-08-06", "gpt-4o"); + add_alias!("gpt-4o-2024-05-13", "gpt-4o-2024-05-13"); + add_alias!("gpt-4.1-mini", "gpt-4.1-mini"); + add_alias!("gpt-4.1-mini-2025-04-14", "gpt-4.1-mini"); + add_alias!("gpt-4.1-nano", "gpt-4.1-nano"); + add_alias!("gpt-4.1-nano-2025-04-14", "gpt-4.1-nano"); + add_alias!("gpt-4o-mini", "gpt-4o-mini"); + add_alias!("gpt-4o-mini-2024-07-18", "gpt-4o-mini"); + add_alias!("codex-mini-latest", "codex-mini-latest"); + add_alias!("gpt-4-turbo", "gpt-4-turbo"); + add_alias!("gpt-4-turbo-2024-04-09", "gpt-4-turbo"); + add_alias!("gpt-5", "gpt-5"); + add_alias!("gpt-5-codex", "gpt-5"); + add_alias!("gpt-5-2025-08-07", "gpt-5"); + add_alias!("gpt-5.1", "gpt-5.1"); + add_alias!("gpt-5.1-2025-08-07", "gpt-5.1"); + add_alias!("gpt-5-mini", "gpt-5-mini"); + add_alias!("gpt-5-mini-2025-08-07", "gpt-5-mini"); + add_alias!("gpt-5-nano", "gpt-5-nano"); + add_alias!("gpt-5-nano-2025-08-07", "gpt-5-nano"); + add_alias!("gpt-5-codex-mini", "gpt-5-codex-mini"); + add_alias!("gpt-5.1-codex", "gpt-5.1-codex"); + add_alias!("gpt-5.1-codex-mini", "gpt-5.1-codex-mini"); + add_alias!("gpt-5.1-codex-max", "gpt-5.1-codex-max"); + add_alias!("gpt-5.2", "gpt-5.2"); + add_alias!("gpt-5.2-2025-12-11", "gpt-5.2"); + add_alias!("gpt-5.2-pro", "gpt-5.2-pro"); + add_alias!("gpt-5.2-codex", "gpt-5.2-codex"); + add_alias!("gpt-5.3-codex", "gpt-5.3-codex"); + add_alias!("gpt-5-pro", "gpt-5-pro"); // Anthropic aliases - "claude-opus-4-6" => "claude-opus-4-6", - "claude-opus-4-5" => "claude-opus-4-5", - "claude-opus-4.5" => "claude-opus-4-5", - "claude-opus-4-5-20251101" => "claude-opus-4-5", - "claude-opus-4" => "claude-opus-4", - "claude-opus-4-20250514" => "claude-opus-4", - "claude-opus-4-0" => "claude-opus-4", - "claude-opus-4.1" => "claude-opus-4-1", - "claude-opus-4-1-20250805" => "claude-opus-4-1", - "claude-sonnet-4" => "claude-sonnet-4", - "claude-sonnet-4-20250514" => "claude-sonnet-4", - "claude-sonnet-4-0" => "claude-sonnet-4", - "claude-sonnet-4.6" => "claude-sonnet-4-6", - "claude-sonnet-4-6" => "claude-sonnet-4-6", - "claude-sonnet-4.5" => "claude-sonnet-4-5", - "claude-sonnet-4-5" => "claude-sonnet-4-5", - "claude-sonnet-4-5-20250929" => "claude-sonnet-4-5", - "claude-3-7-sonnet" => "claude-3-7-sonnet", - "claude-3-7-sonnet-20250219" => "claude-3-7-sonnet", - "claude-3-7-sonnet-latest" => "claude-3-7-sonnet", - "claude-3-5-sonnet" => "claude-3-5-sonnet", - "claude-3-5-sonnet-20241022" => "claude-3-5-sonnet", - "claude-3-5-sonnet-latest" => "claude-3-5-sonnet", - "claude-3-5-sonnet-20240620" => "claude-3-5-sonnet", - "claude-3-5-haiku" => "claude-3-5-haiku", - "claude-3-5-haiku-20241022" => "claude-3-5-haiku", - "claude-3-5-haiku-latest" => "claude-3-5-haiku", - "claude-haiku-4-5" => "claude-haiku-4-5", - "claude-haiku-4.5" => "claude-haiku-4-5", - "claude-haiku-4-5-20251001" => "claude-haiku-4-5", - "claude-3-opus" => "claude-3-opus", - "claude-3-opus-20240229" => "claude-3-opus", - "claude-3-haiku" => "claude-3-haiku", - "claude-3-haiku-20240307" => "claude-3-haiku", + add_alias!("claude-opus-4-6", "claude-opus-4-6"); + add_alias!("claude-opus-4-5", "claude-opus-4-5"); + add_alias!("claude-opus-4.5", "claude-opus-4-5"); + add_alias!("claude-opus-4-5-20251101", "claude-opus-4-5"); + add_alias!("claude-opus-4", "claude-opus-4"); + add_alias!("claude-opus-4-20250514", "claude-opus-4"); + add_alias!("claude-opus-4-0", "claude-opus-4"); + add_alias!("claude-opus-4.1", "claude-opus-4-1"); + add_alias!("claude-opus-4-1-20250805", "claude-opus-4-1"); + add_alias!("claude-sonnet-4", "claude-sonnet-4"); + add_alias!("claude-sonnet-4-20250514", "claude-sonnet-4"); + add_alias!("claude-sonnet-4-0", "claude-sonnet-4"); + add_alias!("claude-sonnet-4.6", "claude-sonnet-4-6"); + add_alias!("claude-sonnet-4.5", "claude-sonnet-4-5"); + add_alias!("claude-sonnet-4-5-20250929", "claude-sonnet-4-5"); + add_alias!("claude-3-7-sonnet", "claude-3-7-sonnet"); + add_alias!("claude-3-7-sonnet-20250219", "claude-3-7-sonnet"); + add_alias!("claude-3-7-sonnet-latest", "claude-3-7-sonnet"); + add_alias!("claude-3-5-sonnet", "claude-3-5-sonnet"); + add_alias!("claude-3-5-sonnet-20241022", "claude-3-5-sonnet"); + add_alias!("claude-3-5-sonnet-latest", "claude-3-5-sonnet"); + add_alias!("claude-3-5-sonnet-20240620", "claude-3-5-sonnet"); + add_alias!("claude-3-5-haiku", "claude-3-5-haiku"); + add_alias!("claude-3-5-haiku-20241022", "claude-3-5-haiku"); + add_alias!("claude-3-5-haiku-latest", "claude-3-5-haiku"); + add_alias!("claude-haiku-4-5", "claude-haiku-4-5"); + add_alias!("claude-haiku-4.5", "claude-haiku-4-5"); + add_alias!("claude-haiku-4-5-20251001", "claude-haiku-4-5"); + add_alias!("claude-3-opus", "claude-3-opus"); + add_alias!("claude-3-opus-20240229", "claude-3-opus"); + add_alias!("claude-3-haiku", "claude-3-haiku"); + add_alias!("claude-3-haiku-20240307", "claude-3-haiku"); // Google aliases - "gemini-3-flash-preview" => "gemini-3-flash-preview", - "gemini-3-flash-preview-12-2025" => "gemini-3-flash-preview", - "gemini-3-flash" => "gemini-3-flash-preview", - "gemini-3.1-pro-preview" => "gemini-3.1-pro-preview", - "gemini-3.1-pro" => "gemini-3.1-pro-preview", - "gemini-3.1-pro-low" => "gemini-3.1-pro-preview", - "gemini-3.1-pro-medium" => "gemini-3.1-pro-preview", - "gemini-3.1-pro-high" => "gemini-3.1-pro-preview", - "gemini-3-pro-preview-11-2025" => "gemini-3-pro-preview-11-2025", - "gemini-3-pro-preview" => "gemini-3-pro-preview-11-2025", - "gemini-3-pro" => "gemini-3-pro-preview-11-2025", - "gemini-2.5-pro" => "gemini-2.5-pro", - "gemini-2.5-pro-preview-06-05" => "gemini-2.5-pro", - "gemini-2.5-pro-preview-05-06" => "gemini-2.5-pro", - "gemini-2.5-pro-preview-03-25" => "gemini-2.5-pro", - "gemini-2.5-flash" => "gemini-2.5-flash", - "gemini-2.5-flash-preview-05-20" => "gemini-2.5-flash", - "gemini-2.5-flash-preview-04-17" => "gemini-2.5-flash", - "gemini-2.5-flash-lite" => "gemini-2.5-flash-lite", - "gemini-2.5-flash-lite-06-17" => "gemini-2.5-flash-lite", - "gemini-2.0-pro-exp-02-05" => "gemini-2.0-pro-exp-02-05", - "gemini-exp-1206" => "gemini-2.0-pro-exp-02-05", - "gemini-2.0-flash" => "gemini-2.0-flash", - "gemini-2.0-flash-001" => "gemini-2.0-flash", - "gemini-2.0-flash-exp" => "gemini-2.0-flash", - "gemini-2.0-flash-lite" => "gemini-2.0-flash-lite", - "gemini-2.0-flash-lite-001" => "gemini-2.0-flash-lite", - "gemini-1.5-flash" => "gemini-1.5-flash", - "gemini-1.5-flash-latest" => "gemini-1.5-flash", - "gemini-1.5-flash-001" => "gemini-1.5-flash", - "gemini-1.5-flash-002" => "gemini-1.5-flash", - "gemini-1.5-flash-8b" => "gemini-1.5-flash-8b", - "gemini-1.5-flash-8b-latest" => "gemini-1.5-flash-8b", - "gemini-1.5-flash-8b-001" => "gemini-1.5-flash-8b", - "gemini-1.5-flash-8b-exp-0924" => "gemini-1.5-flash-8b", - "gemini-1.5-flash-8b-exp-0827" => "gemini-1.5-flash-8b", - "gemini-1.5-pro" => "gemini-1.5-pro", - "gemini-1.5-pro-latest" => "gemini-1.5-pro", - "gemini-1.5-pro-001" => "gemini-1.5-pro", - "gemini-1.5-pro-002" => "gemini-1.5-pro", - "gemini-1.5-pro-exp-0827" => "gemini-1.5-pro", - "gemini-1.5-pro-exp-0801" => "gemini-1.5-pro", + add_alias!("gemini-3-flash-preview", "gemini-3-flash-preview"); + add_alias!("gemini-3-flash-preview-12-2025", "gemini-3-flash-preview"); + add_alias!("gemini-3-flash", "gemini-3-flash-preview"); + add_alias!("gemini-3.1-pro-preview", "gemini-3.1-pro-preview"); + add_alias!("gemini-3.1-pro", "gemini-3.1-pro-preview"); + add_alias!("gemini-3.1-pro-low", "gemini-3.1-pro-preview"); + add_alias!("gemini-3.1-pro-medium", "gemini-3.1-pro-preview"); + add_alias!("gemini-3.1-pro-high", "gemini-3.1-pro-preview"); + add_alias!( + "gemini-3-pro-preview-11-2025", + "gemini-3-pro-preview-11-2025" + ); + add_alias!("gemini-3-pro-preview", "gemini-3-pro-preview-11-2025"); + add_alias!("gemini-3-pro", "gemini-3-pro-preview-11-2025"); + add_alias!("gemini-2.5-pro", "gemini-2.5-pro"); + add_alias!("gemini-2.5-pro-preview-06-05", "gemini-2.5-pro"); + add_alias!("gemini-2.5-pro-preview-05-06", "gemini-2.5-pro"); + add_alias!("gemini-2.5-pro-preview-03-25", "gemini-2.5-pro"); + add_alias!("gemini-2.5-flash", "gemini-2.5-flash"); + add_alias!("gemini-2.5-flash-preview-05-20", "gemini-2.5-flash"); + add_alias!("gemini-2.5-flash-preview-04-17", "gemini-2.5-flash"); + add_alias!("gemini-2.5-flash-lite", "gemini-2.5-flash-lite"); + add_alias!("gemini-2.5-flash-lite-06-17", "gemini-2.5-flash-lite"); + add_alias!("gemini-2.0-pro-exp-02-05", "gemini-2.0-pro-exp-02-05"); + add_alias!("gemini-exp-1206", "gemini-2.0-pro-exp-02-05"); + add_alias!("gemini-2.0-flash", "gemini-2.0-flash"); + add_alias!("gemini-2.0-flash-001", "gemini-2.0-flash"); + add_alias!("gemini-2.0-flash-exp", "gemini-2.0-flash"); + add_alias!("gemini-2.0-flash-lite", "gemini-2.0-flash-lite"); + add_alias!("gemini-2.0-flash-lite-001", "gemini-2.0-flash-lite"); + add_alias!("gemini-1.5-flash", "gemini-1.5-flash"); + add_alias!("gemini-1.5-flash-latest", "gemini-1.5-flash"); + add_alias!("gemini-1.5-flash-001", "gemini-1.5-flash"); + add_alias!("gemini-1.5-flash-002", "gemini-1.5-flash"); + add_alias!("gemini-1.5-flash-8b", "gemini-1.5-flash-8b"); + add_alias!("gemini-1.5-flash-8b-latest", "gemini-1.5-flash-8b"); + add_alias!("gemini-1.5-flash-8b-001", "gemini-1.5-flash-8b"); + add_alias!("gemini-1.5-flash-8b-exp-0924", "gemini-1.5-flash-8b"); + add_alias!("gemini-1.5-flash-8b-exp-0827", "gemini-1.5-flash-8b"); + add_alias!("gemini-1.5-pro", "gemini-1.5-pro"); + add_alias!("gemini-1.5-pro-latest", "gemini-1.5-pro"); + add_alias!("gemini-1.5-pro-001", "gemini-1.5-pro"); + add_alias!("gemini-1.5-pro-002", "gemini-1.5-pro"); + add_alias!("gemini-1.5-pro-exp-0827", "gemini-1.5-pro"); + add_alias!("gemini-1.5-pro-exp-0801", "gemini-1.5-pro"); // Zhipu AI aliases - "zai-glm-4.6" => "glm-4.6", - "glm-5-20260211" => "glm-5", - "glm-5-code" => "glm-5-code", - "glm-5-code-20260211" => "glm-5-code", - "glm-4.5-air-20260211" => "glm-4.5-air", + add_alias!("zai-glm-4.6", "glm-4.6"); + add_alias!("glm-5-20260211", "glm-5"); + add_alias!("glm-5-code", "glm-5-code"); + add_alias!("glm-5-code-20260211", "glm-5-code"); + add_alias!("glm-4.5-air-20260211", "glm-4.5-air"); // OpenAI aliases (continued) - "gpt-5.4" => "gpt-5.4", - "gpt-5.4-2026-03-05" => "gpt-5.4", - "gpt-5.4-mini" => "gpt-5.4-mini", - "gpt-5.4-mini-2026-03-17" => "gpt-5.4-mini", - "gpt-5.4-mini-2026-03-17." => "gpt-5.4-mini", + add_alias!("gpt-5.4", "gpt-5.4"); + add_alias!("gpt-5.4-2026-03-05", "gpt-5.4"); + add_alias!("gpt-5.4-mini", "gpt-5.4-mini"); + add_alias!("gpt-5.4-mini-2026-03-17", "gpt-5.4-mini"); + add_alias!("gpt-5.4-mini-2026-03-17.", "gpt-5.4-mini"); // MiniMax aliases - "minimax-m2.5" => "minimax-m2.5", - "minimax-m2.5-20260211" => "minimax-m2.5", + add_alias!("minimax-m2.5", "minimax-m2.5"); + add_alias!("minimax-m2.5-20260211", "minimax-m2.5"); // StepFun aliases - "step-3.5-flash" => "step-3.5-flash", + add_alias!("step-3.5-flash", "step-3.5-flash"); // Upstage aliases - "solar-pro-3" => "solar-pro-3", + add_alias!("solar-pro-3", "solar-pro-3"); // Aurora aliases - "aurora-alpha" => "aurora-alpha", -}; + add_alias!("aurora-alpha", "aurora-alpha"); +} /// Free-tier model pricing for models accessed via OpenRouter's `:free` suffix /// or other free-tier naming patterns. -static FREE_MODEL_INFO: ModelInfo = ModelInfo { - pricing: PricingStructure::Flat { - input_per_1m: 0.0, - output_per_1m: 0.0, - }, - caching: CachingSupport::None, - is_estimated: false, -}; +fn get_free_model_info() -> ModelInfo { + ModelInfo { + pricing: PricingStructure::Flat { + input_per_1m: 0.0, + output_per_1m: 0.0, + }, + caching: CachingSupport::None, + is_estimated: false, + } +} /// Look up a model name directly in the index and alias tables. -fn lookup_model(name: &str) -> Option<&'static ModelInfo> { - if let Some(model_info) = MODEL_INDEX.get(name) { - return Some(model_info); +fn lookup_model(name: &str) -> Option { + let registry = get_registry_lock().read(); + if let Some(model_info) = registry.index.get(name) { + return Some(model_info.clone()); } - if let Some(&canonical_name) = MODEL_ALIASES.get(name) { - return MODEL_INDEX.get(canonical_name); + if let Some(canonical_name) = registry.aliases.get(name) { + return registry.index.get(canonical_name).cloned(); } None } @@ -1145,7 +1288,7 @@ fn lookup_model(name: &str) -> Option<&'static ModelInfo> { /// `z-ai/glm-5`, `openrouter/aurora-alpha`) by stripping the prefix before /// lookup. Models with a `:free` suffix (OpenRouter free tier) always /// return $0 pricing. -pub fn get_model_info(model_name: &str) -> Option<&'static ModelInfo> { +pub fn get_model_info(model_name: &str) -> Option { // Fast path: direct lookup if let Some(info) = lookup_model(model_name) { return Some(info); @@ -1159,7 +1302,7 @@ pub fn get_model_info(model_name: &str) -> Option<&'static ModelInfo> { // Handle `:free` suffix → always $0 if after_slash.strip_suffix(":free").is_some() { - return Some(&FREE_MODEL_INFO); + return Some(get_free_model_info()); } // Handle other suffixes like `:extended` @@ -1174,7 +1317,7 @@ pub fn get_model_info(model_name: &str) -> Option<&'static ModelInfo> { // Also handle patterns like "minimax-m2.5-free" (without colon) if base_name.strip_suffix("-free").is_some() { - return Some(&FREE_MODEL_INFO); + return Some(get_free_model_info()); } None @@ -1195,7 +1338,7 @@ pub fn calculate_input_cost(model_name: &str, input_tokens: u64) -> f64 { (input_tokens as f64 / 1_000_000.0) * input_per_1m } PricingStructure::Tiered(tiered) => { - calculate_tiered_cost(input_tokens, tiered.tiers, tiered.bracket_pricing, true) + calculate_tiered_cost(input_tokens, &tiered.tiers, tiered.bracket_pricing, true) } }, None => { @@ -1215,7 +1358,7 @@ pub fn calculate_output_cost(model_name: &str, output_tokens: u64) -> f64 { (output_tokens as f64 / 1_000_000.0) * output_per_1m } PricingStructure::Tiered(tiered) => { - calculate_tiered_cost(output_tokens, tiered.tiers, tiered.bracket_pricing, false) + calculate_tiered_cost(output_tokens, &tiered.tiers, tiered.bracket_pricing, false) } }, None => { @@ -1256,7 +1399,7 @@ pub fn calculate_cache_cost( // Google only has read cost, calculate based on tiers calculate_tiered_cache_cost( cache_read_tokens, - tiered.tiers, + &tiered.tiers, tiered.bracket_pricing, ) } @@ -1382,13 +1525,53 @@ where #[cfg(test)] mod tests { use super::{ - calculate_cache_cost, calculate_input_cost, calculate_output_cost, get_model_info, + CachingSupport, ModelInfo, PricingStructure, Registry, calculate_cache_cost, + calculate_input_cost, calculate_output_cost, get_model_info, }; + use std::collections::HashMap; + fn approx_eq(left: f64, right: f64) { assert!((left - right).abs() < 1e-9, "left={left}, right={right}"); } + #[test] + fn test_registry_merging() { + let mut registry = Registry::new_with_defaults(); + let mut custom_models = HashMap::new(); + custom_models.insert( + "super-expensive-o3".to_string(), + ModelInfo { + pricing: PricingStructure::Flat { + input_per_1m: 1000.0, + output_per_1m: 2000.0, + }, + caching: CachingSupport::None, + is_estimated: false, + }, + ); + + let mut custom_aliases = HashMap::new(); + custom_aliases.insert("expensive".to_string(), "super-expensive-o3".to_string()); + + registry.merge(custom_models, custom_aliases); + + let info = registry + .index + .get("super-expensive-o3") + .expect("Should find custom model"); + match info.pricing { + PricingStructure::Flat { input_per_1m, .. } => assert_eq!(input_per_1m, 1000.0), + _ => panic!("Expected flat pricing"), + } + + let canonical = registry + .aliases + .get("expensive") + .expect("Should find aliased model"); + assert_eq!(canonical, "super-expensive-o3"); + } + #[test] fn gemini_3_1_pro_preview_uses_bracket_pricing_for_input() { let cost = calculate_input_cost("gemini-3.1-pro-preview", 250_000); From 32937f103f716573f8a91bc234c1a6094924c431 Mon Sep 17 00:00:00 2001 From: jimyag Date: Wed, 8 Apr 2026 14:37:33 +0800 Subject: [PATCH 2/4] fix: address CodeRabbit review comments --- src/config.rs | 1 + src/models.rs | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/config.rs b/src/config.rs index c399ea3..717b36f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -240,6 +240,7 @@ pub fn set_config_value(key: &str, value: &str) -> Result<()> { #[cfg(test)] mod tests { use super::*; + use crate::models::PricingStructure; use tempfile::TempDir; fn setup_test_config() -> (TempDir, PathBuf) { diff --git a/src/models.rs b/src/models.rs index cee3c83..257deef 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1560,8 +1560,8 @@ mod tests { .index .get("super-expensive-o3") .expect("Should find custom model"); - match info.pricing { - PricingStructure::Flat { input_per_1m, .. } => assert_eq!(input_per_1m, 1000.0), + match &info.pricing { + PricingStructure::Flat { input_per_1m, .. } => assert_eq!(*input_per_1m, 1000.0), _ => panic!("Expected flat pricing"), } From 4a40dfbfddf223a7d495634d9aba555d3da6ea7b Mon Sep 17 00:00:00 2001 From: jimyag Date: Mon, 20 Apr 2026 10:36:39 +0800 Subject: [PATCH 3/4] fix(core): address review feedback on model loading Use an Arc-backed model registry so repeated external config merges are applied without hot-path clones, and stabilize the upload test helpers and test-only HTTP client lifecycle uncovered during verification. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: jimyag --- src/models.rs | 249 ++++++++++++++++++++++++++++++-------------- src/upload.rs | 28 ++++- src/upload/tests.rs | 52 ++++++++- 3 files changed, 244 insertions(+), 85 deletions(-) diff --git a/src/models.rs b/src/models.rs index 1559629..5d4319c 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1,7 +1,7 @@ use parking_lot::RwLock; use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use std::sync::OnceLock; +use std::sync::{Arc, OnceLock}; use crate::utils::warn_once; @@ -82,7 +82,7 @@ pub struct ModelInfo { /// Global registry for models and aliases struct Registry { - index: HashMap, + index: HashMap>, aliases: HashMap, } @@ -100,7 +100,7 @@ impl Registry { external_aliases: HashMap, ) { for (name, info) in external_models { - self.index.insert(name, info); + self.index.insert(name, Arc::new(info)); } for (alias, canonical) in external_aliases { self.aliases.insert(alias, canonical); @@ -109,22 +109,13 @@ impl Registry { } static REGISTRY: OnceLock> = OnceLock::new(); -static EXTERNAL_MODELS_INITIALIZED: std::sync::atomic::AtomicBool = - std::sync::atomic::AtomicBool::new(false); +static FREE_MODEL_INFO: OnceLock> = OnceLock::new(); -/// Initialize the model registry with external configuration. -/// This should be called once at startup. +/// Merge external model configuration into the global registry. pub fn init_external_models( external_models: HashMap, external_aliases: HashMap, ) { - if EXTERNAL_MODELS_INITIALIZED.swap(true, std::sync::atomic::Ordering::SeqCst) { - warn_once( - "WARNING: init_external_models called multiple times. Skipping subsequent loads.", - ); - return; - } - let rwlock = REGISTRY.get_or_init(|| RwLock::new(Registry::new_with_defaults())); let mut registry = rwlock.write(); registry.merge(external_models, external_aliases); @@ -135,18 +126,18 @@ fn get_registry_lock() -> &'static RwLock { } fn populate_defaults( - index: &mut HashMap, + index: &mut HashMap>, aliases: &mut HashMap, ) { macro_rules! add_model { ($name:expr, $pricing:expr, $caching:expr, $est:expr) => { index.insert( $name.to_string(), - ModelInfo { + Arc::new(ModelInfo { pricing: $pricing, caching: $caching, is_estimated: $est, - }, + }), ); }; } @@ -1333,25 +1324,27 @@ fn populate_defaults( /// Free-tier model pricing for models accessed via OpenRouter's `:free` suffix /// or other free-tier naming patterns. -fn get_free_model_info() -> ModelInfo { - ModelInfo { - pricing: PricingStructure::Flat { - input_per_1m: 0.0, - output_per_1m: 0.0, - }, - caching: CachingSupport::None, - is_estimated: false, - } +fn get_free_model_info() -> Arc { + Arc::clone(FREE_MODEL_INFO.get_or_init(|| { + Arc::new(ModelInfo { + pricing: PricingStructure::Flat { + input_per_1m: 0.0, + output_per_1m: 0.0, + }, + caching: CachingSupport::None, + is_estimated: false, + }) + })) } /// Look up a model name directly in the index and alias tables. -fn lookup_model(name: &str) -> Option { +fn lookup_model(name: &str) -> Option> { let registry = get_registry_lock().read(); if let Some(model_info) = registry.index.get(name) { - return Some(model_info.clone()); + return Some(Arc::clone(model_info)); } if let Some(canonical_name) = registry.aliases.get(name) { - return registry.index.get(canonical_name).cloned(); + return registry.index.get(canonical_name).map(Arc::clone); } None } @@ -1362,7 +1355,7 @@ fn lookup_model(name: &str) -> Option { /// `z-ai/glm-5`, `openrouter/aurora-alpha`) by stripping the prefix before /// lookup. Models with a `:free` suffix (OpenRouter free tier) always /// return $0 pricing. -pub fn get_model_info(model_name: &str) -> Option { +pub fn get_model_info(model_name: &str) -> Option> { // Fast path: direct lookup if let Some(info) = lookup_model(model_name) { return Some(info); @@ -1404,17 +1397,21 @@ pub fn is_model_estimated(model_name: &str) -> bool { .unwrap_or(false) } +fn input_cost_for_model(model_info: &ModelInfo, input_tokens: u64) -> f64 { + match &model_info.pricing { + PricingStructure::Flat { input_per_1m, .. } => { + (input_tokens as f64 / 1_000_000.0) * input_per_1m + } + PricingStructure::Tiered(tiered) => { + calculate_tiered_cost(input_tokens, &tiered.tiers, tiered.bracket_pricing, true) + } + } +} + /// Calculate cost for input tokens using the model's pricing structure pub fn calculate_input_cost(model_name: &str, input_tokens: u64) -> f64 { match get_model_info(model_name) { - Some(model_info) => match &model_info.pricing { - PricingStructure::Flat { input_per_1m, .. } => { - (input_tokens as f64 / 1_000_000.0) * input_per_1m - } - PricingStructure::Tiered(tiered) => { - calculate_tiered_cost(input_tokens, &tiered.tiers, tiered.bracket_pricing, true) - } - }, + Some(model_info) => input_cost_for_model(&model_info, input_tokens), None => { warn_once(format!( "WARNING: Unknown model: {model_name}. Defaulting to $0." @@ -1424,17 +1421,21 @@ pub fn calculate_input_cost(model_name: &str, input_tokens: u64) -> f64 { } } +fn output_cost_for_model(model_info: &ModelInfo, output_tokens: u64) -> f64 { + match &model_info.pricing { + PricingStructure::Flat { output_per_1m, .. } => { + (output_tokens as f64 / 1_000_000.0) * output_per_1m + } + PricingStructure::Tiered(tiered) => { + calculate_tiered_cost(output_tokens, &tiered.tiers, tiered.bracket_pricing, false) + } + } +} + /// Calculate cost for output tokens using the model's pricing structure pub fn calculate_output_cost(model_name: &str, output_tokens: u64) -> f64 { match get_model_info(model_name) { - Some(model_info) => match &model_info.pricing { - PricingStructure::Flat { output_per_1m, .. } => { - (output_tokens as f64 / 1_000_000.0) * output_per_1m - } - PricingStructure::Tiered(tiered) => { - calculate_tiered_cost(output_tokens, &tiered.tiers, tiered.bracket_pricing, false) - } - }, + Some(model_info) => output_cost_for_model(&model_info, output_tokens), None => { warn_once(format!( "WARNING: Unknown model: {model_name}. Defaulting to $0." @@ -1444,6 +1445,33 @@ pub fn calculate_output_cost(model_name: &str, output_tokens: u64) -> f64 { } } +fn cache_cost_for_model( + model_info: &ModelInfo, + cache_creation_tokens: u64, + cache_read_tokens: u64, +) -> f64 { + match &model_info.caching { + CachingSupport::None => 0.0, + CachingSupport::OpenAI { + cached_input_per_1m, + } => { + // OpenAI only has cached input cost, no creation cost + (cache_read_tokens as f64 / 1_000_000.0) * cached_input_per_1m + } + CachingSupport::Anthropic { + cache_write_per_1m, + cache_read_per_1m, + } => { + let creation_cost = (cache_creation_tokens as f64 / 1_000_000.0) * cache_write_per_1m; + let read_cost = (cache_read_tokens as f64 / 1_000_000.0) * cache_read_per_1m; + creation_cost + read_cost + } + CachingSupport::Google(tiered) => { + calculate_tiered_cache_cost(cache_read_tokens, &tiered.tiers, tiered.bracket_pricing) + } + } +} + /// Calculate cost for cached tokens pub fn calculate_cache_cost( model_name: &str, @@ -1452,32 +1480,7 @@ pub fn calculate_cache_cost( ) -> f64 { match get_model_info(model_name) { Some(model_info) => { - match &model_info.caching { - CachingSupport::None => 0.0, - CachingSupport::OpenAI { - cached_input_per_1m, - } => { - // OpenAI only has cached input cost, no creation cost - (cache_read_tokens as f64 / 1_000_000.0) * cached_input_per_1m - } - CachingSupport::Anthropic { - cache_write_per_1m, - cache_read_per_1m, - } => { - let creation_cost = - (cache_creation_tokens as f64 / 1_000_000.0) * cache_write_per_1m; - let read_cost = (cache_read_tokens as f64 / 1_000_000.0) * cache_read_per_1m; - creation_cost + read_cost - } - CachingSupport::Google(tiered) => { - // Google only has read cost, calculate based on tiers - calculate_tiered_cache_cost( - cache_read_tokens, - &tiered.tiers, - tiered.bracket_pricing, - ) - } - } + cache_cost_for_model(&model_info, cache_creation_tokens, cache_read_tokens) } None => { warn_once(format!( @@ -1496,11 +1499,19 @@ pub fn calculate_total_cost( cache_creation_tokens: u64, cache_read_tokens: u64, ) -> f64 { - let input_cost = calculate_input_cost(model_name, input_tokens); - let output_cost = calculate_output_cost(model_name, output_tokens); - let cache_cost = calculate_cache_cost(model_name, cache_creation_tokens, cache_read_tokens); - - input_cost + output_cost + cache_cost + match get_model_info(model_name) { + Some(model_info) => { + input_cost_for_model(&model_info, input_tokens) + + output_cost_for_model(&model_info, output_tokens) + + cache_cost_for_model(&model_info, cache_creation_tokens, cache_read_tokens) + } + None => { + warn_once(format!( + "WARNING: Unknown model: {model_name}. Defaulting to $0." + )); + 0.0 + } + } } fn calculate_tiered_cost( @@ -1600,7 +1611,8 @@ where mod tests { use super::{ CachingSupport, ModelInfo, PricingStructure, Registry, calculate_cache_cost, - calculate_input_cost, calculate_output_cost, get_model_info, + calculate_input_cost, calculate_output_cost, get_model_info, get_registry_lock, + init_external_models, }; use std::collections::HashMap; @@ -1609,6 +1621,11 @@ mod tests { assert!((left - right).abs() < 1e-9, "left={left}, right={right}"); } + fn reset_global_registry() { + let registry = get_registry_lock(); + *registry.write() = Registry::new_with_defaults(); + } + #[test] fn test_registry_merging() { let mut registry = Registry::new_with_defaults(); @@ -1646,6 +1663,56 @@ mod tests { assert_eq!(canonical, "super-expensive-o3"); } + #[test] + fn init_external_models_accepts_multiple_calls() { + reset_global_registry(); + + let mut first_models = HashMap::new(); + first_models.insert( + "review-first-model".to_string(), + ModelInfo { + pricing: PricingStructure::Flat { + input_per_1m: 1.0, + output_per_1m: 2.0, + }, + caching: CachingSupport::None, + is_estimated: false, + }, + ); + let mut first_aliases = HashMap::new(); + first_aliases.insert( + "review-first-alias".to_string(), + "review-first-model".to_string(), + ); + + init_external_models(first_models, first_aliases); + assert!(get_model_info("review-first-alias").is_some()); + + let mut second_models = HashMap::new(); + second_models.insert( + "review-second-model".to_string(), + ModelInfo { + pricing: PricingStructure::Flat { + input_per_1m: 3.0, + output_per_1m: 4.0, + }, + caching: CachingSupport::None, + is_estimated: false, + }, + ); + let mut second_aliases = HashMap::new(); + second_aliases.insert( + "review-second-alias".to_string(), + "review-second-model".to_string(), + ); + + init_external_models(second_models, second_aliases); + + assert!(get_model_info("review-second-alias").is_some()); + + reset_global_registry(); + } + #[test] fn gemini_3_1_pro_preview_uses_bracket_pricing_for_input() { let cost = calculate_input_cost("gemini-3.1-pro-preview", 250_000); @@ -1705,4 +1772,30 @@ mod tests { approx_eq(output_cost, 150.0); approx_eq(cache_cost, 37.5); } + + #[test] + fn repeated_tiered_model_lookups_reuse_the_same_tier_storage() { + let first = get_model_info("gemini-2.5-pro").expect("model should exist"); + let second = get_model_info("gemini-2.5-pro").expect("model should exist"); + + match (&first.pricing, &second.pricing) { + (PricingStructure::Tiered(first_tiered), PricingStructure::Tiered(second_tiered)) => { + assert!( + std::ptr::eq(first_tiered.tiers.as_ptr(), second_tiered.tiers.as_ptr()), + "tier pricing should not be reallocated on each lookup" + ); + } + _ => panic!("Expected tiered pricing"), + } + + match (&first.caching, &second.caching) { + (CachingSupport::Google(first_tiered), CachingSupport::Google(second_tiered)) => { + assert!( + std::ptr::eq(first_tiered.tiers.as_ptr(), second_tiered.tiers.as_ptr()), + "cache tiers should not be reallocated on each lookup" + ); + } + _ => panic!("Expected Google caching"), + } + } } diff --git a/src/upload.rs b/src/upload.rs index 8f4d26b..1d5572c 100644 --- a/src/upload.rs +++ b/src/upload.rs @@ -6,9 +6,12 @@ use crate::utils; use anyhow::{Context, Result}; use parking_lot::Mutex; use std::path::PathBuf; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; use std::time::{Duration, Instant}; +#[cfg(not(test))] +use std::sync::OnceLock; + fn upload_log_path() -> PathBuf { std::env::temp_dir().join("SPLITRAIL.log") } @@ -41,17 +44,32 @@ fn upload_debug_log(line: impl Into) { append_upload_log(&line); } +#[cfg(not(test))] static HTTP_CLIENT: OnceLock = OnceLock::new(); /// Get the shared HTTP client singleton -pub fn get_http_client() -> &'static reqwest::Client { - HTTP_CLIENT.get_or_init(|| { +pub fn get_http_client() -> reqwest::Client { + #[cfg(test)] + { reqwest::Client::builder() .timeout(Duration::from_secs(120)) .danger_accept_invalid_certs(true) .build() .expect("Failed to create HTTP client") - }) + } + + #[cfg(not(test))] + { + HTTP_CLIENT + .get_or_init(|| { + reqwest::Client::builder() + .timeout(Duration::from_secs(120)) + .danger_accept_invalid_certs(true) + .build() + .expect("Failed to create HTTP client") + }) + .clone() + } } #[cfg(test)] @@ -155,7 +173,7 @@ where total_messages, upload_debug, }; - match upload_single_chunk(client, config, chunk, &ctx, &mut progress_callback).await { + match upload_single_chunk(&client, config, chunk, &ctx, &mut progress_callback).await { Ok(()) => { last_err = None; break; diff --git a/src/upload/tests.rs b/src/upload/tests.rs index 1c8bea1..d803b90 100644 --- a/src/upload/tests.rs +++ b/src/upload/tests.rs @@ -13,6 +13,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use tempfile::TempDir; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; +use tokio::sync::oneshot; use crate::config::set_test_config_path; @@ -50,6 +51,50 @@ fn make_stats_with_messages(messages: Vec) -> MultiAnalyzer } } +async fn read_http_request(socket: &mut tokio::net::TcpStream) { + let mut buffer = Vec::new(); + let mut content_length = None; + + loop { + let mut chunk = [0u8; 8192]; + let bytes_read = socket.read(&mut chunk).await.expect("read request"); + if bytes_read == 0 { + break; + } + + buffer.extend_from_slice(&chunk[..bytes_read]); + + if content_length.is_none() + && let Some(headers_end) = buffer.windows(4).position(|window| window == b"\r\n\r\n") + { + let headers = &buffer[..headers_end]; + let headers = String::from_utf8_lossy(headers); + content_length = headers.lines().find_map(|line| { + line.strip_prefix("Content-Length: ") + .and_then(|value| value.trim().parse::().ok()) + }); + + if let Some(content_length) = content_length { + let body_bytes = buffer.len() - (headers_end + 4); + if body_bytes >= content_length { + break; + } + } else { + break; + } + } + + if let Some(content_length) = content_length + && let Some(headers_end) = buffer.windows(4).position(|window| window == b"\r\n\r\n") + { + let body_bytes = buffer.len() - (headers_end + 4); + if body_bytes >= content_length { + break; + } + } + } +} + async fn start_test_server( status_line: &str, body: &str, @@ -72,16 +117,17 @@ async fn start_test_server( let base_url = format!("http://{}", addr); let status_line = status_line.to_string(); let body = body.to_string(); + let (ready_tx, ready_rx) = oneshot::channel(); tokio::spawn(async move { + let _ = ready_tx.send(()); for _ in 0..expected_requests { let (mut socket, _) = match listener.accept().await { Ok(conn) => conn, Err(_) => return, }; - let mut buf = [0u8; 4096]; - let _ = socket.read(&mut buf).await; + read_http_request(&mut socket).await; request_counter.fetch_add(1, Ordering::SeqCst); @@ -96,6 +142,8 @@ async fn start_test_server( } }); + let _ = ready_rx.await; + Some(base_url) } From 7a57679805bff4a7799e810bc5c28e7c17fa5637 Mon Sep 17 00:00:00 2001 From: jimyag Date: Mon, 20 Apr 2026 10:54:26 +0800 Subject: [PATCH 4/4] fix(core): harden model registry and upload client Resolve the new review findings by validating external tier configs, following alias chains transitively, restoring the missing doubao canonical model, and tightening upload TLS and header parsing behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: jimyag --- src/models.rs | 188 +++++++++++++++++++++++++++++++++++++++++--- src/upload.rs | 38 ++++++--- src/upload/tests.rs | 29 ++++++- 3 files changed, 229 insertions(+), 26 deletions(-) diff --git a/src/models.rs b/src/models.rs index 5d4319c..06ad0a0 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1,6 +1,6 @@ use parking_lot::RwLock; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::{Arc, OnceLock}; use crate::utils::warn_once; @@ -100,12 +100,59 @@ impl Registry { external_aliases: HashMap, ) { for (name, info) in external_models { + if !Self::validate_model_info(&info) { + warn_once(format!( + "WARNING: init_external_models ignoring invalid tier config for model `{name}`." + )); + continue; + } self.index.insert(name, Arc::new(info)); } for (alias, canonical) in external_aliases { self.aliases.insert(alias, canonical); } } + + fn validate_model_info(info: &ModelInfo) -> bool { + let pricing_ok = match &info.pricing { + PricingStructure::Flat { .. } => true, + PricingStructure::Tiered(tiered) => { + Self::validate_tier_bounds(&tiered.tiers, |tier| tier.max_tokens) + } + }; + + let caching_ok = match &info.caching { + CachingSupport::Google(tiered) => { + Self::validate_tier_bounds(&tiered.tiers, |tier| tier.max_tokens) + } + _ => true, + }; + + pricing_ok && caching_ok + } + + fn validate_tier_bounds(tiers: &[T], max_tokens: F) -> bool + where + F: Fn(&T) -> Option, + { + if tiers.is_empty() { + return false; + } + + let mut previous_limit = 0_u64; + + for (index, tier) in tiers.iter().enumerate() { + match max_tokens(tier) { + Some(limit) if limit > previous_limit && index + 1 < tiers.len() => { + previous_limit = limit; + } + None if index + 1 == tiers.len() => return true, + _ => return false, + } + } + + false + } } static REGISTRY: OnceLock> = OnceLock::new(); @@ -1054,6 +1101,19 @@ fn populate_defaults( false ); + // ByteDance / Doubao Models + add_model!( + "doubao-seed-2.0-code", + PricingStructure::Flat { + input_per_1m: 0.67, + output_per_1m: 3.36 + }, + CachingSupport::OpenAI { + cached_input_per_1m: 0.14 + }, + true + ); + // Z.AI (Zhipu AI) - Additional Models add_model!( "glm-5", @@ -1340,13 +1400,18 @@ fn get_free_model_info() -> Arc { /// Look up a model name directly in the index and alias tables. fn lookup_model(name: &str) -> Option> { let registry = get_registry_lock().read(); - if let Some(model_info) = registry.index.get(name) { - return Some(Arc::clone(model_info)); - } - if let Some(canonical_name) = registry.aliases.get(name) { - return registry.index.get(canonical_name).map(Arc::clone); + let mut current = name; + let mut visited = HashSet::new(); + + loop { + if let Some(model_info) = registry.index.get(current) { + return Some(Arc::clone(model_info)); + } + if !visited.insert(current.to_string()) { + return None; + } + current = registry.aliases.get(current)?.as_str(); } - None } /// Get model info by any valid name (canonical or alias). @@ -1610,9 +1675,9 @@ where #[cfg(test)] mod tests { use super::{ - CachingSupport, ModelInfo, PricingStructure, Registry, calculate_cache_cost, - calculate_input_cost, calculate_output_cost, get_model_info, get_registry_lock, - init_external_models, + CachingSupport, CachingTier, ModelInfo, PricingStructure, PricingTier, Registry, + TieredCaching, TieredPricing, calculate_cache_cost, calculate_input_cost, + calculate_output_cost, get_model_info, get_registry_lock, init_external_models, }; use std::collections::HashMap; @@ -1713,6 +1778,95 @@ mod tests { reset_global_registry(); } + #[test] + fn transitive_aliases_resolve_to_the_final_model() { + reset_global_registry(); + + let mut models = HashMap::new(); + models.insert( + "review-chain-model".to_string(), + ModelInfo { + pricing: PricingStructure::Flat { + input_per_1m: 1.5, + output_per_1m: 2.5, + }, + caching: CachingSupport::None, + is_estimated: false, + }, + ); + + let mut aliases = HashMap::new(); + aliases.insert("review-chain-a".to_string(), "review-chain-b".to_string()); + aliases.insert( + "review-chain-b".to_string(), + "review-chain-model".to_string(), + ); + + init_external_models(models, aliases); + + let model_info = get_model_info("review-chain-a").expect("alias chain should resolve"); + match &model_info.pricing { + PricingStructure::Flat { input_per_1m, .. } => approx_eq(*input_per_1m, 1.5), + _ => panic!("Expected flat pricing"), + } + + reset_global_registry(); + } + + #[test] + fn invalid_external_tier_configs_are_skipped() { + reset_global_registry(); + + let mut models = HashMap::new(); + models.insert( + "review-invalid-tier-model".to_string(), + ModelInfo { + pricing: PricingStructure::Tiered(TieredPricing { + tiers: vec![ + PricingTier { + max_tokens: Some(200), + input_per_1m: 1.0, + output_per_1m: 2.0, + }, + PricingTier { + max_tokens: Some(100), + input_per_1m: 3.0, + output_per_1m: 4.0, + }, + ], + bracket_pricing: false, + }), + caching: CachingSupport::Google(TieredCaching { + tiers: vec![ + CachingTier { + max_tokens: Some(50), + cached_input_per_1m: 0.5, + }, + CachingTier { + max_tokens: Some(25), + cached_input_per_1m: 0.25, + }, + ], + bracket_pricing: false, + }), + is_estimated: false, + }, + ); + + let mut aliases = HashMap::new(); + aliases.insert( + "review-invalid-tier-alias".to_string(), + "review-invalid-tier-model".to_string(), + ); + + init_external_models(models, aliases); + + assert!(get_model_info("review-invalid-tier-model").is_none()); + assert!(get_model_info("review-invalid-tier-alias").is_none()); + + reset_global_registry(); + } + #[test] fn gemini_3_1_pro_preview_uses_bracket_pricing_for_input() { let cost = calculate_input_cost("gemini-3.1-pro-preview", 250_000); @@ -1773,6 +1927,20 @@ mod tests { approx_eq(cache_cost, 37.5); } + #[test] + fn doubao_seed_code_alias_resolves() { + let model_info = get_model_info("doubao-seed-code").expect("alias should resolve"); + assert!(model_info.is_estimated); + + let input_cost = calculate_input_cost("doubao-seed-code", 1_000_000); + let output_cost = calculate_output_cost("doubao-seed-code", 1_000_000); + let cache_cost = calculate_cache_cost("doubao-seed-code", 0, 1_000_000); + + approx_eq(input_cost, 0.67); + approx_eq(output_cost, 3.36); + approx_eq(cache_cost, 0.14); + } + #[test] fn repeated_tiered_model_lookups_reuse_the_same_tier_storage() { let first = get_model_info("gemini-2.5-pro").expect("model should exist"); diff --git a/src/upload.rs b/src/upload.rs index 1d5572c..685216a 100644 --- a/src/upload.rs +++ b/src/upload.rs @@ -47,27 +47,41 @@ fn upload_debug_log(line: impl Into) { #[cfg(not(test))] static HTTP_CLIENT: OnceLock = OnceLock::new(); +fn allow_invalid_certs_from_env(value: Option<&str>) -> bool { + value.is_some_and(|value| value == "1" || value.eq_ignore_ascii_case("true")) +} + +#[cfg(not(test))] +fn upload_accept_invalid_certs_enabled() -> bool { + allow_invalid_certs_from_env( + std::env::var("SPLITRAIL_ACCEPT_INVALID_CERTS") + .ok() + .as_deref(), + ) +} + +fn build_http_client(allow_invalid_certs: bool) -> reqwest::Client { + let builder = reqwest::Client::builder().timeout(Duration::from_secs(120)); + let builder = if allow_invalid_certs { + builder.danger_accept_invalid_certs(true) + } else { + builder + }; + + builder.build().expect("Failed to create HTTP client") +} + /// Get the shared HTTP client singleton pub fn get_http_client() -> reqwest::Client { #[cfg(test)] { - reqwest::Client::builder() - .timeout(Duration::from_secs(120)) - .danger_accept_invalid_certs(true) - .build() - .expect("Failed to create HTTP client") + build_http_client(true) } #[cfg(not(test))] { HTTP_CLIENT - .get_or_init(|| { - reqwest::Client::builder() - .timeout(Duration::from_secs(120)) - .danger_accept_invalid_certs(true) - .build() - .expect("Failed to create HTTP client") - }) + .get_or_init(|| build_http_client(upload_accept_invalid_certs_enabled())) .clone() } } diff --git a/src/upload/tests.rs b/src/upload/tests.rs index d803b90..a7fb0a0 100644 --- a/src/upload/tests.rs +++ b/src/upload/tests.rs @@ -51,6 +51,30 @@ fn make_stats_with_messages(messages: Vec) -> MultiAnalyzer } } +fn parse_content_length(headers: &str) -> Option { + headers.lines().find_map(|line| { + let (name, value) = line.split_once(':')?; + name.trim() + .eq_ignore_ascii_case("Content-Length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) +} + +#[test] +fn parse_content_length_matches_header_names_case_insensitively() { + let headers = "POST / HTTP/1.1\r\nhost: localhost\r\ncontent-length: 123\r\n\r\n"; + assert_eq!(parse_content_length(headers), Some(123)); +} + +#[test] +fn invalid_cert_opt_in_requires_explicit_truthy_value() { + assert!(!allow_invalid_certs_from_env(None)); + assert!(!allow_invalid_certs_from_env(Some("false"))); + assert!(allow_invalid_certs_from_env(Some("true"))); + assert!(allow_invalid_certs_from_env(Some("1"))); +} + async fn read_http_request(socket: &mut tokio::net::TcpStream) { let mut buffer = Vec::new(); let mut content_length = None; @@ -69,10 +93,7 @@ async fn read_http_request(socket: &mut tokio::net::TcpStream) { { let headers = &buffer[..headers_end]; let headers = String::from_utf8_lossy(headers); - content_length = headers.lines().find_map(|line| { - line.strip_prefix("Content-Length: ") - .and_then(|value| value.trim().parse::().ok()) - }); + content_length = parse_content_length(&headers); if let Some(content_length) = content_length { let body_bytes = buffer.len() - (headers_end + 4);