diff --git a/README.md b/README.md index fd3619f..f572cca 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,7 @@ Splitrail is a **fast, cross-platform, real-time token usage tracker and cost mo - [GitHub Copilot CLI](https://github.com/features/copilot) - [OpenCode](https://github.com/sst/opencode) - [Pi Agent](https://github.com/badlogic/pi-mono/tree/main/packages/coding-agent) +- Grok Run one command to instantly review all of your CLI coding agent usage. Upload your usage data to your private account on the [Splitrail Cloud](https://splitrail.dev) for safe-keeping and cross-machine usage aggregation. From the team behind [ **Piebald.**](https://piebald.ai/) diff --git a/src/analyzers/grok.rs b/src/analyzers/grok.rs new file mode 100644 index 0000000..32f406c --- /dev/null +++ b/src/analyzers/grok.rs @@ -0,0 +1,642 @@ +use crate::analyzer::{Analyzer, DataSource}; +use crate::contribution_cache::ContributionStrategy; +use crate::models::calculate_total_cost_for_context_at; +use crate::types::{Application, ConversationMessage, MessageRole, Stats}; +use crate::utils::hash_text; +use anyhow::Result; +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use rayon::prelude::*; +use serde::Deserialize; +use simd_json::prelude::*; +use std::path::{Path, PathBuf}; +use walkdir::WalkDir; + +use super::copilot::count_tokens; + +pub struct GrokAnalyzer; + +impl GrokAnalyzer { + pub fn new() -> Self { + Self + } + + fn data_dir() -> Option { + dirs::home_dir().map(|home| home.join(".grok").join("sessions")) + } +} + +#[derive(Debug, Deserialize)] +struct GrokToolCall { + #[serde(default)] + name: String, +} + +#[derive(Debug, Deserialize)] +struct GrokChatRecord { + #[serde(rename = "type", default)] + record_type: String, + #[serde(default)] + content: Option, + #[serde(default)] + synthetic_reason: Option, + #[serde(default)] + model_id: Option, + #[serde(default)] + tool_calls: Option>, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +struct GrokUsage { + #[serde(default)] + input_tokens: u64, + #[serde(default)] + output_tokens: u64, + #[serde(default)] + cached_read_tokens: u64, + #[serde(default)] + cache_creation_tokens: u64, + #[serde(default)] + reasoning_tokens: u64, +} + +#[derive(Debug, Deserialize)] +struct GrokUpdateRecord { + #[serde(default)] + params: Option, +} + +#[derive(Debug, Deserialize)] +struct GrokUpdateParams { + #[serde(default)] + update: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct GrokSessionUpdate { + #[serde(default)] + session_update: Option, + #[serde(default)] + usage: Option, +} + +#[derive(Debug, Deserialize)] +struct GrokSessionSummary { + #[serde(default)] + created_at: Option>, + #[serde(default)] + generated_title: Option, + #[serde(default)] + session_summary: Option, +} + +fn is_grok_chat_path(path: &Path) -> bool { + path.is_file() + && path + .file_name() + .is_some_and(|name| name == "chat_history.jsonl") + && path + .parent() + .and_then(Path::parent) + .is_some_and(|project_dir| project_dir.parent().is_some()) +} + +fn truncate_session_name(text: &str) -> Option { + let text = text.trim(); + if text.is_empty() { + return None; + } + + let truncated: String = text.chars().take(50).collect(); + Some(if truncated.chars().count() < text.chars().count() { + format!("{truncated}...") + } else { + truncated + }) +} + +fn user_text_for_session_name(text: &str) -> Option { + let text = text.trim(); + if let Some(start) = text.find("") { + let query = &text[start + "".len()..]; + return query + .split_once("") + .map(|(query, _)| query.trim().to_string()); + } + + if text.starts_with("") || text.starts_with("") { + return None; + } + + Some(text.to_string()) +} + +fn text_from_content(content: &simd_json::OwnedValue) -> Option { + match content { + simd_json::OwnedValue::String(text) => Some(text.to_string()), + simd_json::OwnedValue::Array(items) => items.iter().find_map(|item| { + item.get("text") + .and_then(|text| text.as_str()) + .map(ToOwned::to_owned) + }), + simd_json::OwnedValue::Object(object) => object + .get("text") + .and_then(|text| text.as_str()) + .map(ToOwned::to_owned), + _ => None, + } +} + +fn extract_tool_stats(tool_calls: &[GrokToolCall]) -> Stats { + let mut stats = Stats { + tool_calls: tool_calls.len() as u32, + ..Default::default() + }; + + for tool_call in tool_calls { + match tool_call.name.as_str() { + "read_file" | "list_dir" | "list_directory" => stats.files_read += 1, + "grep" | "search" | "rg" => stats.file_content_searches += 1, + "run_terminal_command" | "run_terminal" => stats.terminal_commands += 1, + "write_file" | "create_file" => stats.files_added += 1, + "apply_patch" | "edit_file" => stats.files_edited += 1, + _ => {} + } + } + + stats +} + +fn parse_turn_usages(chat_history_path: &Path) -> Vec { + let Some(session_dir) = chat_history_path.parent() else { + return Vec::new(); + }; + let updates_path = session_dir.join("updates.jsonl"); + let Ok(content) = std::fs::read_to_string(&updates_path) else { + return Vec::new(); + }; + + content + .lines() + .enumerate() + .filter_map(|(line_index, line)| { + if line.trim().is_empty() { + return None; + } + + let mut bytes = line.as_bytes().to_vec(); + let record = match simd_json::from_slice::(&mut bytes) { + Ok(record) => record, + Err(error) => { + crate::utils::warn_once(format!( + "WARNING: Failed to parse Grok updates file `{}` line {}: {error}", + updates_path.display(), + line_index + 1 + )); + return None; + } + }; + let update = record.params?.update?; + if update.session_update.as_deref() != Some("turn_completed") { + return None; + } + Some(update.usage.unwrap_or_default()) + }) + .collect() +} + +fn distribute_u64(total: u64, indices: &[usize], weights: &[u64]) -> Vec { + if indices.is_empty() { + return Vec::new(); + } + + let total_weight: u128 = indices + .iter() + .map(|index| u128::from(weights[*index].max(1))) + .sum(); + let mut distributed = 0_u64; + + indices + .iter() + .enumerate() + .map(|(position, index)| { + if position + 1 == indices.len() { + return total.saturating_sub(distributed); + } + + let weight = u128::from(weights[*index].max(1)); + let share = ((u128::from(total) * weight) / total_weight) as u64; + distributed = distributed.saturating_add(share); + share + }) + .collect() +} + +fn distribute_f64(total: f64, indices: &[usize], weights: &[u64]) -> Vec { + if indices.is_empty() { + return Vec::new(); + } + + let total_weight: f64 = indices + .iter() + .map(|index| weights[*index].max(1) as f64) + .sum(); + let mut distributed = 0.0; + + indices + .iter() + .enumerate() + .map(|(position, index)| { + if position + 1 == indices.len() { + return total - distributed; + } + + let weight = weights[*index].max(1) as f64; + let share = total * weight / total_weight; + distributed += share; + share + }) + .collect() +} + +fn apply_turn_usage( + messages: &mut [ConversationMessage], + assistant_groups: &[Vec], + message_weights: &[u64], + usages: &[GrokUsage], +) { + for (usage_index, group) in assistant_groups + .iter() + .filter(|group| !group.is_empty()) + .enumerate() + { + let Some(usage) = usages.get(usage_index) else { + break; + }; + + let input_tokens = distribute_u64(usage.input_tokens, group, message_weights); + let output_tokens = distribute_u64(usage.output_tokens, group, message_weights); + let cached_read_tokens = distribute_u64(usage.cached_read_tokens, group, message_weights); + let cache_creation_tokens = + distribute_u64(usage.cache_creation_tokens, group, message_weights); + let reasoning_tokens = distribute_u64(usage.reasoning_tokens, group, message_weights); + // Grok's local CLI may report a discounted or subscription cost in + // costUsdTicks. Splitrail uses the public API standard price table so + // the value is comparable with the other API-based analyzers. + let model_index = group + .iter() + .copied() + .find(|index| messages[*index].model.is_some()); + let estimated_costs = model_index.and_then(|index| { + messages[index].model.as_deref().map(|model| { + let context_tokens = usage + .input_tokens + .saturating_add(usage.cached_read_tokens) + .saturating_add(usage.cache_creation_tokens); + let total_cost = calculate_total_cost_for_context_at( + model, + usage.input_tokens, + usage.output_tokens, + usage.cache_creation_tokens, + usage.cached_read_tokens, + context_tokens, + Some(messages[index].date), + ); + distribute_f64(total_cost, group, message_weights) + }) + }); + + for (position, index) in group.iter().enumerate() { + let stats = &mut messages[*index].stats; + stats.input_tokens = input_tokens[position]; + stats.output_tokens = output_tokens[position]; + stats.cache_read_tokens = cached_read_tokens[position]; + stats.cache_creation_tokens = cache_creation_tokens[position]; + stats.cached_tokens = stats.cache_read_tokens + stats.cache_creation_tokens; + stats.reasoning_tokens = reasoning_tokens[position]; + if let Some(estimated_costs) = &estimated_costs { + stats.cost = estimated_costs[position]; + } + } + } +} + +fn session_metadata(path: &Path) -> (DateTime, Option) { + let summary_path = path + .parent() + .map(|session_dir| session_dir.join("summary.json")); + + if let Some(summary_path) = summary_path + && let Ok(mut bytes) = std::fs::read(summary_path) + && let Ok(summary) = simd_json::from_slice::(&mut bytes) + { + let date = summary.created_at.unwrap_or_else(Utc::now); + return ( + date, + summary + .generated_title + .as_deref() + .and_then(truncate_session_name) + .or_else(|| { + summary + .session_summary + .as_deref() + .and_then(truncate_session_name) + }), + ); + } + + let date = path + .metadata() + .and_then(|metadata| metadata.modified()) + .map(DateTime::::from) + .unwrap_or_else(|_| Utc::now()); + (date, None) +} + +pub fn parse_chat_history_file(path: &Path) -> Result> { + let project_dir = path + .parent() + .and_then(Path::parent) + .and_then(Path::file_name) + .map(|name| name.to_string_lossy().into_owned()) + .unwrap_or_default(); + let project_hash = hash_text(&project_dir); + let conversation_hash = path + .parent() + .and_then(Path::file_name) + .map(|name| hash_text(&name.to_string_lossy())) + .unwrap_or_else(|| hash_text(&path.to_string_lossy())); + let file_path = path.to_string_lossy(); + let (date, mut session_name) = session_metadata(path); + let content = std::fs::read_to_string(path)?; + let mut messages = Vec::new(); + // Keep a leading group so assistant records before the first user record + // still receive usage. Empty groups are skipped when usage is applied. + let mut assistant_groups: Vec> = vec![Vec::new()]; + let mut message_weights = Vec::new(); + + for (line_index, line) in content.lines().enumerate() { + if line.trim().is_empty() { + continue; + } + + let mut bytes = line.as_bytes().to_vec(); + let record = match simd_json::from_slice::(&mut bytes) { + Ok(record) => record, + Err(error) => { + crate::utils::warn_once(format!( + "WARNING: Failed to parse Grok chat history `{file_path}` line {}: {error}", + line_index + 1 + )); + continue; + } + }; + + if record.synthetic_reason.is_some() { + continue; + } + + let (role, model, stats, weight) = match record.record_type.as_str() { + "user" => { + let user_text = record + .content + .as_ref() + .and_then(text_from_content) + .and_then(|text| user_text_for_session_name(&text)); + if user_text.is_none() { + continue; + } + if session_name.is_none() { + session_name = user_text.and_then(|text| truncate_session_name(&text)); + } + assistant_groups.push(Vec::new()); + (MessageRole::User, None, Stats::default(), 0) + } + "assistant" => { + let tool_calls = record.tool_calls.unwrap_or_default(); + let output_text = record + .content + .as_ref() + .and_then(text_from_content) + .unwrap_or_default(); + let weight = count_tokens(&output_text) + .saturating_add(tool_calls.len() as u64) + .max(1); + ( + MessageRole::Assistant, + record.model_id, + extract_tool_stats(&tool_calls), + weight, + ) + } + _ => continue, + }; + + let message_index = messages.len(); + if role == MessageRole::Assistant + && let Some(group) = assistant_groups.last_mut() + { + group.push(message_index); + } + message_weights.push(weight); + + messages.push(ConversationMessage { + application: Application::Grok, + date, + project_hash: project_hash.clone(), + conversation_hash: conversation_hash.clone(), + local_hash: Some(format!("{conversation_hash}:{line_index}")), + global_hash: hash_text(&format!("{file_path}:{line_index}")), + model, + stats, + role, + uuid: None, + session_name: session_name.clone(), + }); + } + + // `updates.jsonl` contains authoritative per-turn usage. Chat history does + // not carry token counts, so distribute each turn total across its model + // calls using the existing tokenizer as a stable weight. + let usages = parse_turn_usages(path); + apply_turn_usage(&mut messages, &assistant_groups, &message_weights, &usages); + + Ok(messages) +} + +#[async_trait] +impl Analyzer for GrokAnalyzer { + fn display_name(&self) -> &'static str { + "Grok" + } + + fn get_data_glob_patterns(&self) -> Vec { + Self::data_dir() + .map(|dir| format!("{}/*/*/chat_history.jsonl", dir.to_string_lossy())) + .into_iter() + .collect() + } + + fn discover_data_sources(&self) -> Result> { + let sources = Self::data_dir() + .filter(|dir| dir.is_dir()) + .into_iter() + .flat_map(|dir| WalkDir::new(dir).min_depth(3).max_depth(3).into_iter()) + .filter_map(|entry| entry.ok()) + .filter(|entry| is_grok_chat_path(entry.path())) + .map(|entry| DataSource { + path: entry.into_path(), + }) + .collect(); + + Ok(sources) + } + + fn is_available(&self) -> bool { + Self::data_dir() + .filter(|dir| dir.is_dir()) + .into_iter() + .flat_map(|dir| WalkDir::new(dir).min_depth(3).max_depth(3).into_iter()) + .filter_map(|entry| entry.ok()) + .any(|entry| is_grok_chat_path(entry.path())) + } + + fn parse_source(&self, source: &DataSource) -> Result> { + parse_chat_history_file(&source.path) + } + + fn parse_sources_parallel(&self, sources: &[DataSource]) -> Vec { + let messages: Vec<_> = sources + .par_iter() + .flat_map(|source| self.parse_source(source).unwrap_or_default()) + .collect(); + crate::utils::deduplicate_by_global_hash(messages) + } + + fn get_watch_directories(&self) -> Vec { + Self::data_dir() + .filter(|dir| dir.is_dir()) + .into_iter() + .collect() + } + + fn is_valid_data_path(&self, path: &Path) -> bool { + is_grok_chat_path(path) + } + + fn contribution_strategy(&self) -> ContributionStrategy { + ContributionStrategy::SingleSession + } + + fn requires_full_reload_for_source_change(&self) -> bool { + // Usage is stored beside chat_history.jsonl in updates.jsonl, so a + // change to either file can change the contribution for the session. + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn parses_real_messages_and_skips_synthetic_context() { + let dir = tempdir().expect("temporary directory should be created"); + let project_dir = dir.path().join("project"); + let session_dir = project_dir.join("session"); + std::fs::create_dir_all(&session_dir).expect("session directory should be created"); + std::fs::write( + session_dir.join("summary.json"), + r#"{"generated_title":"Grok title","session_summary":"Grok summary","created_at":"2026-08-01T12:00:00Z"}"#, + ) + .expect("summary should be written"); + std::fs::write( + session_dir.join("chat_history.jsonl"), + concat!( + r#"{"type":"user","synthetic_reason":"project_instructions","content":[{"type":"text","text":"ignored"}]}"#, "\n", + r#"{"type":"user","content":[{"type":"text","text":"Implement Grok support"}]}"#, "\n", + r#"{"type":"assistant","model_id":"grok-4.5","tool_calls":[{"name":"read_file"},{"name":"run_terminal_command"}]}"#, "\n", + ), + ) + .expect("chat history should be written"); + std::fs::write( + session_dir.join("updates.jsonl"), + r#"{"params":{"update":{"sessionUpdate":"turn_completed","usage":{"inputTokens":100,"outputTokens":20,"cachedReadTokens":30,"cacheCreationTokens":4,"reasoningTokens":5,"costUsdTicks":10000000000}}}}"#, + ) + .expect("updates should be written"); + + let messages = parse_chat_history_file(&session_dir.join("chat_history.jsonl")) + .expect("chat history should parse"); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].role, MessageRole::User); + assert_eq!(messages[1].role, MessageRole::Assistant); + assert_eq!(messages[1].model.as_deref(), Some("grok-4.5")); + assert_eq!(messages[1].stats.tool_calls, 2); + assert_eq!(messages[1].stats.files_read, 1); + assert_eq!(messages[1].stats.terminal_commands, 1); + assert_eq!(messages[1].stats.input_tokens, 100); + assert_eq!(messages[1].stats.output_tokens, 20); + assert_eq!(messages[1].stats.cache_read_tokens, 30); + assert_eq!(messages[1].stats.cache_creation_tokens, 4); + assert_eq!(messages[1].stats.reasoning_tokens, 5); + assert!((messages[1].stats.cost - 0.000329).abs() < f64::EPSILON); + assert_eq!(messages[0].session_name.as_deref(), Some("Grok title")); + } + + #[test] + fn aligns_usage_across_empty_groups_and_late_model_ids() { + let dir = tempdir().expect("temporary directory should be created"); + let project_dir = dir.path().join("project"); + let session_dir = project_dir.join("session"); + std::fs::create_dir_all(&session_dir).expect("session directory should be created"); + std::fs::write( + session_dir.join("chat_history.jsonl"), + concat!( + r#"{"type":"assistant","content":"leading"}"#, + "\n", + r#"{"type":"user","content":"first"}"#, + "\n", + r#"{"type":"assistant","content":"unlabeled"}"#, + "\n", + r#"{"type":"assistant","model_id":"grok-4.5","content":"labeled"}"#, + "\n", + r#"{"type":"user","content":"empty turn"}"#, + "\n", + r#"{"type":"user","content":"second"}"#, + "\n", + r#"{"type":"assistant","model_id":"grok-4.5","content":"final"}"#, + "\n", + ), + ) + .expect("chat history should be written"); + std::fs::write( + session_dir.join("updates.jsonl"), + concat!( + r#"{"params":{"update":{"sessionUpdate":"turn_completed","usage":{"inputTokens":100,"outputTokens":10}}}}"#, "\n", + r#"{"params":{"update":{"sessionUpdate":"turn_completed","usage":{"inputTokens":200,"outputTokens":20}}}}"#, "\n", + r#"{"params":{"update":{"sessionUpdate":"turn_completed","usage":{"inputTokens":300,"outputTokens":30}}}}"#, + ), + ) + .expect("updates should be written"); + + let messages = parse_chat_history_file(&session_dir.join("chat_history.jsonl")) + .expect("chat history should parse"); + let assistants: Vec<_> = messages + .iter() + .filter(|message| message.role == MessageRole::Assistant) + .collect(); + + assert_eq!(assistants.len(), 4); + assert_eq!(assistants[0].stats.input_tokens, 100); + assert_eq!( + assistants[1].stats.input_tokens + assistants[2].stats.input_tokens, + 200 + ); + assert!(assistants[1].stats.cost > 0.0); + assert_eq!(assistants[3].stats.input_tokens, 300); + } +} diff --git a/src/analyzers/mod.rs b/src/analyzers/mod.rs index b9be065..d571a03 100644 --- a/src/analyzers/mod.rs +++ b/src/analyzers/mod.rs @@ -6,6 +6,7 @@ pub mod codex_cli; pub mod copilot; pub mod copilot_cli; pub mod gemini_cli; +pub mod grok; pub mod kilo_cli; pub mod kilo_code; pub mod opencode; @@ -23,6 +24,7 @@ pub use codex_cli::CodexCliAnalyzer; pub use copilot::CopilotAnalyzer; pub use copilot_cli::CopilotCliAnalyzer; pub use gemini_cli::GeminiCliAnalyzer; +pub use grok::GrokAnalyzer; pub use kilo_cli::KiloCliAnalyzer; pub use kilo_code::KiloCodeAnalyzer; pub use opencode::OpenCodeAnalyzer; diff --git a/src/main.rs b/src/main.rs index f2b3634..b9afd48 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,8 +6,9 @@ use std::sync::Arc; use analyzer::AnalyzerRegistry; use analyzers::{ AntigravityCliAnalyzer, ClaudeCodeAnalyzer, ClineAnalyzer, CodexCliAnalyzer, CopilotAnalyzer, - CopilotCliAnalyzer, GeminiCliAnalyzer, KiloCliAnalyzer, KiloCodeAnalyzer, OpenCodeAnalyzer, - PiAgentAnalyzer, PiebaldAnalyzer, QwenCodeAnalyzer, RooCodeAnalyzer, ZooCodeAnalyzer, + CopilotCliAnalyzer, GeminiCliAnalyzer, GrokAnalyzer, KiloCliAnalyzer, KiloCodeAnalyzer, + OpenCodeAnalyzer, PiAgentAnalyzer, PiebaldAnalyzer, QwenCodeAnalyzer, RooCodeAnalyzer, + ZooCodeAnalyzer, }; mod analyzer; @@ -212,6 +213,7 @@ pub fn create_analyzer_registry() -> AnalyzerRegistry { registry.register(PiAgentAnalyzer::new()); registry.register(PiebaldAnalyzer::new()); registry.register(AntigravityCliAnalyzer::new()); + registry.register(GrokAnalyzer::new()); registry } diff --git a/src/models.rs b/src/models.rs index 512a314..cf2c7ab 100644 --- a/src/models.rs +++ b/src/models.rs @@ -1535,15 +1535,69 @@ fn populate_defaults( ); // xAI Models + // Source: https://docs.x.ai/developers/pricing add_model!( - "grok-code-fast-1", - PricingStructure::Flat { - input_per_1m: 0.20, - output_per_1m: 1.50 - }, - CachingSupport::OpenAI { - cached_input_per_1m: 0.02 - }, + "grok-4.5", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ + PricingTier { + max_tokens: Some(200_000), + input_per_1m: 2.00, + output_per_1m: 6.00, + }, + PricingTier { + max_tokens: None, + input_per_1m: 4.00, + output_per_1m: 12.00, + }, + ], + bracket_pricing: true, + }), + CachingSupport::Tiered(TieredCaching { + tiers: vec![ + CachingTier { + max_tokens: Some(200_000), + cached_input_per_1m: 0.30, + }, + CachingTier { + max_tokens: None, + cached_input_per_1m: 0.60, + }, + ], + bracket_pricing: true, + }), + false + ); + add_model!( + "grok-build-0.1", + PricingStructure::Tiered(TieredPricing { + tiers: vec![ + PricingTier { + max_tokens: Some(200_000), + input_per_1m: 1.00, + output_per_1m: 2.00, + }, + PricingTier { + max_tokens: None, + input_per_1m: 2.00, + output_per_1m: 4.00, + }, + ], + bracket_pricing: true, + }), + CachingSupport::Tiered(TieredCaching { + tiers: vec![ + CachingTier { + max_tokens: Some(200_000), + cached_input_per_1m: 0.20, + }, + CachingTier { + max_tokens: None, + cached_input_per_1m: 0.40, + }, + ], + bracket_pricing: true, + }), false ); @@ -2075,6 +2129,11 @@ fn populate_defaults( // Aurora aliases add_alias!("aurora-alpha", "aurora-alpha"); + + // xAI aliases + add_alias!("grok-code-fast-1", "grok-build-0.1"); + add_alias!("grok-code-fast", "grok-build-0.1"); + add_alias!("grok-code-fast-1-0825", "grok-build-0.1"); } /// Free-tier model pricing for models accessed via OpenRouter's `:free` suffix @@ -2324,6 +2383,8 @@ fn cache_cost_for_caching( creation_cost + read_cost } CachingSupport::Tiered(tiered) => { + // Tiered caching models currently publish cached-read rates only; + // cache creation tokens are intentionally not charged here. calculate_tiered_cache_cost(cache_read_tokens, &tiered.tiers, tiered.bracket_pricing) } } @@ -2450,6 +2511,39 @@ pub fn calculate_total_cost_for_service_tier_at( } } +/// Calculate standard cost when a model's tiers are selected by total prompt +/// context rather than by each token category independently. +pub fn calculate_total_cost_for_context_at( + model_name: &str, + input_tokens: u64, + output_tokens: u64, + cache_creation_tokens: u64, + cache_read_tokens: u64, + context_tokens: u64, + effective_at: Option>, +) -> f64 { + match get_model_info(model_name) { + Some(model_info) => { + let (pricing, caching) = standard_pricing_for_date(&model_info, effective_at); + calculate_context_cost( + pricing, + caching, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + context_tokens, + ) + } + None => { + warn_once(format!( + "WARNING: Unknown model: {model_name}. Defaulting to $0." + )); + 0.0 + } + } +} + fn calculate_tiered_cost( tokens: u64, tiers: &[PricingTier], @@ -2543,6 +2637,47 @@ where None } +fn calculate_context_cost( + pricing: &PricingStructure, + caching: &CachingSupport, + input_tokens: u64, + output_tokens: u64, + cache_creation_tokens: u64, + cache_read_tokens: u64, + context_tokens: u64, +) -> f64 { + let token_cost = match pricing { + PricingStructure::Flat { + input_per_1m, + output_per_1m, + } => { + (input_tokens as f64 / 1_000_000.0) * input_per_1m + + (output_tokens as f64 / 1_000_000.0) * output_per_1m + } + PricingStructure::Tiered(tiered) => { + find_tier(context_tokens, &tiered.tiers, |tier| tier.max_tokens) + .map(|tier| { + (input_tokens as f64 / 1_000_000.0) * tier.input_per_1m + + (output_tokens as f64 / 1_000_000.0) * tier.output_per_1m + }) + .unwrap_or(0.0) + } + }; + + let cache_cost = match caching { + CachingSupport::Tiered(tiered) => { + // Tiered caching models currently publish cached-read rates only; + // cache creation tokens are intentionally not charged here. + find_tier(context_tokens, &tiered.tiers, |tier| tier.max_tokens) + .map(|tier| (cache_read_tokens as f64 / 1_000_000.0) * tier.cached_input_per_1m) + .unwrap_or(0.0) + } + _ => cache_cost_for_caching(caching, cache_creation_tokens, cache_read_tokens), + }; + + token_cost + cache_cost +} + #[cfg(test)] mod tests { use super::{ @@ -2552,8 +2687,8 @@ mod tests { calculate_input_cost, calculate_input_cost_for_service_tier, calculate_input_cost_for_service_tier_at, calculate_output_cost, calculate_output_cost_for_service_tier, calculate_output_cost_for_service_tier_at, - calculate_total_cost_for_service_tier_at, get_model_info, get_registry_lock, - init_external_models, + calculate_total_cost_for_context_at, calculate_total_cost_for_service_tier_at, + get_model_info, get_registry_lock, init_external_models, }; use chrono::{TimeZone, Utc}; @@ -3362,6 +3497,42 @@ mod tests { approx_eq(cache_cost, 37.5); } + #[test] + fn xai_standard_pricing_uses_context_tiers_and_aliases() { + assert!( + !get_model_info("grok-4.5") + .expect("Grok 4.5 should exist") + .is_estimated + ); + + approx_eq( + calculate_total_cost_for_context_at( + "grok-4.5", 1_000_000, 1_000_000, 0, 1_000_000, 199_999, None, + ), + 8.3, + ); + approx_eq( + calculate_total_cost_for_context_at( + "grok-code-fast-1", + 1_000_000, + 1_000_000, + 0, + 1_000_000, + 200_001, + None, + ), + 6.4, + ); + approx_eq( + calculate_total_cost_for_context_at( + "grok-4.5", 1_000_000, 1_000_000, 999_999, 1_000_000, 2_000_000, None, + ), + calculate_total_cost_for_context_at( + "grok-4.5", 1_000_000, 1_000_000, 0, 1_000_000, 2_000_000, None, + ), + ); + } + #[test] fn doubao_seed_code_alias_resolves() { let model_info = get_model_info("doubao-seed-code").expect("alias should resolve"); diff --git a/src/types.rs b/src/types.rs index 1c67d44..cfceb64 100644 --- a/src/types.rs +++ b/src/types.rs @@ -233,6 +233,7 @@ pub enum Application { PiAgent, Piebald, AntigravityCli, + Grok, } #[derive(Debug, Clone, Serialize, Deserialize)]