From a408d5467ed5b5f6c142f3ac1c7ef3931e1f4570 Mon Sep 17 00:00:00 2001
From: OnlyYu1996 <1158673577@qq.com>
Date: Sun, 17 May 2026 07:23:59 +0800
Subject: [PATCH] fix(stats): support provider-prefixed pricing env keys

---
 src/cortex-cli/src/stats_cmd.rs | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/cortex-cli/src/stats_cmd.rs b/src/cortex-cli/src/stats_cmd.rs
index 1e407503e..8caaf990c 100644
--- a/src/cortex-cli/src/stats_cmd.rs
+++ b/src/cortex-cli/src/stats_cmd.rs
@@ -118,9 +118,11 @@ fn load_custom_pricing() -> std::collections::HashMap<String, ModelPricing> {
     // Try to load from environment variables in format:
     // CORTEX_PRICING_<MODEL>=<input_price>,<output_price>
     // Example: CORTEX_PRICING_GPT4O=2.5,10.0
+    // Use double underscore to encode provider-prefixed model IDs:
+    // CORTEX_PRICING_ANTHROPIC__CLAUDE_3_5_SONNET=3.0,15.0
     for (key, value) in std::env::vars() {
         if let Some(model_suffix) = key.strip_prefix("CORTEX_PRICING_") {
-            let model_name = model_suffix.to_lowercase().replace('_', "-");
+            let model_name = pricing_env_model_name(model_suffix);
             let parts: Vec<&str> = value.split(',').collect();
             if parts.len() == 2
                 && let (Ok(input), Ok(output)) = (
@@ -142,6 +144,14 @@ fn load_custom_pricing() -> std::collections::HashMap<String, ModelPricing> {
     custom
 }
 
+fn pricing_env_model_name(model_suffix: &str) -> String {
+    model_suffix
+        .to_lowercase()
+        .replace("__", "\0")
+        .replace('_', "-")
+        .replace('\0', "/")
+}
+
 impl StatsCli {
     /// Run the stats command.
     pub async fn run(self) -> Result<()> {
@@ -735,6 +745,27 @@ mod tests {
         assert!((cost - 12.5).abs() < 0.001);
     }
 
+    #[test]
+    fn test_pricing_env_model_name_supports_provider_prefixed_models() {
+        assert_eq!(
+            pricing_env_model_name("ANTHROPIC__CLAUDE_3_5_SONNET"),
+            "anthropic/claude-3-5-sonnet"
+        );
+        assert_eq!(
+            pricing_env_model_name("OPENAI__GPT_4O_MINI"),
+            "openai/gpt-4o-mini"
+        );
+    }
+
+    #[test]
+    fn test_pricing_env_model_name_preserves_legacy_underscore_mapping() {
+        assert_eq!(pricing_env_model_name("GPT_4O_MINI"), "gpt-4o-mini");
+        assert_eq!(
+            pricing_env_model_name("CLAUDE_3_5_HAIKU"),
+            "claude-3-5-haiku"
+        );
+    }
+
     #[test]
     fn test_validate_days_range() {
         // Valid values