feat: Add comprehensive cost tracking display for LLM usage (#2992)

Co-authored-by: jack <jack@deck.local>
Co-authored-by: Bradley Axen <baxen@squareup.com>

crates/goose/Cargo.toml

@@ -17,6 +17,7 @@ mcp-core = { path = "../mcp-core" }
 anyhow = "1.0"
 thiserror = "1.0"
 futures = "0.3"
+dirs = "5.0"
 reqwest = { version = "0.12.9", features = [
         "rustls-tls-native-roots",
         "json",

crates/goose/src/providers/anthropic.rs

@@ -5,7 +5,7 @@ use reqwest::{Client, StatusCode};
 use serde_json::Value;
 use std::time::Duration;
 
-use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
+use super::base::{ConfigKey, ModelInfo, Provider, ProviderMetadata, ProviderUsage};
 use super::errors::ProviderError;
 use super::formats::anthropic::{create_request, get_usage, response_to_message};
 use super::utils::{emit_debug_trace, get_model};
@@ -122,12 +122,18 @@ impl AnthropicProvider {
 #[async_trait]
 impl Provider for AnthropicProvider {
     fn metadata() -> ProviderMetadata {
-        ProviderMetadata::new(
+        ProviderMetadata::with_models(
             "anthropic",
             "Anthropic",
             "Claude and other models from Anthropic",
             ANTHROPIC_DEFAULT_MODEL,
-            ANTHROPIC_KNOWN_MODELS.to_vec(),
+            vec![
+                ModelInfo::with_cost("claude-3-5-sonnet-20241022", 200000, 0.000003, 0.000015),
+                ModelInfo::with_cost("claude-3-5-haiku-20241022", 200000, 0.000001, 0.000005),
+                ModelInfo::with_cost("claude-3-opus-20240229", 200000, 0.000015, 0.000075),
+                ModelInfo::with_cost("claude-3-sonnet-20240229", 200000, 0.000003, 0.000015),
+                ModelInfo::with_cost("claude-3-haiku-20240307", 200000, 0.00000025, 0.00000125),
+            ],
             ANTHROPIC_DOC_URL,
             vec![
                 ConfigKey::new("ANTHROPIC_API_KEY", true, true, None),

crates/goose/src/providers/base.rs

@@ -32,6 +32,41 @@ pub struct ModelInfo {
     pub name: String,
     /// The maximum context length this model supports
     pub context_limit: usize,
+    /// Cost per token for input (optional)
+    pub input_token_cost: Option<f64>,
+    /// Cost per token for output (optional)
+    pub output_token_cost: Option<f64>,
+    /// Currency for the costs (default: "$")
+    pub currency: Option<String>,
+}
+
+impl ModelInfo {
+    /// Create a new ModelInfo with just name and context limit
+    pub fn new(name: impl Into<String>, context_limit: usize) -> Self {
+        Self {
+            name: name.into(),
+            context_limit,
+            input_token_cost: None,
+            output_token_cost: None,
+            currency: None,
+        }
+    }
+
+    /// Create a new ModelInfo with cost information (per token)
+    pub fn with_cost(
+        name: impl Into<String>,
+        context_limit: usize,
+        input_cost: f64,
+        output_cost: f64,
+    ) -> Self {
+        Self {
+            name: name.into(),
+            context_limit,
+            input_token_cost: Some(input_cost),
+            output_token_cost: Some(output_cost),
+            currency: Some("$".to_string()),
+        }
+    }
 }
 
 /// Metadata about a provider's configuration requirements and capabilities
@@ -74,6 +109,9 @@ impl ProviderMetadata {
                 .map(|&name| ModelInfo {
                     name: name.to_string(),
                     context_limit: ModelConfig::new(name.to_string()).context_limit(),
+                    input_token_cost: None,
+                    output_token_cost: None,
+                    currency: None,
                 })
                 .collect(),
             model_doc_link: model_doc_link.to_string(),
@@ -81,6 +119,27 @@ impl ProviderMetadata {
         }
     }
 
+    /// Create a new ProviderMetadata with ModelInfo objects that include cost data
+    pub fn with_models(
+        name: &str,
+        display_name: &str,
+        description: &str,
+        default_model: &str,
+        models: Vec<ModelInfo>,
+        model_doc_link: &str,
+        config_keys: Vec<ConfigKey>,
+    ) -> Self {
+        Self {
+            name: name.to_string(),
+            display_name: display_name.to_string(),
+            description: description.to_string(),
+            default_model: default_model.to_string(),
+            known_models: models,
+            model_doc_link: model_doc_link.to_string(),
+            config_keys,
+        }
+    }
+
     pub fn empty() -> Self {
         Self {
             name: "".to_string(),
@@ -313,6 +372,9 @@ mod tests {
         let info = ModelInfo {
             name: "test-model".to_string(),
             context_limit: 1000,
+            input_token_cost: None,
+            output_token_cost: None,
+            currency: None,
         };
         assert_eq!(info.context_limit, 1000);
 
@@ -320,6 +382,9 @@ mod tests {
         let info2 = ModelInfo {
             name: "test-model".to_string(),
             context_limit: 1000,
+            input_token_cost: None,
+            output_token_cost: None,
+            currency: None,
         };
         assert_eq!(info, info2);
 
@@ -327,7 +392,20 @@ mod tests {
         let info3 = ModelInfo {
             name: "test-model".to_string(),
             context_limit: 2000,
+            input_token_cost: None,
+            output_token_cost: None,
+            currency: None,
         };
         assert_ne!(info, info3);
     }
+
+    #[test]
+    fn test_model_info_with_cost() {
+        let info = ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001);
+        assert_eq!(info.name, "gpt-4o");
+        assert_eq!(info.context_limit, 128000);
+        assert_eq!(info.input_token_cost, Some(0.0000025));
+        assert_eq!(info.output_token_cost, Some(0.00001));
+        assert_eq!(info.currency, Some("$".to_string()));
+    }
 }

crates/goose/src/providers/mod.rs

@@ -18,6 +18,7 @@ pub mod oauth;
 pub mod ollama;
 pub mod openai;
 pub mod openrouter;
+pub mod pricing;
 pub mod sagemaker_tgi;
 pub mod snowflake;
 pub mod toolshim;

crates/goose/src/providers/openai.rs

@@ -5,7 +5,7 @@ use serde_json::Value;
 use std::collections::HashMap;
 use std::time::Duration;
 
-use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
+use super::base::{ConfigKey, ModelInfo, Provider, ProviderMetadata, ProviderUsage, Usage};
 use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse};
 use super::errors::ProviderError;
 use super::formats::openai::{create_request, get_usage, response_to_message};
@@ -126,12 +126,20 @@ impl OpenAiProvider {
 #[async_trait]
 impl Provider for OpenAiProvider {
     fn metadata() -> ProviderMetadata {
-        ProviderMetadata::new(
+        ProviderMetadata::with_models(
             "openai",
             "OpenAI",
             "GPT-4 and other OpenAI models, including OpenAI compatible ones",
             OPEN_AI_DEFAULT_MODEL,
-            OPEN_AI_KNOWN_MODELS.to_vec(),
+            vec![
+                ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001),
+                ModelInfo::with_cost("gpt-4o-mini", 128000, 0.00000015, 0.0000006),
+                ModelInfo::with_cost("gpt-4-turbo", 128000, 0.00001, 0.00003),
+                ModelInfo::with_cost("gpt-3.5-turbo", 16385, 0.0000005, 0.0000015),
+                ModelInfo::with_cost("o1", 200000, 0.000015, 0.00006),
+                ModelInfo::with_cost("o3", 200000, 0.000015, 0.00006), // Using o1 pricing as placeholder
+                ModelInfo::with_cost("o4-mini", 128000, 0.000003, 0.000012), // Using o1-mini pricing as placeholder
+            ],
             OPEN_AI_DOC_URL,
             vec![
                 ConfigKey::new("OPENAI_API_KEY", true, true, None),

crates/goose/src/providers/pricing.rs

@@ -0,0 +1,387 @@
+use anyhow::Result;
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use tokio::sync::RwLock;
+
+/// Disk cache configuration
+const CACHE_FILE_NAME: &str = "pricing_cache.json";
+const CACHE_TTL_DAYS: u64 = 7; // Cache for 7 days
+
+/// Get the cache directory path
+fn get_cache_dir() -> Result<PathBuf> {
+    let cache_dir = if let Ok(goose_dir) = std::env::var("GOOSE_CACHE_DIR") {
+        PathBuf::from(goose_dir)
+    } else {
+        dirs::cache_dir()
+            .ok_or_else(|| anyhow::anyhow!("Could not determine cache directory"))?
+            .join("goose")
+    };
+    Ok(cache_dir)
+}
+
+/// Cached pricing data structure for disk storage
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CachedPricingData {
+    /// Nested HashMap: provider -> model -> pricing info
+    pub pricing: HashMap<String, HashMap<String, PricingInfo>>,
+    /// Unix timestamp when data was fetched
+    pub fetched_at: u64,
+}
+
+/// Simplified pricing info for efficient storage
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PricingInfo {
+    pub input_cost: f64,  // Cost per token
+    pub output_cost: f64, // Cost per token
+    pub context_length: Option<u32>,
+}
+
+/// Cache for OpenRouter pricing data with disk persistence
+pub struct PricingCache {
+    /// In-memory cache
+    memory_cache: Arc<RwLock<Option<CachedPricingData>>>,
+}
+
+impl PricingCache {
+    pub fn new() -> Self {
+        Self {
+            memory_cache: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    /// Load pricing from disk cache
+    async fn load_from_disk(&self) -> Result<Option<CachedPricingData>> {
+        let cache_path = get_cache_dir()?.join(CACHE_FILE_NAME);
+
+        if !cache_path.exists() {
+            return Ok(None);
+        }
+
+        match tokio::fs::read(&cache_path).await {
+            Ok(data) => {
+                match serde_json::from_slice::<CachedPricingData>(&data) {
+                    Ok(cached) => {
+                        // Check if cache is still valid
+                        let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
+                        let age_days = (now - cached.fetched_at) / (24 * 60 * 60);
+
+                        if age_days < CACHE_TTL_DAYS {
+                            tracing::info!(
+                                "Loaded pricing data from disk cache (age: {} days)",
+                                age_days
+                            );
+                            Ok(Some(cached))
+                        } else {
+                            tracing::info!("Disk cache expired (age: {} days)", age_days);
+                            Ok(None)
+                        }
+                    }
+                    Err(e) => {
+                        tracing::warn!("Failed to parse pricing cache: {}", e);
+                        Ok(None)
+                    }
+                }
+            }
+            Err(e) => {
+                tracing::warn!("Failed to read pricing cache: {}", e);
+                Ok(None)
+            }
+        }
+    }
+
+    /// Save pricing data to disk
+    async fn save_to_disk(&self, data: &CachedPricingData) -> Result<()> {
+        let cache_dir = get_cache_dir()?;
+        tokio::fs::create_dir_all(&cache_dir).await?;
+
+        let cache_path = cache_dir.join(CACHE_FILE_NAME);
+        let json_data = serde_json::to_vec_pretty(data)?;
+        tokio::fs::write(&cache_path, json_data).await?;
+
+        tracing::info!("Saved pricing data to disk cache");
+        Ok(())
+    }
+
+    /// Get pricing for a specific model
+    pub async fn get_model_pricing(&self, provider: &str, model: &str) -> Option<PricingInfo> {
+        // Try memory cache first
+        {
+            let cache = self.memory_cache.read().await;
+            if let Some(cached) = &*cache {
+                return cached
+                    .pricing
+                    .get(&provider.to_lowercase())
+                    .and_then(|models| models.get(model))
+                    .cloned();
+            }
+        }
+
+        // Try loading from disk
+        if let Ok(Some(disk_cache)) = self.load_from_disk().await {
+            // Update memory cache
+            {
+                let mut cache = self.memory_cache.write().await;
+                *cache = Some(disk_cache.clone());
+            }
+
+            return disk_cache
+                .pricing
+                .get(&provider.to_lowercase())
+                .and_then(|models| models.get(model))
+                .cloned();
+        }
+
+        None
+    }
+
+    /// Force refresh pricing data from OpenRouter
+    pub async fn refresh(&self) -> Result<()> {
+        let pricing = fetch_openrouter_pricing_internal().await?;
+
+        // Convert to our efficient structure
+        let mut structured_pricing: HashMap<String, HashMap<String, PricingInfo>> = HashMap::new();
+
+        for (model_id, model) in pricing {
+            if let Some((provider, model_name)) = parse_model_id(&model_id) {
+                if let (Some(input_cost), Some(output_cost)) = (
+                    convert_pricing(&model.pricing.prompt),
+                    convert_pricing(&model.pricing.completion),
+                ) {
+                    let provider_lower = provider.to_lowercase();
+                    let provider_models = structured_pricing.entry(provider_lower).or_default();
+
+                    provider_models.insert(
+                        model_name,
+                        PricingInfo {
+                            input_cost,
+                            output_cost,
+                            context_length: model.context_length,
+                        },
+                    );
+                }
+            }
+        }
+
+        let cached_data = CachedPricingData {
+            pricing: structured_pricing,
+            fetched_at: SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(),
+        };
+
+        // Log how many models we fetched
+        let total_models: usize = cached_data
+            .pricing
+            .values()
+            .map(|models| models.len())
+            .sum();
+        tracing::info!(
+            "Fetched pricing for {} providers with {} total models from OpenRouter",
+            cached_data.pricing.len(),
+            total_models
+        );
+
+        // Save to disk
+        self.save_to_disk(&cached_data).await?;
+
+        // Update memory cache
+        {
+            let mut cache = self.memory_cache.write().await;
+            *cache = Some(cached_data);
+        }
+
+        Ok(())
+    }
+
+    /// Initialize cache (load from disk or fetch if needed)
+    pub async fn initialize(&self) -> Result<()> {
+        // Try loading from disk first
+        if let Ok(Some(cached)) = self.load_from_disk().await {
+            // Log how many models we have cached
+            let total_models: usize = cached.pricing.values().map(|models| models.len()).sum();
+            tracing::info!(
+                "Loaded {} providers with {} total models from disk cache",
+                cached.pricing.len(),
+                total_models
+            );
+
+            // Update memory cache
+            {
+                let mut cache = self.memory_cache.write().await;
+                *cache = Some(cached);
+            }
+
+            return Ok(());
+        }
+
+        // If no disk cache, fetch from OpenRouter
+        tracing::info!("No valid disk cache found, fetching from OpenRouter");
+        self.refresh().await
+    }
+}
+
+impl Default for PricingCache {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+// Global cache instance
+lazy_static::lazy_static! {
+    static ref PRICING_CACHE: PricingCache = PricingCache::new();
+    static ref HTTP_CLIENT: Client = Client::builder()
+        .timeout(Duration::from_secs(30))
+        .pool_idle_timeout(Duration::from_secs(90))
+        .pool_max_idle_per_host(10)
+        .build()
+        .unwrap();
+}
+
+/// OpenRouter model pricing information
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OpenRouterModel {
+    pub id: String,
+    pub name: String,
+    pub pricing: OpenRouterPricing,
+    pub context_length: Option<u32>,
+    pub architecture: Option<Architecture>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OpenRouterPricing {
+    pub prompt: String,     // Cost per token for input (in USD)
+    pub completion: String, // Cost per token for output (in USD)
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Architecture {
+    pub modality: String,
+    pub tokenizer: String,
+    pub instruct_type: Option<String>,
+}
+
+/// Response from OpenRouter models endpoint
+#[derive(Debug, Deserialize)]
+pub struct OpenRouterModelsResponse {
+    pub data: Vec<OpenRouterModel>,
+}
+
+/// Internal function to fetch pricing data
+async fn fetch_openrouter_pricing_internal() -> Result<HashMap<String, OpenRouterModel>> {
+    let response = HTTP_CLIENT
+        .get("https://openrouter.ai/api/v1/models")
+        .send()
+        .await?;
+
+    if !response.status().is_success() {
+        anyhow::bail!(
+            "Failed to fetch OpenRouter models: HTTP {}",
+            response.status()
+        );
+    }
+
+    let models_response: OpenRouterModelsResponse = response.json().await?;
+
+    // Create a map for easy lookup
+    let mut pricing_map = HashMap::new();
+    for model in models_response.data {
+        pricing_map.insert(model.id.clone(), model);
+    }
+
+    Ok(pricing_map)
+}
+
+/// Initialize pricing cache on startup
+pub async fn initialize_pricing_cache() -> Result<()> {
+    PRICING_CACHE.initialize().await
+}
+
+/// Get pricing for a specific model
+pub async fn get_model_pricing(provider: &str, model: &str) -> Option<PricingInfo> {
+    PRICING_CACHE.get_model_pricing(provider, model).await
+}
+
+/// Force refresh pricing data
+pub async fn refresh_pricing() -> Result<()> {
+    PRICING_CACHE.refresh().await
+}
+
+/// Get all cached pricing data
+pub async fn get_all_pricing() -> HashMap<String, HashMap<String, PricingInfo>> {
+    let cache = PRICING_CACHE.memory_cache.read().await;
+    if let Some(cached) = &*cache {
+        cached.pricing.clone()
+    } else {
+        // Try loading from disk
+        if let Ok(Some(disk_cache)) = PRICING_CACHE.load_from_disk().await {
+            // Update memory cache
+            drop(cache);
+            let mut write_cache = PRICING_CACHE.memory_cache.write().await;
+            *write_cache = Some(disk_cache.clone());
+            disk_cache.pricing
+        } else {
+            HashMap::new()
+        }
+    }
+}
+
+/// Convert OpenRouter model ID to provider/model format
+/// e.g., "anthropic/claude-3.5-sonnet" -> ("anthropic", "claude-3.5-sonnet")
+pub fn parse_model_id(model_id: &str) -> Option<(String, String)> {
+    let parts: Vec<&str> = model_id.splitn(2, '/').collect();
+    if parts.len() == 2 {
+        // Normalize provider names to match our internal naming
+        let provider = match parts[0] {
+            "openai" => "openai",
+            "anthropic" => "anthropic",
+            "google" => "google",
+            "meta-llama" => "ollama", // Meta models often run via Ollama
+            "mistralai" => "mistral",
+            "cohere" => "cohere",
+            "perplexity" => "perplexity",
+            "deepseek" => "deepseek",
+            "groq" => "groq",
+            "nvidia" => "nvidia",
+            "microsoft" => "azure",
+            "replicate" => "replicate",
+            "huggingface" => "huggingface",
+            _ => parts[0],
+        };
+        Some((provider.to_string(), parts[1].to_string()))
+    } else {
+        None
+    }
+}
+
+/// Convert OpenRouter pricing to cost per token (already in that format)
+pub fn convert_pricing(price_str: &str) -> Option<f64> {
+    // OpenRouter prices are already in USD per token
+    price_str.parse::<f64>().ok()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_model_id() {
+        assert_eq!(
+            parse_model_id("anthropic/claude-3.5-sonnet"),
+            Some(("anthropic".to_string(), "claude-3.5-sonnet".to_string()))
+        );
+        assert_eq!(
+            parse_model_id("openai/gpt-4"),
+            Some(("openai".to_string(), "gpt-4".to_string()))
+        );
+        assert_eq!(parse_model_id("invalid-format"), None);
+    }
+
+    #[test]
+    fn test_convert_pricing() {
+        assert_eq!(convert_pricing("0.000003"), Some(0.000003));
+        assert_eq!(convert_pricing("0.015"), Some(0.015));
+        assert_eq!(convert_pricing("invalid"), None);
+    }
+}