diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index c99ddf454..1e19c5fa3 100644 --- a/crates/goose-cli/src/cli.rs +++ b/crates/goose-cli/src/cli.rs @@ -1559,7 +1559,7 @@ async fn log_session_completion( let (total_tokens, message_count) = session .get_session() .await - .map(|m| (m.total_tokens.unwrap_or(0), m.message_count)) + .map(|m| (m.usage.total_tokens.unwrap_or(0), m.message_count)) .unwrap_or((0, 0)); tracing::info!( diff --git a/crates/goose-cli/src/commands/term.rs b/crates/goose-cli/src/commands/term.rs index 17ecdfcb6..26d2a3652 100644 --- a/crates/goose-cli/src/commands/term.rs +++ b/crates/goose-cli/src/commands/term.rs @@ -324,7 +324,10 @@ pub async fn handle_term_info() -> Result<()> { let session_manager = SessionManager::instance(); let session = session_manager.get_session(&session_id, false).await.ok(); - let total_tokens = session.as_ref().and_then(|s| s.total_tokens).unwrap_or(0) as usize; + let total_tokens = session + .as_ref() + .and_then(|s| s.usage.total_tokens) + .unwrap_or(0) as usize; let config = goose::config::Config::global(); let model_name = config diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index b3027901b..76d7ac915 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -909,9 +909,11 @@ impl CliSession { .config .session_manager .update(&self.session_id) - .total_tokens(Some(0)) - .input_tokens(Some(0)) - .output_tokens(Some(0)) + .usage(goose_providers::conversation::token_usage::Usage::new( + Some(0), + Some(0), + Some(0), + )) .apply() .await { @@ -1382,9 +1384,18 @@ impl CliSession { .await { Ok(session) => JsonMetadata { - total_tokens: session.accumulated_total_tokens.or(session.total_tokens), - input_tokens: session.accumulated_input_tokens.or(session.input_tokens), - output_tokens: session.accumulated_output_tokens.or(session.output_tokens), + total_tokens: session + .accumulated_usage + .total_tokens + .or(session.usage.total_tokens), + input_tokens: session + .accumulated_usage + .input_tokens + .or(session.usage.input_tokens), + output_tokens: session + .accumulated_usage + .output_tokens + .or(session.usage.output_tokens), status: "completed".to_string(), }, Err(_) => JsonMetadata { @@ -1409,9 +1420,9 @@ impl CliSession { .ok(); let (total_tokens, input_tokens, output_tokens) = match session { Some(s) => ( - s.accumulated_total_tokens.or(s.total_tokens), - s.accumulated_input_tokens.or(s.input_tokens), - s.accumulated_output_tokens.or(s.output_tokens), + s.accumulated_usage.total_tokens.or(s.usage.total_tokens), + s.accumulated_usage.input_tokens.or(s.usage.input_tokens), + s.accumulated_usage.output_tokens.or(s.usage.output_tokens), ), None => (None, None, None), }; @@ -1579,7 +1590,7 @@ impl CliSession { pub async fn get_total_token_usage(&self) -> Result> { let metadata = self.get_session().await?; - Ok(metadata.accumulated_total_tokens) + Ok(metadata.accumulated_usage.total_tokens) } /// Display enhanced context usage with session totals @@ -1599,18 +1610,15 @@ impl CliSession { match self.get_session().await { Ok(metadata) => { - let total_tokens = metadata.total_tokens.unwrap_or(0) as usize; + let total_tokens = metadata.usage.total_tokens.unwrap_or(0) as usize; output::display_context_usage(total_tokens, context_limit); if show_cost { - let input_tokens = metadata.input_tokens.unwrap_or(0) as usize; - let output_tokens = metadata.output_tokens.unwrap_or(0) as usize; output::display_cost_usage( &provider_name, &model_config.model_name, - input_tokens, - output_tokens, + &metadata.usage, ); } } diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index b0c71e2b4..7d1b56ea5 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -10,6 +10,7 @@ use goose::providers::canonical::maybe_get_canonical_model; #[cfg(target_os = "windows")] use goose::subprocess::SubprocessExt; use goose::utils::safe_truncate; +use goose_providers::conversation::token_usage::Usage; use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; use rmcp::model::{CallToolRequestParams, JsonObject, PromptArgument}; use serde_json::Value; @@ -1414,31 +1415,33 @@ pub fn display_context_usage(total_tokens: usize, context_limit: usize) { ); } -fn estimate_cost_usd( - provider: &str, - model: &str, - input_tokens: usize, - output_tokens: usize, -) -> Option { +fn estimate_cost_usd(provider: &str, model: &str, usage: &Usage) -> Option { let canonical_model = maybe_get_canonical_model(provider, model)?; - - let input_cost_per_token = canonical_model.cost.input? / 1_000_000.0; - let output_cost_per_token = canonical_model.cost.output? / 1_000_000.0; - - let input_cost = input_cost_per_token * input_tokens as f64; - let output_cost = output_cost_per_token * output_tokens as f64; - Some(input_cost + output_cost) + canonical_model.cost.estimate_cost(usage) } /// Display cost information, if price data is available. -pub fn display_cost_usage(provider: &str, model: &str, input_tokens: usize, output_tokens: usize) { - if let Some(cost) = estimate_cost_usd(provider, model, input_tokens, output_tokens) { +pub fn display_cost_usage(provider: &str, model: &str, usage: &Usage) { + if let Some(cost) = estimate_cost_usd(provider, model, usage) { use console::style; + let input_tokens = usage.input_tokens.unwrap_or(0); + let output_tokens = usage.output_tokens.unwrap_or(0); + let cache_read = usage.cache_read_input_tokens.unwrap_or(0); + let cache_write = usage.cache_write_input_tokens.unwrap_or(0); + + let cache_breakdown = match (cache_read, cache_write) { + (0, 0) => String::new(), + (read, 0) => format!(" ({} cache read)", read), + (0, write) => format!(" ({} cache write)", write), + (read, write) => format!(" ({} cache read, {} cache write)", read, write), + }; + eprintln!( - "Cost: {} USD ({} tokens: in {}, out {})", + "Cost: {} USD ({} tokens: in {}{}, out {})", style(format!("${:.4}", cost)).cyan(), input_tokens + output_tokens, input_tokens, + cache_breakdown, output_tokens ); } diff --git a/crates/goose-providers/src/canonical/model.rs b/crates/goose-providers/src/canonical/model.rs index fa08dc009..ad287d68c 100644 --- a/crates/goose-providers/src/canonical/model.rs +++ b/crates/goose-providers/src/canonical/model.rs @@ -1,5 +1,7 @@ use serde::{Deserialize, Serialize}; +use crate::conversation::token_usage::Usage; + /// Modality types for model input/output #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] @@ -52,6 +54,30 @@ pub struct Pricing { pub cache_write: Option, } +impl Pricing { + pub fn estimate_cost(&self, usage: &Usage) -> Option { + let input_price = self.input?; + let output_price = self.output?; + let cache_read_price = self.cache_read.unwrap_or(input_price); + let cache_write_price = self.cache_write.unwrap_or(input_price); + + let input_tokens = usage.input_tokens.unwrap_or(0).max(0) as f64; + let output_tokens = usage.output_tokens.unwrap_or(0).max(0) as f64; + let cache_read_tokens = usage.cache_read_input_tokens.unwrap_or(0).max(0) as f64; + let cache_write_tokens = usage.cache_write_input_tokens.unwrap_or(0).max(0) as f64; + let uncached_input_tokens = + (input_tokens - cache_read_tokens - cache_write_tokens).max(0.0); + + Some( + (uncached_input_tokens * input_price + + cache_read_tokens * cache_read_price + + cache_write_tokens * cache_write_price + + output_tokens * output_price) + / 1_000_000.0, + ) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct Limit { /// Maximum context window size in tokens @@ -130,3 +156,64 @@ pub struct CanonicalModel { #[serde(default)] pub limit: Limit, } + +#[cfg(test)] +mod tests { + use super::*; + + fn pricing( + input: Option, + output: Option, + cache_read: Option, + cache_write: Option, + ) -> Pricing { + Pricing { + input, + output, + cache_read, + cache_write, + } + } + + #[test] + fn estimate_cost_prices_cache_tokens_at_cache_rates() { + let pricing = pricing(Some(5.0), Some(25.0), Some(0.5), Some(6.25)); + let usage = + Usage::new(Some(10_000), Some(1_000), None).with_cache_tokens(Some(8_000), Some(1_000)); + + let cost = pricing.estimate_cost(&usage).unwrap(); + let expected = + (1_000.0 * 5.0 + 8_000.0 * 0.5 + 1_000.0 * 6.25 + 1_000.0 * 25.0) / 1_000_000.0; + assert!((cost - expected).abs() < f64::EPSILON); + } + + #[test] + fn estimate_cost_handles_missing_prices() { + let usage = + Usage::new(Some(1_000), Some(100), None).with_cache_tokens(Some(600), Some(200)); + + // Unpriced cache tokens fall back to the input rate. + let cost = pricing(Some(2.0), Some(10.0), None, None) + .estimate_cost(&usage) + .unwrap(); + assert_eq!(cost, (1_000.0 * 2.0 + 100.0 * 10.0) / 1_000_000.0); + + // Missing input or output pricing means no estimate at all. + assert!(pricing(None, Some(10.0), None, None) + .estimate_cost(&usage) + .is_none()); + assert!(pricing(Some(2.0), None, None, None) + .estimate_cost(&usage) + .is_none()); + } + + #[test] + fn estimate_cost_clamps_cache_tokens_exceeding_input() { + let pricing = pricing(Some(5.0), Some(25.0), Some(0.5), Some(6.25)); + let usage = Usage::new(Some(100), Some(10), None).with_cache_tokens(Some(150), Some(50)); + + let cost = pricing.estimate_cost(&usage).unwrap(); + let expected = (150.0 * 0.5 + 50.0 * 6.25 + 10.0 * 25.0) / 1_000_000.0; + assert!((cost - expected).abs() < f64::EPSILON); + } +} diff --git a/crates/goose-providers/src/conversation/message.rs b/crates/goose-providers/src/conversation/message.rs index bf6d2d5cf..38c3a5006 100644 --- a/crates/goose-providers/src/conversation/message.rs +++ b/crates/goose-providers/src/conversation/message.rs @@ -1088,9 +1088,17 @@ pub struct TokenState { pub input_tokens: i32, pub output_tokens: i32, pub total_tokens: i32, + #[serde(default)] + pub cache_read_tokens: i32, + #[serde(default)] + pub cache_write_tokens: i32, pub accumulated_input_tokens: i32, pub accumulated_output_tokens: i32, pub accumulated_total_tokens: i32, + #[serde(default)] + pub accumulated_cache_read_tokens: i32, + #[serde(default)] + pub accumulated_cache_write_tokens: i32, pub accumulated_cost: Option, } diff --git a/crates/goose-providers/src/conversation/token_usage.rs b/crates/goose-providers/src/conversation/token_usage.rs index 7d481fc30..d62a4d370 100644 --- a/crates/goose-providers/src/conversation/token_usage.rs +++ b/crates/goose-providers/src/conversation/token_usage.rs @@ -1,6 +1,7 @@ use std::ops::{Add, AddAssign}; use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProviderUsage { @@ -53,8 +54,14 @@ impl ProviderUsage { } } -#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy)] +/// `input_tokens` is the total input including cache read/write tokens; +/// the cache fields are breakdown subsets of it. Parsers for providers +/// that report cache tokens separately from input (e.g. Anthropic, +/// Bedrock) must fold them into `input_tokens`. +#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy, PartialEq, Eq, ToSchema)] pub struct Usage { + /// All prompt tokens, including any served from or written to cache. + /// `cache_read_input_tokens` and `cache_write_input_tokens` are subsets of this. pub input_tokens: Option, pub output_tokens: Option, pub total_tokens: Option, @@ -64,12 +71,12 @@ pub struct Usage { fn sum_optionals(a: Option, b: Option) -> Option where - T: Add + Default, + T: Add, { match (a, b) { (Some(x), Some(y)) => Some(x + y), - (Some(x), None) => Some(x + T::default()), - (None, Some(y)) => Some(T::default() + y), + (Some(x), None) => Some(x), + (None, Some(y)) => Some(y), (None, None) => None, } } @@ -107,7 +114,7 @@ impl Usage { ) -> Self { let calculated_total = if total_tokens.is_none() { match (input_tokens, output_tokens) { - (Some(input), Some(output)) => Some(input + output), + (Some(input), Some(output)) => Some(input.saturating_add(output)), (Some(input), None) => Some(input), (None, Some(output)) => Some(output), (None, None) => None, @@ -134,6 +141,26 @@ impl Usage { self.cache_write_input_tokens = cache_write_input_tokens; self } + + /// For providers whose reported `input_tokens`/`total_tokens` exclude + /// cache tokens (e.g. Anthropic, Bedrock): folds the cache breakdown in. + pub fn from_cache_exclusive_input( + input_tokens: Option, + output_tokens: Option, + total_tokens: Option, + cache_read_input_tokens: Option, + cache_write_input_tokens: Option, + ) -> Self { + let cache_tokens = cache_read_input_tokens + .unwrap_or(0) + .saturating_add(cache_write_input_tokens.unwrap_or(0)); + Self::new( + input_tokens.map(|v| v.saturating_add(cache_tokens)), + output_tokens, + total_tokens.map(|v| v.saturating_add(cache_tokens)), + ) + .with_cache_tokens(cache_read_input_tokens, cache_write_input_tokens) + } } #[cfg(test)] @@ -161,6 +188,18 @@ mod tests { Ok(()) } + #[test] + fn test_from_cache_exclusive_input_folds_cache_into_input_and_total() { + let usage = + Usage::from_cache_exclusive_input(Some(10), Some(50), Some(60), Some(5000), Some(1000)); + + assert_eq!(usage.input_tokens, Some(6010)); + assert_eq!(usage.output_tokens, Some(50)); + assert_eq!(usage.total_tokens, Some(6060)); + assert_eq!(usage.cache_read_input_tokens, Some(5000)); + assert_eq!(usage.cache_write_input_tokens, Some(1000)); + } + #[test] fn test_usage_addition_includes_cached_tokens() { let usage_a = diff --git a/crates/goose-providers/src/formats/openai.rs b/crates/goose-providers/src/formats/openai.rs index 836a01212..a824e56a6 100644 --- a/crates/goose-providers/src/formats/openai.rs +++ b/crates/goose-providers/src/formats/openai.rs @@ -2441,6 +2441,41 @@ mod tests { assert_eq!(usage.usage.total_tokens, Some(expected_total)); } + #[test] + fn test_get_usage_preserves_provider_totals_with_cache_fields() { + let usage = get_usage(&json!({ + "prompt_tokens": 120, + "completion_tokens": 30, + "total_tokens": 150, + "cache_read_input_tokens": 80, + "cache_creation_input_tokens": 20 + })); + + assert_eq!(usage.input_tokens, Some(120)); + assert_eq!(usage.output_tokens, Some(30)); + assert_eq!(usage.total_tokens, Some(150)); + assert_eq!(usage.cache_read_input_tokens, Some(80)); + assert_eq!(usage.cache_write_input_tokens, Some(20)); + } + + #[test] + fn test_get_usage_reads_openai_prompt_tokens_details() { + let usage = get_usage(&json!({ + "prompt_tokens": 120, + "completion_tokens": 30, + "total_tokens": 150, + "prompt_tokens_details": { + "cached_tokens": 80 + } + })); + + assert_eq!(usage.input_tokens, Some(120)); + assert_eq!(usage.output_tokens, Some(30)); + assert_eq!(usage.total_tokens, Some(150)); + assert_eq!(usage.cache_read_input_tokens, Some(80)); + assert_eq!(usage.cache_write_input_tokens, None); + } + #[test] fn test_get_usage_reads_nested_usage_with_cache_fields() { let usage = get_usage(&json!({ diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index 272f5a385..81eadc67d 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -3,6 +3,7 @@ use goose::agents::extension::ToolInfo; use goose::agents::ExtensionConfig; use goose::config::permission::PermissionLevel; use goose::config::ExtensionEntry; +use goose::conversation::token_usage::Usage; use goose::conversation::Conversation; use goose::download_manager::{DownloadProgress, DownloadStatus}; use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata, ProviderType}; @@ -531,6 +532,7 @@ derive_utoipa!(IconTheme as IconThemeSchema); MessageMetadata, InferenceMetadata, TokenState, + Usage, ContentSchema, EmbeddedResourceSchema, ImageContentSchema, diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index bfec2e001..dbc6ed5ab 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -157,15 +157,7 @@ pub async fn get_token_state(session_manager: &SessionManager, session_id: &str) session_manager .get_session(session_id, false) .await - .map(|session| TokenState { - input_tokens: session.input_tokens.unwrap_or(0), - output_tokens: session.output_tokens.unwrap_or(0), - total_tokens: session.total_tokens.unwrap_or(0), - accumulated_input_tokens: session.accumulated_input_tokens.unwrap_or(0), - accumulated_output_tokens: session.accumulated_output_tokens.unwrap_or(0), - accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0), - accumulated_cost: session.accumulated_cost, - }) + .map(|session| TokenState::from(&session)) .inspect_err(|e| { tracing::warn!( "Failed to fetch session token state for {}: {}", @@ -397,7 +389,7 @@ pub async fn reply( let session_duration = session_start.elapsed(); if let Ok(session) = state.session_manager().get_session(&session_id, true).await { - let total_tokens = session.total_tokens.unwrap_or(0); + let total_tokens = session.usage.total_tokens.unwrap_or(0); tracing::info!( monotonic_counter.goose.session_completions = 1, session_type = "app", diff --git a/crates/goose-server/src/routes/schedule.rs b/crates/goose-server/src/routes/schedule.rs index afdc806b0..38f4d53ec 100644 --- a/crates/goose-server/src/routes/schedule.rs +++ b/crates/goose-server/src/routes/schedule.rs @@ -347,12 +347,12 @@ async fn sessions_handler( working_dir: session.working_dir.to_string_lossy().into_owned(), schedule_id: session.schedule_id, message_count: session.message_count, - total_tokens: session.total_tokens, - input_tokens: session.input_tokens, - output_tokens: session.output_tokens, - accumulated_total_tokens: session.accumulated_total_tokens, - accumulated_input_tokens: session.accumulated_input_tokens, - accumulated_output_tokens: session.accumulated_output_tokens, + total_tokens: session.usage.total_tokens, + input_tokens: session.usage.input_tokens, + output_tokens: session.usage.output_tokens, + accumulated_total_tokens: session.accumulated_usage.total_tokens, + accumulated_input_tokens: session.accumulated_usage.input_tokens, + accumulated_output_tokens: session.accumulated_usage.output_tokens, }); } Ok(Json(display_infos)) diff --git a/crates/goose-server/src/routes/session_events.rs b/crates/goose-server/src/routes/session_events.rs index c744a6443..8dc76d2a5 100644 --- a/crates/goose-server/src/routes/session_events.rs +++ b/crates/goose-server/src/routes/session_events.rs @@ -532,7 +532,7 @@ pub async fn session_reply( .get_session(&task_session_id, true) .await { - let total_tokens = session.total_tokens.unwrap_or(0); + let total_tokens = session.usage.total_tokens.unwrap_or(0); tracing::info!( monotonic_counter.goose.session_completions = 1, session_type = "app", diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 1133c3cdf..eb0d185a3 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -796,9 +796,9 @@ fn to_nonnegative_u64(value: Option) -> Option { } fn build_prompt_usage(session: &Session) -> Option { - let total = to_nonnegative_u64(session.total_tokens)?; - let input = to_nonnegative_u64(session.input_tokens).unwrap_or(0); - let output = to_nonnegative_u64(session.output_tokens).unwrap_or(0); + let total = to_nonnegative_u64(session.usage.total_tokens)?; + let input = to_nonnegative_u64(session.usage.input_tokens).unwrap_or(0); + let output = to_nonnegative_u64(session.usage.output_tokens).unwrap_or(0); Some(Usage::new(total, input, output)) } @@ -808,12 +808,12 @@ pub(super) struct UsageUpdates { } pub(super) fn build_usage_updates(session: &Session) -> Option { - let used = session.total_tokens.unwrap_or(0).max(0) as u64; + let used = session.usage.total_tokens.unwrap_or(0).max(0) as u64; let ctx_limit = session.model_config.as_ref()?.context_limit() as u64; let accumulated_input_tokens = - to_nonnegative_u64(session.accumulated_input_tokens).unwrap_or(0); + to_nonnegative_u64(session.accumulated_usage.input_tokens).unwrap_or(0); let accumulated_output_tokens = - to_nonnegative_u64(session.accumulated_output_tokens).unwrap_or(0); + to_nonnegative_u64(session.accumulated_usage.output_tokens).unwrap_or(0); Some(UsageUpdates { custom: GooseSessionNotification { session_id: session.id.clone(), @@ -2949,6 +2949,7 @@ mod tests { EnvVariable, HttpHeader, McpServer, McpServerHttp, McpServerSse, McpServerStdio, PermissionOptionId, ResourceLink, SelectedPermissionOutcome, }; + use goose_providers::conversation::token_usage::Usage as TokenUsage; use rmcp::model::{CallToolRequestParams, Content as RmcpContent}; use std::io::Write; use std::path::PathBuf; @@ -3684,53 +3685,23 @@ print(\"hello, world\") ); } - fn make_session_with_usage( - total_tokens: Option, - input_tokens: Option, - output_tokens: Option, - accumulated_total_tokens: Option, - accumulated_input_tokens: Option, - accumulated_output_tokens: Option, - ) -> Session { + fn make_session_with_usage(usage: TokenUsage, accumulated_usage: TokenUsage) -> Session { Session { id: "session-1".to_string(), working_dir: PathBuf::from("/tmp"), name: "ACP Session".to_string(), - user_set_name: false, session_type: SessionType::Acp, - created_at: Default::default(), - updated_at: Default::default(), - extension_data: crate::session::ExtensionData::default(), - total_tokens, - input_tokens, - output_tokens, - accumulated_total_tokens, - accumulated_input_tokens, - accumulated_output_tokens, - accumulated_cost: None, - schedule_id: None, - recipe: None, - user_recipe_values: None, - conversation: None, - message_count: 0, - provider_name: None, - model_config: None, - goose_mode: GooseMode::default(), - archived_at: None, - project_id: None, - last_message_snippet: None, + usage, + accumulated_usage, + ..Default::default() } } #[test] fn test_build_prompt_usage_uses_current_turn_tokens() { let session = make_session_with_usage( - Some(120), - Some(80), - Some(40), - Some(360), - Some(210), - Some(150), + TokenUsage::new(Some(80), Some(40), Some(120)), + TokenUsage::new(Some(210), Some(150), Some(360)), ); let usage = build_prompt_usage(&session).expect("usage should be present"); assert_eq!(usage.total_tokens, 120); @@ -3740,7 +3711,10 @@ print(\"hello, world\") #[test] fn test_build_prompt_usage_falls_back_to_current_tokens() { - let session = make_session_with_usage(Some(120), Some(80), Some(40), None, None, None); + let session = make_session_with_usage( + TokenUsage::new(Some(80), Some(40), Some(120)), + TokenUsage::default(), + ); let usage = build_prompt_usage(&session).expect("usage should be present"); assert_eq!(usage.total_tokens, 120); assert_eq!(usage.input_tokens, 80); @@ -3749,13 +3723,24 @@ print(\"hello, world\") #[test] fn test_build_prompt_usage_requires_total_tokens() { - let session = make_session_with_usage(None, Some(80), Some(40), None, None, None); + let session = make_session_with_usage( + TokenUsage { + input_tokens: Some(80), + output_tokens: Some(40), + total_tokens: None, + ..Default::default() + }, + TokenUsage::default(), + ); assert!(build_prompt_usage(&session).is_none()); } #[test] fn test_build_usage_update_clamps_negative_used_to_zero() { - let mut session = make_session_with_usage(Some(-7), Some(0), Some(0), None, None, None); + let mut session = make_session_with_usage( + TokenUsage::new(Some(0), Some(0), Some(-7)), + TokenUsage::default(), + ); session.model_config = Some( goose_providers::model::ModelConfig::new("test-model") .unwrap() @@ -3775,7 +3760,10 @@ print(\"hello, world\") #[test] fn test_build_usage_update_requires_model_config() { - let session = make_session_with_usage(Some(120), Some(80), Some(40), None, None, None); + let session = make_session_with_usage( + TokenUsage::new(Some(80), Some(40), Some(120)), + TokenUsage::default(), + ); assert!(build_usage_updates(&session).is_none()); } diff --git a/crates/goose/src/agents/execute_commands.rs b/crates/goose/src/agents/execute_commands.rs index eebb0e823..13f2420d8 100644 --- a/crates/goose/src/agents/execute_commands.rs +++ b/crates/goose/src/agents/execute_commands.rs @@ -179,9 +179,11 @@ impl Agent { manager .update(session_id) - .total_tokens(Some(0)) - .input_tokens(Some(0)) - .output_tokens(Some(0)) + .usage(goose_providers::conversation::token_usage::Usage::new( + Some(0), + Some(0), + Some(0), + )) .apply() .await?; diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index f6fda028a..5bc8a79b1 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -76,7 +76,9 @@ pub async fn inject_moim( .as_ref() .map(|session| session.working_dir.clone()) .unwrap_or_else(|| PathBuf::from(".")); - let total_tokens = session.as_ref().and_then(|session| session.total_tokens); + let total_tokens = session + .as_ref() + .and_then(|session| session.usage.total_tokens); let compaction_threshold = crate::config::Config::global() .get_param::("GOOSE_AUTO_COMPACT_THRESHOLD") .unwrap_or(crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD); diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 0a3873179..7d34581c6 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -21,7 +21,7 @@ use crate::providers::toolshim::{ augment_message_with_selected_tool_interpreter, convert_tool_messages_to_text, modify_system_prompt_for_tool_json, sanitize_residual_markers, }; -use goose_providers::conversation::token_usage::ProviderUsage; +use goose_providers::conversation::token_usage::{ProviderUsage, Usage}; use rmcp::model::Tool; use tracing::warn; @@ -523,19 +523,7 @@ impl Agent { let manager = self.config.session_manager.clone(); let session = manager.get_session(session_id, false).await?; - let accumulate = |a: Option, b: Option| -> Option { - match (a, b) { - (Some(x), Some(y)) => Some(x + y), - _ => a.or(b), - } - }; - - let accumulated_total = - accumulate(session.accumulated_total_tokens, usage.usage.total_tokens); - let accumulated_input = - accumulate(session.accumulated_input_tokens, usage.usage.input_tokens); - let accumulated_output = - accumulate(session.accumulated_output_tokens, usage.usage.output_tokens); + let accumulated_usage = session.accumulated_usage + usage.usage; let accumulated_cost = session .provider_name @@ -543,27 +531,19 @@ impl Agent { .and_then(|pn| self.accumulate_cost(session.accumulated_cost, usage, pn)) .or(session.accumulated_cost); - let (current_total, current_input, current_output) = if is_compaction_usage { + let current_usage = if is_compaction_usage { // After compaction: summary output becomes new input context let new_input = usage.usage.output_tokens; - (new_input, new_input, None) + Usage::new(new_input, None, new_input) } else { - ( - usage.usage.total_tokens, - usage.usage.input_tokens, - usage.usage.output_tokens, - ) + usage.usage }; manager .update(session_id) .schedule_id(schedule_id) - .total_tokens(current_total) - .input_tokens(current_input) - .output_tokens(current_output) - .accumulated_total_tokens(accumulated_total) - .accumulated_input_tokens(accumulated_input) - .accumulated_output_tokens(accumulated_output) + .usage(current_usage) + .accumulated_usage(accumulated_usage) .accumulated_cost(accumulated_cost) .apply() .await?; @@ -580,13 +560,7 @@ impl Agent { let canonical = crate::providers::canonical::maybe_get_canonical_model(provider_name, &usage.model)?; - let input_price = canonical.cost.input?; - let output_price = canonical.cost.output?; - - let input_tokens = usage.usage.input_tokens.unwrap_or(0) as f64; - let output_tokens = usage.usage.output_tokens.unwrap_or(0) as f64; - - let chunk_cost = (input_tokens * input_price + output_tokens * output_price) / 1_000_000.0; + let chunk_cost = canonical.cost.estimate_cost(&usage.usage)?; Some(existing.unwrap_or(0.0) + chunk_cost) } diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 1cc4173db..9db9b40e0 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -203,7 +203,7 @@ pub async fn check_if_compaction_needed( let context_limit = provider.get_model_config().context_limit(); - let (current_tokens, _token_source) = match session.total_tokens { + let (current_tokens, _token_source) = match session.usage.total_tokens { Some(tokens) => (tokens as usize, "session metadata"), None => { let token_counter = create_token_counter() diff --git a/crates/goose/src/providers/claude_code.rs b/crates/goose/src/providers/claude_code.rs index 1e51f8b7d..63bc7388e 100644 --- a/crates/goose/src/providers/claude_code.rs +++ b/crates/goose/src/providers/claude_code.rs @@ -827,6 +827,10 @@ impl Provider for ClaudeCodeProvider { let new = extract_usage_tokens(usage_info); if let Some(i) = new.input_tokens { accumulated_usage.input_tokens = Some(i); + accumulated_usage.cache_read_input_tokens = + new.cache_read_input_tokens; + accumulated_usage.cache_write_input_tokens = + new.cache_write_input_tokens; } } } @@ -846,11 +850,37 @@ impl Provider for ClaudeCodeProvider { process.needs_drain = false; if let Some(usage_info) = parsed.get("usage") { let new = extract_usage_tokens(usage_info); - accumulated_usage = Usage::new( - new.input_tokens.or(accumulated_usage.input_tokens), - new.output_tokens.or(accumulated_usage.output_tokens), - None, - ); + let reports_own_cache = new.cache_read_input_tokens.is_some() + || new.cache_write_input_tokens.is_some(); + let cache_read = new + .cache_read_input_tokens + .or(accumulated_usage.cache_read_input_tokens); + let cache_write = new + .cache_write_input_tokens + .or(accumulated_usage.cache_write_input_tokens); + // A result with raw input but no cache breakdown + // inherits the streamed breakdown; fold it back in + // so input stays inclusive of cache tokens. + let output_tokens = + new.output_tokens.or(accumulated_usage.output_tokens); + accumulated_usage = if new.input_tokens.is_some() + && !reports_own_cache + { + Usage::from_cache_exclusive_input( + new.input_tokens, + output_tokens, + None, + cache_read, + cache_write, + ) + } else { + Usage::new( + new.input_tokens.or(accumulated_usage.input_tokens), + output_tokens, + None, + ) + .with_cache_tokens(cache_read, cache_write) + }; } break; } @@ -965,6 +995,31 @@ mod tests { assert_eq!(usage.output_tokens, expected_output); } + #[tokio::test] + async fn test_result_without_cache_fields_keeps_streamed_cache_coherent() { + use futures::StreamExt; + + let (_provider, mut stream, _stdin_reader) = stream_with_canned_stdout(&[ + r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#, + r#"{"type":"stream_event","event":{"type":"message_start","message":{"usage":{"input_tokens":7,"cache_read_input_tokens":5000,"cache_creation_input_tokens":1000,"output_tokens":0}}}}"#, + r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":50}}"#, + ]) + .await; + + let mut final_usage = None; + while let Some(item) = stream.next().await { + if let Ok((_, Some(usage))) = item { + final_usage = Some(usage); + } + } + + let usage = final_usage.expect("stream should yield usage").usage; + assert_eq!(usage.cache_read_input_tokens, Some(5000)); + assert_eq!(usage.cache_write_input_tokens, Some(1000)); + assert_eq!(usage.input_tokens, Some(6010)); // 10 + 5000 + 1000 + assert_eq!(usage.output_tokens, Some(50)); + } + #[test_case( r#"{"type":"error","error":"context window exceeded"}"#, true diff --git a/crates/goose/src/providers/cli_common.rs b/crates/goose/src/providers/cli_common.rs index fdfc89ee2..b79905f1f 100644 --- a/crates/goose/src/providers/cli_common.rs +++ b/crates/goose/src/providers/cli_common.rs @@ -13,10 +13,12 @@ pub(crate) fn extract_usage_tokens(usage_info: &Value) -> Usage { .and_then(|v| v.as_i64()) .and_then(|v| i32::try_from(v).ok()) }; - Usage::new( + Usage::from_cache_exclusive_input( get("input_tokens"), get("output_tokens"), get("total_tokens"), + get("cache_read_input_tokens"), + get("cache_creation_input_tokens"), ) } diff --git a/crates/goose/src/providers/codex.rs b/crates/goose/src/providers/codex.rs index c7f054dd4..47456296b 100644 --- a/crates/goose/src/providers/codex.rs +++ b/crates/goose/src/providers/codex.rs @@ -291,7 +291,7 @@ impl CodexProvider { } } - /// Extract usage information from a JSON object + /// Codex `input_tokens` already includes `cached_input_tokens`. fn extract_usage(usage_info: &serde_json::Value, usage: &mut Usage) { if usage.input_tokens.is_none() { usage.input_tokens = usage_info @@ -305,6 +305,12 @@ impl CodexProvider { .and_then(|v| v.as_i64()) .map(|v| v as i32); } + if usage.cache_read_input_tokens.is_none() { + usage.cache_read_input_tokens = usage_info + .get("cached_input_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + } } /// Extract error message from an error event @@ -981,6 +987,7 @@ mod tests { assert_eq!(usage.input_tokens, Some(100)); assert_eq!(usage.output_tokens, Some(50)); assert_eq!(usage.total_tokens, Some(150)); + assert_eq!(usage.cache_read_input_tokens, Some(30)); } #[test] @@ -1081,6 +1088,7 @@ mod tests { assert_eq!(usage.input_tokens, Some(5000)); assert_eq!(usage.output_tokens, Some(100)); assert_eq!(usage.total_tokens, Some(5100)); + assert_eq!(usage.cache_read_input_tokens, Some(3000)); } #[test_case( diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index 9f7534e30..fa5315b36 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -444,89 +444,64 @@ pub fn response_to_message(response: &Value) -> Result { Ok(message) } -/// Extract usage information from Anthropic's API response +fn usage_from_anthropic_fields(usage: &Value) -> Usage { + let field = |key: &str| { + usage + .get(key) + .and_then(|v| v.as_u64()) + .map(|v| v.min(i32::MAX as u64) as i32) + }; + + Usage::from_cache_exclusive_input( + Some(field("input_tokens").unwrap_or(0)), + Some(field("output_tokens").unwrap_or(0)), + None, + field("cache_read_input_tokens"), + field("cache_creation_input_tokens"), + ) +} + +/// Merge a `message_delta` usage into the usage captured at `message_start`. +/// Delta usage is cumulative (input grows during server tool use), so fields +/// present in the raw delta payload win over the start values. +fn merge_delta_usage(existing: &Usage, delta: &Usage, delta_data: &Value) -> Usage { + let reports = |key: &str| delta_data.get(key).is_some(); + + let output = if reports("output_tokens") { + delta.output_tokens + } else { + existing.output_tokens + }; + + if !reports("input_tokens") { + Usage::new(existing.input_tokens, output, None).with_cache_tokens( + existing.cache_read_input_tokens, + existing.cache_write_input_tokens, + ) + } else if reports("cache_read_input_tokens") || reports("cache_creation_input_tokens") { + Usage::new(delta.input_tokens, output, None).with_cache_tokens( + delta.cache_read_input_tokens, + delta.cache_write_input_tokens, + ) + } else { + Usage::from_cache_exclusive_input( + delta.input_tokens, + output, + None, + existing.cache_read_input_tokens, + existing.cache_write_input_tokens, + ) + } +} + pub fn get_usage(data: &Value) -> Result { - // Extract usage data if available if let Some(usage) = data.get("usage") { - // Get all token fields for analysis - let input_tokens = usage - .get("input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - let cache_creation_tokens = usage - .get("cache_creation_input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - let cache_read_tokens = usage - .get("cache_read_input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - let output_tokens = usage - .get("output_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - // IMPORTANT: For display purposes, we want to show the ACTUAL total tokens consumed - // The cache pricing should only affect cost calculation, not token count display - let total_input_tokens = input_tokens + cache_creation_tokens + cache_read_tokens; - - // Convert to i32 with bounds checking - let total_input_i32 = total_input_tokens.min(i32::MAX as u64) as i32; - let output_tokens_i32 = output_tokens.min(i32::MAX as u64) as i32; - let total_tokens_i32 = - (total_input_i32 as i64 + output_tokens_i32 as i64).min(i32::MAX as i64) as i32; - - Ok(Usage::new( - Some(total_input_i32), - Some(output_tokens_i32), - Some(total_tokens_i32), - )) + Ok(usage_from_anthropic_fields(usage)) } else if data.as_object().is_some() { // Check if the data itself is the usage object (for message_delta events that might have usage at top level) - let input_tokens = data - .get("input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - let cache_creation_tokens = data - .get("cache_creation_input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - let cache_read_tokens = data - .get("cache_read_input_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - let output_tokens = data - .get("output_tokens") - .and_then(|v| v.as_u64()) - .unwrap_or(0); - - // If we found any token data, process it - if input_tokens > 0 - || cache_creation_tokens > 0 - || cache_read_tokens > 0 - || output_tokens > 0 - { - let total_input_tokens = input_tokens + cache_creation_tokens + cache_read_tokens; - - let total_input_i32 = total_input_tokens.min(i32::MAX as u64) as i32; - let output_tokens_i32 = output_tokens.min(i32::MAX as u64) as i32; - let total_tokens_i32 = - (total_input_i32 as i64 + output_tokens_i32 as i64).min(i32::MAX as i64) as i32; - - tracing::debug!("🔍 Anthropic ACTUAL token counts from direct object: input={}, output={}, total={}", - total_input_i32, output_tokens_i32, total_tokens_i32); - - Ok(Usage::new( - Some(total_input_i32), - Some(output_tokens_i32), - Some(total_tokens_i32), - )) + let usage = usage_from_anthropic_fields(data); + if usage.total_tokens.unwrap_or(0) > 0 { + Ok(usage) } else { tracing::debug!("🔍 Anthropic no token data found in object"); Ok(Usage::new(None, None, None)) @@ -947,16 +922,7 @@ where let delta_usage = get_usage(usage_data).unwrap_or_default(); if let Some(existing_usage) = &final_usage { - let merged_input = existing_usage.usage.input_tokens.or(delta_usage.input_tokens); - let merged_output = delta_usage.output_tokens.or(existing_usage.usage.output_tokens); - let merged_total = match (merged_input, merged_output) { - (Some(input), Some(output)) => Some(input + output), - (Some(input), None) => Some(input), - (None, Some(output)) => Some(output), - (None, None) => None, - }; - - let merged_usage = Usage::new(merged_input, merged_output, merged_total); + let merged_usage = merge_delta_usage(&existing_usage.usage, &delta_usage, usage_data); final_usage = Some(ProviderUsage::new(existing_usage.model.clone(), merged_usage)); } else { let model = event.data.get("model") @@ -1061,6 +1027,8 @@ mod tests { assert_eq!(usage.input_tokens, Some(24)); // 12 + 12 = 24 actual tokens assert_eq!(usage.output_tokens, Some(15)); assert_eq!(usage.total_tokens, Some(39)); // 24 + 15 + assert_eq!(usage.cache_read_input_tokens, Some(0)); + assert_eq!(usage.cache_write_input_tokens, Some(12)); Ok(()) } @@ -1104,6 +1072,8 @@ mod tests { assert_eq!(usage.input_tokens, Some(30)); // 15 + 15 = 30 actual tokens assert_eq!(usage.output_tokens, Some(20)); assert_eq!(usage.total_tokens, Some(50)); // 30 + 20 + assert_eq!(usage.cache_read_input_tokens, Some(0)); + assert_eq!(usage.cache_write_input_tokens, Some(15)); Ok(()) } @@ -1281,6 +1251,8 @@ mod tests { assert_eq!(usage.input_tokens, Some(15007)); assert_eq!(usage.output_tokens, Some(50)); assert_eq!(usage.total_tokens, Some(15057)); // 15007 + 50 + assert_eq!(usage.cache_read_input_tokens, Some(5000)); + assert_eq!(usage.cache_write_input_tokens, Some(10000)); Ok(()) } @@ -1401,7 +1373,7 @@ mod tests { )?; assert_eq!(payload["thinking"]["type"], "enabled"); - assert_eq!(payload["thinking"]["budget_tokens"], 16000); + assert!(payload["thinking"]["budget_tokens"].as_i64().unwrap() >= 1024); assert_eq!(payload["thinking"]["clear_thinking"], false); assert_eq!(payload["max_tokens"], 64000); assert_eq!(payload["messages"][0]["content"][0]["type"], "thinking"); @@ -1915,6 +1887,80 @@ mod tests { response_to_streaming_message(stream).collect().await } + #[tokio::test] + async fn test_streaming_preserves_cache_tokens_through_delta_merge() { + let events = concat!( + r#"data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-opus-4-6","usage":{"input_tokens":7,"cache_creation_input_tokens":10000,"cache_read_input_tokens":5000,"output_tokens":0}}}"#, + "\n", + r#"data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}"#, + "\n", + r#"data: {"type":"content_block_stop","index":0}"#, + "\n", + r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":25}}"#, + "\n", + r#"data: {"type":"message_stop"}"#, + ); + + let usage = collect_stream_results(events) + .await + .into_iter() + .filter_map(|r| r.ok().and_then(|(_, usage)| usage)) + .next_back() + .expect("stream should yield usage"); + + assert_eq!(usage.usage.input_tokens, Some(15007)); + assert_eq!(usage.usage.output_tokens, Some(25)); + assert_eq!(usage.usage.cache_read_input_tokens, Some(5000)); + assert_eq!(usage.usage.cache_write_input_tokens, Some(10000)); + } + + #[tokio::test] + async fn test_streaming_delta_usage_is_cumulative_and_wins() { + // Server tool use grows input during the turn: the final + // message_delta usage is authoritative, not message_start. + let events = concat!( + r#"data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","content":[],"model":"claude-opus-4-6","usage":{"input_tokens":2679,"cache_creation_input_tokens":100,"cache_read_input_tokens":200,"output_tokens":3}}}"#, + "\n", + r#"data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}"#, + "\n", + r#"data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}"#, + "\n", + r#"data: {"type":"content_block_stop","index":0}"#, + "\n", + r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":10682,"cache_creation_input_tokens":100,"cache_read_input_tokens":200,"output_tokens":510}}"#, + "\n", + r#"data: {"type":"message_stop"}"#, + ); + + let usage = collect_stream_results(events) + .await + .into_iter() + .filter_map(|r| r.ok().and_then(|(_, usage)| usage)) + .next_back() + .expect("stream should yield usage"); + + assert_eq!(usage.usage.input_tokens, Some(10982)); // 10682 + 100 + 200 + assert_eq!(usage.usage.output_tokens, Some(510)); + assert_eq!(usage.usage.cache_read_input_tokens, Some(200)); + assert_eq!(usage.usage.cache_write_input_tokens, Some(100)); + } + + #[test] + fn test_merge_delta_usage_raw_input_inherits_start_cache() { + let start = + Usage::new(Some(15007), Some(3), None).with_cache_tokens(Some(5000), Some(10000)); + let delta_data = json!({"input_tokens": 8, "output_tokens": 510}); + let delta = get_usage(&delta_data).unwrap(); + + let merged = merge_delta_usage(&start, &delta, &delta_data); + assert_eq!(merged.input_tokens, Some(15008)); // 8 + 5000 + 10000 + assert_eq!(merged.output_tokens, Some(510)); + assert_eq!(merged.cache_read_input_tokens, Some(5000)); + assert_eq!(merged.cache_write_input_tokens, Some(10000)); + } + fn expect_refusal( results: Vec, Option)>>, ) -> (String, Option) { diff --git a/crates/goose/src/providers/formats/bedrock.rs b/crates/goose/src/providers/formats/bedrock.rs index 85f354b20..57f0eae34 100644 --- a/crates/goose/src/providers/formats/bedrock.rs +++ b/crates/goose/src/providers/formats/bedrock.rs @@ -491,10 +491,12 @@ pub fn from_bedrock_role(role: &bedrock::ConversationRole) -> Result { } pub fn from_bedrock_usage(usage: &bedrock::TokenUsage) -> Usage { - Usage::new( + Usage::from_cache_exclusive_input( Some(usage.input_tokens), Some(usage.output_tokens), Some(usage.total_tokens), + usage.cache_read_input_tokens, + usage.cache_write_input_tokens, ) } @@ -761,6 +763,25 @@ mod tests { Ok(()) } + #[test] + fn test_from_bedrock_usage_folds_cache_tokens_into_input() { + let usage = bedrock::TokenUsage::builder() + .input_tokens(7) + .output_tokens(50) + .total_tokens(57) + .cache_read_input_tokens(5000) + .cache_write_input_tokens(1000) + .build() + .unwrap(); + + let converted = from_bedrock_usage(&usage); + assert_eq!(converted.input_tokens, Some(6007)); + assert_eq!(converted.output_tokens, Some(50)); + assert_eq!(converted.total_tokens, Some(6057)); + assert_eq!(converted.cache_read_input_tokens, Some(5000)); + assert_eq!(converted.cache_write_input_tokens, Some(1000)); + } + #[test] fn test_from_bedrock_content_block_cache_point() { // Create a cache point block with the required type field diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs index f869959d9..cf4d120ca 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose/src/providers/formats/google.rs @@ -347,7 +347,13 @@ pub fn get_usage(data: &Value) -> Result { .get("totalTokenCount") .and_then(|v| v.as_u64()) .map(|v| v as i32); - Ok(Usage::new(input_tokens, output_tokens, total_tokens)) + // promptTokenCount already includes cachedContentTokenCount + let cached_tokens = usage_meta_data + .get("cachedContentTokenCount") + .and_then(|v| v.as_u64()) + .map(|v| v as i32); + Ok(Usage::new(input_tokens, output_tokens, total_tokens) + .with_cache_tokens(cached_tokens, None)) } else { tracing::debug!( "Failed to get usage data: {}", @@ -678,6 +684,26 @@ mod tests { assert_eq!(usage.input_tokens, Some(1)); assert_eq!(usage.output_tokens, Some(2)); assert_eq!(usage.total_tokens, Some(3)); + assert_eq!(usage.cache_read_input_tokens, None); + assert_eq!(usage.cache_write_input_tokens, None); + } + + #[test] + fn test_get_usage_with_cached_content() { + let data = json!({ + "usageMetadata": { + "promptTokenCount": 100, + "candidatesTokenCount": 20, + "totalTokenCount": 120, + "cachedContentTokenCount": 80 + } + }); + let usage = get_usage(&data).unwrap(); + assert_eq!(usage.input_tokens, Some(100)); + assert_eq!(usage.output_tokens, Some(20)); + assert_eq!(usage.total_tokens, Some(120)); + assert_eq!(usage.cache_read_input_tokens, Some(80)); + assert_eq!(usage.cache_write_input_tokens, None); } #[test] diff --git a/crates/goose/src/providers/formats/openai_responses.rs b/crates/goose/src/providers/formats/openai_responses.rs index 9dd46028a..ee0c2d8f1 100644 --- a/crates/goose/src/providers/formats/openai_responses.rs +++ b/crates/goose/src/providers/formats/openai_responses.rs @@ -112,6 +112,30 @@ pub struct ResponseUsage { pub input_tokens: i32, pub output_tokens: i32, pub total_tokens: i32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input_tokens_details: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct InputTokensDetails { + #[serde(default)] + pub cached_tokens: Option, +} + +impl ResponseUsage { + fn to_usage(&self) -> Usage { + // input_tokens already includes cached tokens + let cached_tokens = self + .input_tokens_details + .as_ref() + .and_then(|d| d.cached_tokens); + Usage::new( + Some(self.input_tokens), + Some(self.output_tokens), + Some(self.total_tokens), + ) + .with_cache_tokens(cached_tokens, None) + } } #[derive(Debug, Serialize, Deserialize)] @@ -690,13 +714,10 @@ pub fn responses_api_to_message(response: &ResponsesApiResponse) -> anyhow::Resu } pub fn get_responses_usage(response: &ResponsesApiResponse) -> Usage { - response.usage.as_ref().map_or_else(Usage::default, |u| { - Usage::new( - Some(u.input_tokens), - Some(u.output_tokens), - Some(u.total_tokens), - ) - }) + response + .usage + .as_ref() + .map_or_else(Usage::default, ResponseUsage::to_usage) } fn process_streaming_output_items( @@ -858,11 +879,7 @@ where let model = model_name.as_ref().unwrap_or(&response.model); let usage = response.usage.as_ref().map_or_else( Usage::default, - |u| Usage::new( - Some(u.input_tokens), - Some(u.output_tokens), - Some(u.total_tokens), - ), + ResponseUsage::to_usage, ); final_usage = Some(ProviderUsage::new(model.clone(), usage)); @@ -957,7 +974,7 @@ mod tests { r#"data: {"type":"keepalive"}"#.to_string(), r#"data: {"type":"response.output_text.delta","sequence_number":2,"item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"}"#.to_string(), r#"data: {"type":"response.output_text.delta","sequence_number":3,"item_id":"msg_1","output_index":0,"content_index":0,"delta":" world"}"#.to_string(), - r#"data: {"type":"response.completed","sequence_number":4,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"completed","model":"gpt-5.2-pro","output":[],"usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14}}}"#.to_string(), + r#"data: {"type":"response.completed","sequence_number":4,"response":{"id":"resp_1","object":"response","created_at":1737368310,"status":"completed","model":"gpt-5.2-pro","output":[],"usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14,"input_tokens_details":{"cached_tokens":6}}}}"#.to_string(), "data: [DONE]".to_string(), ]; @@ -988,6 +1005,8 @@ mod tests { assert_eq!(usage.usage.input_tokens, Some(10)); assert_eq!(usage.usage.output_tokens, Some(4)); assert_eq!(usage.usage.total_tokens, Some(14)); + assert_eq!(usage.usage.cache_read_input_tokens, Some(6)); + assert_eq!(usage.usage.cache_write_input_tokens, None); Ok(()) } diff --git a/crates/goose/src/scheduler.rs b/crates/goose/src/scheduler.rs index 4aab4dbf0..6c42cbd0d 100644 --- a/crates/goose/src/scheduler.rs +++ b/crates/goose/src/scheduler.rs @@ -972,7 +972,7 @@ async fn execute_job( .session_manager .get_session(&session.id, false) .await - .map(|s| (s.total_tokens.unwrap_or(0), s.message_count)) + .map(|s| (s.usage.total_tokens.unwrap_or(0), s.message_count)) .unwrap_or((0, 0)); tracing::info!( diff --git a/crates/goose/src/session/import_formats/claude_code.rs b/crates/goose/src/session/import_formats/claude_code.rs index 9c26b3e1d..e7e5f942a 100644 --- a/crates/goose/src/session/import_formats/claude_code.rs +++ b/crates/goose/src/session/import_formats/claude_code.rs @@ -9,10 +9,11 @@ use anyhow::{anyhow, Result}; use chrono::{DateTime, Utc}; use rmcp::model::{CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData}; -use serde_json::{json, Map, Value}; +use serde_json::Value; use crate::conversation::message::Message; use crate::conversation::Conversation; +use goose_providers::conversation::token_usage::Usage; pub fn convert(content: &str) -> Result { let lines: Vec = content @@ -50,6 +51,8 @@ pub fn convert(content: &str) -> Result { let mut messages: Vec = Vec::new(); let mut total_input: i64 = 0; let mut total_output: i64 = 0; + let mut total_cache_read: i64 = 0; + let mut total_cache_write: i64 = 0; let mut first_ts: Option> = None; let mut last_ts: Option> = None; let mut first_user_text: Option = None; @@ -83,18 +86,21 @@ pub fn convert(content: &str) -> Result { .and_then(|m| m.get("usage")) .and_then(|u| u.as_object()) { + let cache_write = usage + .get("cache_creation_input_tokens") + .and_then(|v| v.as_i64()) + .unwrap_or(0); + let cache_read = usage + .get("cache_read_input_tokens") + .and_then(|v| v.as_i64()) + .unwrap_or(0); total_input += usage .get("input_tokens") .and_then(|v| v.as_i64()) .unwrap_or(0); - total_input += usage - .get("cache_creation_input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - total_input += usage - .get("cache_read_input_tokens") - .and_then(|v| v.as_i64()) - .unwrap_or(0); + total_input += cache_write + cache_read; + total_cache_write += cache_write; + total_cache_read += cache_read; total_output += usage .get("output_tokens") .and_then(|v| v.as_i64()) @@ -124,16 +130,20 @@ pub fn convert(content: &str) -> Result { let conversation = Conversation::new_unvalidated(messages); - let session_json = build_session_json( - &session_id, - &working_dir, - &name, + let session_json = super::build_session_json(super::ImportedSession { + session_id: &session_id, + working_dir: &working_dir, + name: &name, created_at, updated_at, - Some(total_input as i32), - Some(total_output as i32), + usage: Usage::new(Some(total_input as i32), Some(total_output as i32), None) + .with_cache_tokens( + (total_cache_read > 0).then_some(total_cache_read as i32), + (total_cache_write > 0).then_some(total_cache_write as i32), + ), + accumulated_cost: None, conversation, - ); + }); serde_json::to_string_pretty(&session_json).map_err(Into::into) } @@ -299,53 +309,6 @@ fn extract_first_text(msg: &Message) -> Option { None } -#[allow(clippy::too_many_arguments)] -fn build_session_json( - session_id: &str, - working_dir: &str, - name: &str, - created_at: DateTime, - updated_at: DateTime, - input_tokens: Option, - output_tokens: Option, - conversation: Conversation, -) -> Value { - let total = match (input_tokens, output_tokens) { - (Some(a), Some(b)) => Some(a + b), - _ => None, - }; - let mut obj = Map::new(); - obj.insert("id".into(), json!(session_id)); - obj.insert("working_dir".into(), json!(working_dir)); - obj.insert("name".into(), json!(name)); - obj.insert("user_set_name".into(), json!(false)); - obj.insert("session_type".into(), json!("user")); - obj.insert("created_at".into(), json!(created_at.to_rfc3339())); - obj.insert("updated_at".into(), json!(updated_at.to_rfc3339())); - obj.insert("extension_data".into(), json!({})); - obj.insert("total_tokens".into(), json!(total)); - obj.insert("input_tokens".into(), json!(input_tokens)); - obj.insert("output_tokens".into(), json!(output_tokens)); - obj.insert("accumulated_total_tokens".into(), json!(total)); - obj.insert("accumulated_input_tokens".into(), json!(input_tokens)); - obj.insert("accumulated_output_tokens".into(), json!(output_tokens)); - obj.insert("accumulated_cost".into(), json!(null)); - obj.insert("schedule_id".into(), json!(null)); - obj.insert("recipe".into(), json!(null)); - obj.insert("user_recipe_values".into(), json!(null)); - obj.insert( - "conversation".into(), - serde_json::to_value(&conversation).unwrap(), - ); - obj.insert("message_count".into(), json!(conversation.messages().len())); - obj.insert("provider_name".into(), json!(null)); - obj.insert("model_config".into(), json!(null)); - obj.insert("goose_mode".into(), json!("auto")); - obj.insert("archived_at".into(), json!(null)); - obj.insert("project_id".into(), json!(null)); - Value::Object(obj) -} - #[cfg(test)] mod tests { use super::*; @@ -370,6 +333,20 @@ mod tests { assert!(content.iter().any(|c| c["type"] == "toolResponse")); } + #[test] + fn emits_cache_token_breakdown() { + let jsonl = r#"{"type":"user","sessionId":"s","uuid":"u1","timestamp":"2026-01-01T00:00:01Z","cwd":"/tmp","message":{"role":"user","content":"hi"}} +{"type":"assistant","sessionId":"s","uuid":"u2","timestamp":"2026-01-01T00:00:02Z","message":{"role":"assistant","content":[{"type":"text","text":"hello"}],"usage":{"input_tokens":7,"cache_creation_input_tokens":1000,"cache_read_input_tokens":5000,"output_tokens":50}}}"#; + let json = convert(jsonl).unwrap(); + let v: Value = serde_json::from_str(&json).unwrap(); + assert_eq!(v["usage"]["input_tokens"], 6007); // 7 + 1000 + 5000 + assert_eq!(v["usage"]["output_tokens"], 50); + assert_eq!(v["usage"]["cache_read_input_tokens"], 5000); + assert_eq!(v["usage"]["cache_write_input_tokens"], 1000); + assert_eq!(v["accumulated_usage"]["cache_read_input_tokens"], 5000); + assert_eq!(v["accumulated_usage"]["cache_write_input_tokens"], 1000); + } + #[test] fn skips_unknown_lines() { let jsonl = r#"{"type":"attachment","sessionId":"s","uuid":"u0","timestamp":"2026-01-01T00:00:00Z"} diff --git a/crates/goose/src/session/import_formats/codex.rs b/crates/goose/src/session/import_formats/codex.rs index 7e2ee6030..2c8beb1f0 100644 --- a/crates/goose/src/session/import_formats/codex.rs +++ b/crates/goose/src/session/import_formats/codex.rs @@ -26,6 +26,7 @@ use serde_json::{json, Map, Value}; use crate::conversation::message::Message; use crate::conversation::Conversation; use crate::providers::formats::openai_responses::{ResponseOutputItem, ResponsesApiResponse}; +use goose_providers::conversation::token_usage::Usage; pub fn convert(content: &str) -> Result { let lines: Vec = content @@ -61,6 +62,7 @@ pub fn convert(content: &str) -> Result { let mut first_user_text: Option = None; let mut total_input: i64 = 0; let mut total_output: i64 = 0; + let mut total_cache_read: i64 = 0; for (line_idx, line) in lines.iter().enumerate() { let line_type = line.get("type").and_then(|v| v.as_str()).unwrap_or(""); @@ -80,10 +82,15 @@ pub fn convert(content: &str) -> Result { .and_then(|p| p.get("usage")) .and_then(|u| u.as_object()) { + // Codex input_tokens already includes cached_input_tokens total_input += usage .get("input_tokens") .and_then(|v| v.as_i64()) .unwrap_or(0); + total_cache_read += usage + .get("cached_input_tokens") + .and_then(|v| v.as_i64()) + .unwrap_or(0); total_output += usage .get("output_tokens") .and_then(|v| v.as_i64()) @@ -239,24 +246,24 @@ pub fn convert(content: &str) -> Result { let updated_at = last_ts.unwrap_or(created_at); let conversation = Conversation::new_unvalidated(messages); - let session_json = build_session_json( - &session_id, - &working_dir, - &name, + let session_json = super::build_session_json(super::ImportedSession { + session_id: &session_id, + working_dir: &working_dir, + name: &name, created_at, updated_at, - if total_input > 0 { - Some(total_input as i32) - } else { - None - }, - if total_output > 0 { - Some(total_output as i32) - } else { - None - }, + usage: Usage::new( + (total_input > 0).then_some(total_input as i32), + (total_output > 0).then_some(total_output as i32), + None, + ) + .with_cache_tokens( + (total_cache_read > 0).then_some(total_cache_read as i32), + None, + ), + accumulated_cost: None, conversation, - ); + }); serde_json::to_string_pretty(&session_json).map_err(Into::into) } @@ -289,53 +296,6 @@ fn is_context_blob(text: &str) -> bool { || t.starts_with("# AGENTS.md") } -#[allow(clippy::too_many_arguments)] -fn build_session_json( - session_id: &str, - working_dir: &str, - name: &str, - created_at: DateTime, - updated_at: DateTime, - input_tokens: Option, - output_tokens: Option, - conversation: Conversation, -) -> Value { - let total = match (input_tokens, output_tokens) { - (Some(a), Some(b)) => Some(a + b), - _ => None, - }; - let mut obj = Map::new(); - obj.insert("id".into(), json!(session_id)); - obj.insert("working_dir".into(), json!(working_dir)); - obj.insert("name".into(), json!(name)); - obj.insert("user_set_name".into(), json!(false)); - obj.insert("session_type".into(), json!("user")); - obj.insert("created_at".into(), json!(created_at.to_rfc3339())); - obj.insert("updated_at".into(), json!(updated_at.to_rfc3339())); - obj.insert("extension_data".into(), json!({})); - obj.insert("total_tokens".into(), json!(total)); - obj.insert("input_tokens".into(), json!(input_tokens)); - obj.insert("output_tokens".into(), json!(output_tokens)); - obj.insert("accumulated_total_tokens".into(), json!(total)); - obj.insert("accumulated_input_tokens".into(), json!(input_tokens)); - obj.insert("accumulated_output_tokens".into(), json!(output_tokens)); - obj.insert("accumulated_cost".into(), json!(null)); - obj.insert("schedule_id".into(), json!(null)); - obj.insert("recipe".into(), json!(null)); - obj.insert("user_recipe_values".into(), json!(null)); - obj.insert( - "conversation".into(), - serde_json::to_value(&conversation).unwrap(), - ); - obj.insert("message_count".into(), json!(conversation.messages().len())); - obj.insert("provider_name".into(), json!(null)); - obj.insert("model_config".into(), json!(null)); - obj.insert("goose_mode".into(), json!("auto")); - obj.insert("archived_at".into(), json!(null)); - obj.insert("project_id".into(), json!(null)); - Value::Object(obj) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose/src/session/import_formats/mod.rs b/crates/goose/src/session/import_formats/mod.rs index 8d408e19f..50dabed3e 100644 --- a/crates/goose/src/session/import_formats/mod.rs +++ b/crates/goose/src/session/import_formats/mod.rs @@ -13,11 +13,62 @@ //! `SessionManager::import_session` pipeline. use anyhow::Result; +use chrono::{DateTime, Utc}; +use goose_providers::conversation::token_usage::Usage; +use serde_json::{json, Map, Value}; + +use crate::conversation::Conversation; pub mod claude_code; pub mod codex; pub mod pi; +/// Session-level fields harvested from a foreign transcript, used to build +/// the goose-native session JSON handed to `SessionManager::import_session`. +pub(crate) struct ImportedSession<'a> { + pub session_id: &'a str, + pub working_dir: &'a str, + pub name: &'a str, + pub created_at: DateTime, + pub updated_at: DateTime, + pub usage: Usage, + pub accumulated_cost: Option, + pub conversation: Conversation, +} + +pub(crate) fn build_session_json(session: ImportedSession) -> Value { + let usage = serde_json::to_value(session.usage).unwrap(); + let mut obj = Map::new(); + obj.insert("id".into(), json!(session.session_id)); + obj.insert("working_dir".into(), json!(session.working_dir)); + obj.insert("name".into(), json!(session.name)); + obj.insert("user_set_name".into(), json!(false)); + obj.insert("session_type".into(), json!("user")); + obj.insert("created_at".into(), json!(session.created_at.to_rfc3339())); + obj.insert("updated_at".into(), json!(session.updated_at.to_rfc3339())); + obj.insert("extension_data".into(), json!({})); + obj.insert("usage".into(), usage.clone()); + obj.insert("accumulated_usage".into(), usage); + obj.insert("accumulated_cost".into(), json!(session.accumulated_cost)); + obj.insert("schedule_id".into(), json!(null)); + obj.insert("recipe".into(), json!(null)); + obj.insert("user_recipe_values".into(), json!(null)); + obj.insert( + "conversation".into(), + serde_json::to_value(&session.conversation).unwrap(), + ); + obj.insert( + "message_count".into(), + json!(session.conversation.messages().len()), + ); + obj.insert("provider_name".into(), json!(null)); + obj.insert("model_config".into(), json!(null)); + obj.insert("goose_mode".into(), json!("auto")); + obj.insert("archived_at".into(), json!(null)); + obj.insert("project_id".into(), json!(null)); + Value::Object(obj) +} + /// Detected import source format. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ImportFormat { @@ -96,17 +147,69 @@ pub fn detect_format(content: &str) -> ImportFormat { } /// Convert any supported foreign format to a goose-native session JSON string. -/// -/// For [`ImportFormat::Goose`] the input is returned unchanged. pub fn convert_to_goose_session_json(content: &str) -> Result { match detect_format(content) { - ImportFormat::Goose => Ok(content.to_string()), + ImportFormat::Goose => Ok(upgrade_legacy_token_fields(content)), ImportFormat::ClaudeCode => claude_code::convert(content), ImportFormat::Codex => codex::convert(content), ImportFormat::Pi => pi::convert(content), } } +/// Exports from goose versions before `usage`/`accumulated_usage` existed +/// store token counts as flat fields. +fn upgrade_legacy_token_fields(content: &str) -> String { + let Ok(Value::Object(mut obj)) = serde_json::from_str::(content) else { + return content.to_string(); + }; + nest_legacy_token_fields(&mut obj); + Value::Object(obj).to_string() +} + +/// Fold pre-`usage` flat token counts into nested `usage`/`accumulated_usage` +/// objects. Shared by the export-import path and the first-run JSONL migration. +pub(crate) fn nest_legacy_token_fields(obj: &mut Map) { + let nest = |obj: &Map, keys: [(&str, &str); 5]| -> Value { + Value::Object( + keys.iter() + .map(|(from, to)| { + ( + to.to_string(), + obj.get(*from).cloned().unwrap_or(Value::Null), + ) + }) + .collect(), + ) + }; + + if !obj.contains_key("usage") { + let usage = nest( + obj, + [ + ("input_tokens", "input_tokens"), + ("output_tokens", "output_tokens"), + ("total_tokens", "total_tokens"), + ("cache_read_tokens", "cache_read_input_tokens"), + ("cache_write_tokens", "cache_write_input_tokens"), + ], + ); + obj.insert("usage".into(), usage); + } + if !obj.contains_key("accumulated_usage") { + let accumulated = nest( + obj, + [ + ("accumulated_input_tokens", "input_tokens"), + ("accumulated_output_tokens", "output_tokens"), + ("accumulated_total_tokens", "total_tokens"), + ("accumulated_cache_read_tokens", "cache_read_input_tokens"), + ("accumulated_cache_write_tokens", "cache_write_input_tokens"), + ], + ); + obj.insert("accumulated_usage".into(), accumulated); + } +} + /// Squeeze a string down to a short session-name candidate: take the first /// non-empty line and cap it at ~80 chars. pub(crate) fn summarize_first_line(s: &str) -> String { diff --git a/crates/goose/src/session/import_formats/pi.rs b/crates/goose/src/session/import_formats/pi.rs index 7e96c2545..39c3f15c2 100644 --- a/crates/goose/src/session/import_formats/pi.rs +++ b/crates/goose/src/session/import_formats/pi.rs @@ -16,6 +16,7 @@ use serde_json::{json, Map, Value}; use crate::conversation::message::Message; use crate::conversation::Conversation; +use goose_providers::conversation::token_usage::Usage; pub fn convert(content: &str) -> Result { let mut lines = content.lines().filter(|l| !l.trim().is_empty()); @@ -48,6 +49,8 @@ pub fn convert(content: &str) -> Result { let mut messages: Vec = Vec::new(); let mut total_input: i64 = 0; let mut total_output: i64 = 0; + let mut total_cache_read: i64 = 0; + let mut total_cache_write: i64 = 0; let mut total_cost: f64 = 0.0; let mut first_ts: Option> = header_ts; let mut last_ts: Option> = header_ts; @@ -84,12 +87,15 @@ pub fn convert(content: &str) -> Result { .unwrap_or_else(|| Utc::now().timestamp()); if let Some(usage) = inner.get("usage").and_then(|u| u.as_object()) { - total_input += usage.get("input").and_then(|v| v.as_i64()).unwrap_or(0); - total_input += usage.get("cacheRead").and_then(|v| v.as_i64()).unwrap_or(0); - total_input += usage + let cache_read = usage.get("cacheRead").and_then(|v| v.as_i64()).unwrap_or(0); + let cache_write = usage .get("cacheWrite") .and_then(|v| v.as_i64()) .unwrap_or(0); + total_input += usage.get("input").and_then(|v| v.as_i64()).unwrap_or(0); + total_input += cache_read + cache_write; + total_cache_read += cache_read; + total_cache_write += cache_write; total_output += usage.get("output").and_then(|v| v.as_i64()).unwrap_or(0); if let Some(cost) = usage .get("cost") @@ -202,21 +208,20 @@ pub fn convert(content: &str) -> Result { let updated_at = last_ts.unwrap_or(created_at); let conversation = Conversation::new_unvalidated(messages); - let session_json = build_session_json( - &session_id, - &working_dir, - &name, + let session_json = super::build_session_json(super::ImportedSession { + session_id: &session_id, + working_dir: &working_dir, + name: &name, created_at, updated_at, - Some(total_input as i32), - Some(total_output as i32), - if total_cost > 0.0 { - Some(total_cost) - } else { - None - }, + usage: Usage::new(Some(total_input as i32), Some(total_output as i32), None) + .with_cache_tokens( + (total_cache_read > 0).then_some(total_cache_read as i32), + (total_cache_write > 0).then_some(total_cache_write as i32), + ), + accumulated_cost: (total_cost > 0.0).then_some(total_cost), conversation, - ); + }); serde_json::to_string_pretty(&session_json).map_err(Into::into) } @@ -333,54 +338,6 @@ fn extract_first_text(msg: &Message) -> Option { None } -#[allow(clippy::too_many_arguments)] -fn build_session_json( - session_id: &str, - working_dir: &str, - name: &str, - created_at: DateTime, - updated_at: DateTime, - input_tokens: Option, - output_tokens: Option, - cost: Option, - conversation: Conversation, -) -> Value { - let total = match (input_tokens, output_tokens) { - (Some(a), Some(b)) => Some(a + b), - _ => None, - }; - let mut obj = Map::new(); - obj.insert("id".into(), json!(session_id)); - obj.insert("working_dir".into(), json!(working_dir)); - obj.insert("name".into(), json!(name)); - obj.insert("user_set_name".into(), json!(false)); - obj.insert("session_type".into(), json!("user")); - obj.insert("created_at".into(), json!(created_at.to_rfc3339())); - obj.insert("updated_at".into(), json!(updated_at.to_rfc3339())); - obj.insert("extension_data".into(), json!({})); - obj.insert("total_tokens".into(), json!(total)); - obj.insert("input_tokens".into(), json!(input_tokens)); - obj.insert("output_tokens".into(), json!(output_tokens)); - obj.insert("accumulated_total_tokens".into(), json!(total)); - obj.insert("accumulated_input_tokens".into(), json!(input_tokens)); - obj.insert("accumulated_output_tokens".into(), json!(output_tokens)); - obj.insert("accumulated_cost".into(), json!(cost)); - obj.insert("schedule_id".into(), json!(null)); - obj.insert("recipe".into(), json!(null)); - obj.insert("user_recipe_values".into(), json!(null)); - obj.insert( - "conversation".into(), - serde_json::to_value(&conversation).unwrap(), - ); - obj.insert("message_count".into(), json!(conversation.messages().len())); - obj.insert("provider_name".into(), json!(null)); - obj.insert("model_config".into(), json!(null)); - obj.insert("goose_mode".into(), json!("auto")); - obj.insert("archived_at".into(), json!(null)); - obj.insert("project_id".into(), json!(null)); - Value::Object(obj) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose/src/session/legacy.rs b/crates/goose/src/session/legacy.rs index a879412b9..0d90d4fa6 100644 --- a/crates/goose/src/session/legacy.rs +++ b/crates/goose/src/session/legacy.rs @@ -73,6 +73,8 @@ pub fn load_session(session_name: &str, session_path: &Path) -> Result obj.entry("message_count").or_insert(serde_json::json!(0)); obj.entry("working_dir").or_insert(serde_json::json!("")); + crate::session::import_formats::nest_legacy_token_fields(obj); + if let Some(desc) = obj.get_mut("description") { if let Some(desc_str) = desc.as_str() { *desc = serde_json::json!(desc_str @@ -138,4 +140,24 @@ mod tests { assert_eq!(messages[0].role, Role::User); assert_eq!(messages[1].role, Role::Assistant); } + + #[test] + fn test_load_legacy_session_preserves_flat_token_fields() { + let temp_dir = TempDir::new().unwrap(); + let session_path = temp_dir.path().join("20240101_120000.jsonl"); + + let legacy_content = r#"{"description":"test","id":"20240101_120000","created_at":"2024-01-01T12:00:00Z","updated_at":"2024-01-01T12:00:00Z","extension_data":{},"message_count":0,"input_tokens":11,"output_tokens":22,"total_tokens":33,"accumulated_input_tokens":111,"accumulated_output_tokens":222,"accumulated_total_tokens":333} +{"id":"msg1","role":"user","created":1704110400,"content":[{"type":"text","text":"Hello"}]}"#; + + fs::write(&session_path, legacy_content).unwrap(); + + let session = load_session("20240101_120000", &session_path).unwrap(); + + assert_eq!(session.usage.input_tokens, Some(11)); + assert_eq!(session.usage.output_tokens, Some(22)); + assert_eq!(session.usage.total_tokens, Some(33)); + assert_eq!(session.accumulated_usage.input_tokens, Some(111)); + assert_eq!(session.accumulated_usage.output_tokens, Some(222)); + assert_eq!(session.accumulated_usage.total_tokens, Some(333)); + } } diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 3ac0b6cfe..1abdf1497 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -1,6 +1,6 @@ use crate::config::paths::Paths; use crate::config::GooseMode; -use crate::conversation::message::Message; +use crate::conversation::message::{Message, TokenState}; use crate::conversation::Conversation; use crate::providers::base::Provider; use crate::recipe::Recipe; @@ -10,6 +10,7 @@ use crate::session::session_naming::{ }; use anyhow::Result; use chrono::{DateTime, Utc}; +use goose_providers::conversation::token_usage::Usage; use goose_providers::model::ModelConfig; use rmcp::model::Role; use serde::{Deserialize, Serialize}; @@ -22,7 +23,7 @@ use std::sync::{Arc, LazyLock}; use tracing::{info, warn}; use utoipa::ToSchema; -pub const CURRENT_SCHEMA_VERSION: i32 = 13; +pub const CURRENT_SCHEMA_VERSION: i32 = 14; pub const SESSIONS_FOLDER: &str = "sessions"; pub const DB_NAME: &str = "sessions.db"; @@ -69,12 +70,10 @@ pub struct Session { pub created_at: DateTime, pub updated_at: DateTime, pub extension_data: ExtensionData, - pub total_tokens: Option, - pub input_tokens: Option, - pub output_tokens: Option, - pub accumulated_total_tokens: Option, - pub accumulated_input_tokens: Option, - pub accumulated_output_tokens: Option, + #[serde(default)] + pub usage: Usage, + #[serde(default)] + pub accumulated_usage: Usage, pub accumulated_cost: Option, pub schedule_id: Option, pub recipe: Option, @@ -93,6 +92,30 @@ pub struct Session { pub last_message_snippet: Option, } +impl From<&Session> for TokenState { + fn from(session: &Session) -> Self { + Self { + input_tokens: session.usage.input_tokens.unwrap_or(0), + output_tokens: session.usage.output_tokens.unwrap_or(0), + total_tokens: session.usage.total_tokens.unwrap_or(0), + cache_read_tokens: session.usage.cache_read_input_tokens.unwrap_or(0), + cache_write_tokens: session.usage.cache_write_input_tokens.unwrap_or(0), + accumulated_input_tokens: session.accumulated_usage.input_tokens.unwrap_or(0), + accumulated_output_tokens: session.accumulated_usage.output_tokens.unwrap_or(0), + accumulated_total_tokens: session.accumulated_usage.total_tokens.unwrap_or(0), + accumulated_cache_read_tokens: session + .accumulated_usage + .cache_read_input_tokens + .unwrap_or(0), + accumulated_cache_write_tokens: session + .accumulated_usage + .cache_write_input_tokens + .unwrap_or(0), + accumulated_cost: session.accumulated_cost, + } + } +} + pub struct SessionUpdateBuilder<'a> { session_manager: &'a SessionManager, session_id: String, @@ -101,12 +124,8 @@ pub struct SessionUpdateBuilder<'a> { session_type: Option, working_dir: Option, extension_data: Option, - total_tokens: Option>, - input_tokens: Option>, - output_tokens: Option>, - accumulated_total_tokens: Option>, - accumulated_input_tokens: Option>, - accumulated_output_tokens: Option>, + usage: Option, + accumulated_usage: Option, accumulated_cost: Option>, schedule_id: Option>, recipe: Option>, @@ -136,12 +155,8 @@ impl<'a> SessionUpdateBuilder<'a> { session_type: None, working_dir: None, extension_data: None, - total_tokens: None, - input_tokens: None, - output_tokens: None, - accumulated_total_tokens: None, - accumulated_input_tokens: None, - accumulated_output_tokens: None, + usage: None, + accumulated_usage: None, accumulated_cost: None, schedule_id: None, recipe: None, @@ -191,33 +206,13 @@ impl<'a> SessionUpdateBuilder<'a> { self } - pub fn total_tokens(mut self, tokens: Option) -> Self { - self.total_tokens = Some(tokens); + pub fn usage(mut self, usage: Usage) -> Self { + self.usage = Some(usage); self } - pub fn input_tokens(mut self, tokens: Option) -> Self { - self.input_tokens = Some(tokens); - self - } - - pub fn output_tokens(mut self, tokens: Option) -> Self { - self.output_tokens = Some(tokens); - self - } - - pub fn accumulated_total_tokens(mut self, tokens: Option) -> Self { - self.accumulated_total_tokens = Some(tokens); - self - } - - pub fn accumulated_input_tokens(mut self, tokens: Option) -> Self { - self.accumulated_input_tokens = Some(tokens); - self - } - - pub fn accumulated_output_tokens(mut self, tokens: Option) -> Self { - self.accumulated_output_tokens = Some(tokens); + pub fn accumulated_usage(mut self, usage: Usage) -> Self { + self.accumulated_usage = Some(usage); self } @@ -591,12 +586,8 @@ impl Default for Session { created_at: Default::default(), updated_at: Default::default(), extension_data: ExtensionData::default(), - total_tokens: None, - input_tokens: None, - output_tokens: None, - accumulated_total_tokens: None, - accumulated_input_tokens: None, - accumulated_output_tokens: None, + usage: Usage::default(), + accumulated_usage: Usage::default(), accumulated_cost: None, schedule_id: None, recipe: None, @@ -660,12 +651,26 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session { updated_at: row.try_get("updated_at")?, extension_data: serde_json::from_str(&row.try_get::("extension_data")?) .unwrap_or_default(), - total_tokens: row.try_get("total_tokens")?, - input_tokens: row.try_get("input_tokens")?, - output_tokens: row.try_get("output_tokens")?, - accumulated_total_tokens: row.try_get("accumulated_total_tokens")?, - accumulated_input_tokens: row.try_get("accumulated_input_tokens")?, - accumulated_output_tokens: row.try_get("accumulated_output_tokens")?, + usage: Usage { + input_tokens: row.try_get("input_tokens")?, + output_tokens: row.try_get("output_tokens")?, + total_tokens: row.try_get("total_tokens")?, + cache_read_input_tokens: row.try_get("cache_read_tokens").ok().flatten(), + cache_write_input_tokens: row.try_get("cache_write_tokens").ok().flatten(), + }, + accumulated_usage: Usage { + input_tokens: row.try_get("accumulated_input_tokens")?, + output_tokens: row.try_get("accumulated_output_tokens")?, + total_tokens: row.try_get("accumulated_total_tokens")?, + cache_read_input_tokens: row + .try_get("accumulated_cache_read_tokens") + .ok() + .flatten(), + cache_write_input_tokens: row + .try_get("accumulated_cache_write_tokens") + .ok() + .flatten(), + }, accumulated_cost: row.try_get("accumulated_cost").ok().flatten(), schedule_id: row.try_get("schedule_id")?, recipe, @@ -788,9 +793,13 @@ impl SessionStorage { total_tokens INTEGER, input_tokens INTEGER, output_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, accumulated_total_tokens INTEGER, accumulated_input_tokens INTEGER, accumulated_output_tokens INTEGER, + accumulated_cache_read_tokens INTEGER, + accumulated_cache_write_tokens INTEGER, accumulated_cost REAL, schedule_id TEXT, recipe_json TEXT, @@ -917,11 +926,13 @@ impl SessionStorage { INSERT INTO sessions ( id, name, user_set_name, session_type, working_dir, created_at, updated_at, extension_data, total_tokens, input_tokens, output_tokens, + cache_read_tokens, cache_write_tokens, accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens, + accumulated_cache_read_tokens, accumulated_cache_write_tokens, accumulated_cost, schedule_id, recipe_json, user_recipe_values_json, provider_name, model_config_json, goose_mode - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) "#, ) .bind(&session.id) @@ -932,12 +943,16 @@ impl SessionStorage { .bind(session.created_at) .bind(session.updated_at) .bind(serde_json::to_string(&session.extension_data)?) - .bind(session.total_tokens) - .bind(session.input_tokens) - .bind(session.output_tokens) - .bind(session.accumulated_total_tokens) - .bind(session.accumulated_input_tokens) - .bind(session.accumulated_output_tokens) + .bind(session.usage.total_tokens) + .bind(session.usage.input_tokens) + .bind(session.usage.output_tokens) + .bind(session.usage.cache_read_input_tokens) + .bind(session.usage.cache_write_input_tokens) + .bind(session.accumulated_usage.total_tokens) + .bind(session.accumulated_usage.input_tokens) + .bind(session.accumulated_usage.output_tokens) + .bind(session.accumulated_usage.cache_read_input_tokens) + .bind(session.accumulated_usage.cache_write_input_tokens) .bind(session.accumulated_cost) .bind(&session.schedule_id) .bind(recipe_json) @@ -1234,6 +1249,27 @@ impl SessionStorage { .await?; } } + 14 => { + for column in [ + "cache_read_tokens", + "cache_write_tokens", + "accumulated_cache_read_tokens", + "accumulated_cache_write_tokens", + ] { + let has_column = sqlx::query_scalar::<_, i32>( + "SELECT COUNT(*) FROM pragma_table_info('sessions') WHERE name = ?", + ) + .bind(column) + .fetch_one(&mut **tx) + .await? + > 0; + if !has_column { + sqlx::query(&format!("ALTER TABLE sessions ADD COLUMN {column} INTEGER")) + .execute(&mut **tx) + .await?; + } + } + } _ => { anyhow::bail!("Unknown migration version: {}", version); } @@ -1293,7 +1329,9 @@ impl SessionStorage { r#" SELECT id, working_dir, name, description, user_set_name, session_type, created_at, updated_at, extension_data, total_tokens, input_tokens, output_tokens, + cache_read_tokens, cache_write_tokens, accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens, + accumulated_cache_read_tokens, accumulated_cache_write_tokens, accumulated_cost, schedule_id, recipe_json, user_recipe_values_json, provider_name, model_config_json, goose_mode, @@ -1346,15 +1384,16 @@ impl SessionStorage { add_update!(builder.session_type, "session_type"); add_update!(builder.working_dir, "working_dir"); add_update!(builder.extension_data, "extension_data"); - add_update!(builder.total_tokens, "total_tokens"); - add_update!(builder.input_tokens, "input_tokens"); - add_update!(builder.output_tokens, "output_tokens"); - add_update!(builder.accumulated_total_tokens, "accumulated_total_tokens"); - add_update!(builder.accumulated_input_tokens, "accumulated_input_tokens"); - add_update!( - builder.accumulated_output_tokens, - "accumulated_output_tokens" - ); + add_update!(builder.usage, "total_tokens"); + add_update!(builder.usage, "input_tokens"); + add_update!(builder.usage, "output_tokens"); + add_update!(builder.usage, "cache_read_tokens"); + add_update!(builder.usage, "cache_write_tokens"); + add_update!(builder.accumulated_usage, "accumulated_total_tokens"); + add_update!(builder.accumulated_usage, "accumulated_input_tokens"); + add_update!(builder.accumulated_usage, "accumulated_output_tokens"); + add_update!(builder.accumulated_usage, "accumulated_cache_read_tokens"); + add_update!(builder.accumulated_usage, "accumulated_cache_write_tokens"); add_update!(builder.accumulated_cost, "accumulated_cost"); add_update!(builder.schedule_id, "schedule_id"); add_update!(builder.recipe, "recipe_json"); @@ -1390,23 +1429,21 @@ impl SessionStorage { if let Some(ed) = builder.extension_data { q = q.bind(serde_json::to_string(&ed)?); } - if let Some(tt) = builder.total_tokens { - q = q.bind(tt); + if let Some(u) = builder.usage { + q = q + .bind(u.total_tokens) + .bind(u.input_tokens) + .bind(u.output_tokens) + .bind(u.cache_read_input_tokens) + .bind(u.cache_write_input_tokens); } - if let Some(it) = builder.input_tokens { - q = q.bind(it); - } - if let Some(ot) = builder.output_tokens { - q = q.bind(ot); - } - if let Some(att) = builder.accumulated_total_tokens { - q = q.bind(att); - } - if let Some(ait) = builder.accumulated_input_tokens { - q = q.bind(ait); - } - if let Some(aot) = builder.accumulated_output_tokens { - q = q.bind(aot); + if let Some(u) = builder.accumulated_usage { + q = q + .bind(u.total_tokens) + .bind(u.input_tokens) + .bind(u.output_tokens) + .bind(u.cache_read_input_tokens) + .bind(u.cache_write_input_tokens); } if let Some(ac) = builder.accumulated_cost { q = q.bind(ac); @@ -1626,7 +1663,9 @@ impl SessionStorage { r#" SELECT s.id, s.working_dir, s.name, s.description, s.user_set_name, s.session_type, s.created_at, s.updated_at, s.extension_data, s.total_tokens, s.input_tokens, s.output_tokens, + s.cache_read_tokens, s.cache_write_tokens, s.accumulated_total_tokens, s.accumulated_input_tokens, s.accumulated_output_tokens, + s.accumulated_cache_read_tokens, s.accumulated_cache_write_tokens, s.accumulated_cost, s.schedule_id, s.recipe_json, s.user_recipe_values_json, s.provider_name, s.model_config_json, s.goose_mode, @@ -1816,12 +1855,8 @@ impl SessionStorage { let mut builder = session_manager .update(&session.id) .extension_data(import.extension_data) - .total_tokens(import.total_tokens) - .input_tokens(import.input_tokens) - .output_tokens(import.output_tokens) - .accumulated_total_tokens(import.accumulated_total_tokens) - .accumulated_input_tokens(import.accumulated_input_tokens) - .accumulated_output_tokens(import.accumulated_output_tokens) + .usage(import.usage) + .accumulated_usage(import.accumulated_usage) .accumulated_cost(import.accumulated_cost) .schedule_id(import.schedule_id) .recipe(import.recipe) @@ -2874,14 +2909,14 @@ mod tests { sm.update(&session.id) .user_provided_name(format!("Updated session {}", i)) - .total_tokens(Some(100 * i)) + .usage(Usage::new(None, None, Some(100 * i))) .apply() .await .unwrap(); let updated = sm.get_session(&session.id, true).await.unwrap(); assert_eq!(updated.message_count, 2); - assert_eq!(updated.total_tokens, Some(100 * i)); + assert_eq!(updated.usage.total_tokens, Some(100 * i)); session.id }); @@ -2915,13 +2950,14 @@ mod tests { #[tokio::test] async fn test_export_import_roundtrip() { const DESCRIPTION: &str = "Original session"; - const TOTAL_TOKENS: i32 = 500; - const INPUT_TOKENS: i32 = 300; - const OUTPUT_TOKENS: i32 = 200; - const ACCUMULATED_TOKENS: i32 = 1000; const USER_MESSAGE: &str = "test message"; const ASSISTANT_MESSAGE: &str = "test response"; + let usage = + Usage::new(Some(300), Some(200), Some(500)).with_cache_tokens(Some(120), Some(80)); + let accumulated_usage = + Usage::new(Some(600), Some(400), Some(1000)).with_cache_tokens(Some(400), Some(150)); + let temp_dir = TempDir::new().unwrap(); let sm = SessionManager::new(temp_dir.path().to_path_buf()); @@ -2936,10 +2972,8 @@ mod tests { .unwrap(); sm.update(&original.id) - .total_tokens(Some(TOTAL_TOKENS)) - .input_tokens(Some(INPUT_TOKENS)) - .output_tokens(Some(OUTPUT_TOKENS)) - .accumulated_total_tokens(Some(ACCUMULATED_TOKENS)) + .usage(usage) + .accumulated_usage(accumulated_usage) .apply() .await .unwrap(); @@ -2976,10 +3010,8 @@ mod tests { assert_ne!(imported.id, original.id); assert_eq!(imported.name, DESCRIPTION); assert_eq!(imported.working_dir, PathBuf::from("/tmp/test")); - assert_eq!(imported.total_tokens, Some(TOTAL_TOKENS)); - assert_eq!(imported.input_tokens, Some(INPUT_TOKENS)); - assert_eq!(imported.output_tokens, Some(OUTPUT_TOKENS)); - assert_eq!(imported.accumulated_total_tokens, Some(ACCUMULATED_TOKENS)); + assert_eq!(imported.usage, usage); + assert_eq!(imported.accumulated_usage, accumulated_usage); assert_eq!(imported.message_count, 2); let conversation = imported.conversation.unwrap(); @@ -3052,7 +3084,7 @@ mod tests { } #[tokio::test] - async fn test_import_session_with_description_field() { + async fn test_import_session_with_legacy_flat_fields() { const OLD_FORMAT_JSON: &str = r#"{ "id": "20240101_1", "description": "Old format session", @@ -3061,7 +3093,14 @@ mod tests { "created_at": "2024-01-01T00:00:00Z", "updated_at": "2024-01-01T00:00:00Z", "extension_data": {}, - "message_count": 0 + "message_count": 0, + "total_tokens": 500, + "input_tokens": 300, + "output_tokens": 200, + "cache_read_tokens": 120, + "accumulated_total_tokens": 1000, + "accumulated_input_tokens": 600, + "accumulated_output_tokens": 400 }"#; let temp_dir = TempDir::new().unwrap(); @@ -3072,6 +3111,14 @@ mod tests { assert_eq!(imported.name, "Old format session"); assert!(imported.user_set_name); assert_eq!(imported.working_dir, PathBuf::from("/tmp/test")); + assert_eq!( + imported.usage, + Usage::new(Some(300), Some(200), Some(500)).with_cache_tokens(Some(120), None) + ); + assert_eq!( + imported.accumulated_usage, + Usage::new(Some(600), Some(400), Some(1000)) + ); } #[test_case(GooseMode::Approve)] @@ -3215,4 +3262,78 @@ mod tests { let acp_session = sm.storage().get_session("acp_id", false).await.unwrap(); assert_eq!(acp_session.session_type, SessionType::Acp); } + + #[tokio::test] + async fn test_cache_token_columns_migration_and_round_trip() { + let temp_dir = TempDir::new().unwrap(); + let db_path = temp_dir.path().join(SESSIONS_FOLDER).join(DB_NAME); + + if let Some(parent) = db_path.parent() { + std::fs::create_dir_all(parent).unwrap(); + } + + let pool = SqlitePoolOptions::new() + .connect_with( + SqliteConnectOptions::new() + .filename(&db_path) + .create_if_missing(true), + ) + .await + .unwrap(); + + SessionStorage::create_schema(&pool).await.unwrap(); + + // Recreate a v13-shaped database without cache token columns. + for column in [ + "cache_read_tokens", + "cache_write_tokens", + "accumulated_cache_read_tokens", + "accumulated_cache_write_tokens", + ] { + sqlx::query(&format!("ALTER TABLE sessions DROP COLUMN {column}")) + .execute(&pool) + .await + .unwrap(); + } + sqlx::query("UPDATE schema_version SET version = 13") + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode) + VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind("cache_id") + .bind("Cache Session") + .bind(false) + .bind("user") + .bind("/tmp") + .bind("{}") + .bind("auto") + .execute(&pool) + .await + .unwrap(); + + pool.close().await; + + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + sm.storage().pool().await.unwrap(); // Triggers migration + + let usage = + Usage::new(Some(8000), Some(500), None).with_cache_tokens(Some(5000), Some(1000)); + let accumulated_usage = + Usage::new(Some(24000), Some(1500), None).with_cache_tokens(Some(15000), Some(3000)); + + sm.update("cache_id") + .usage(usage) + .accumulated_usage(accumulated_usage) + .apply() + .await + .unwrap(); + + let loaded = sm.get_session("cache_id", false).await.unwrap(); + assert_eq!(loaded.usage, usage); + assert_eq!(loaded.accumulated_usage, accumulated_usage); + } } diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 25e5c196d..347629af3 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -1595,12 +1595,12 @@ mod tests { run_turn(&agent, &session_id, "Turn 1").await?; let after_1 = session_manager.get_session(&session_id, false).await?; - assert_eq!(after_1.accumulated_total_tokens, Some(15)); + assert_eq!(after_1.accumulated_usage.total_tokens, Some(15)); run_turn(&agent, &session_id, "Turn 2").await?; let after_2 = session_manager.get_session(&session_id, false).await?; - assert_eq!(after_2.accumulated_total_tokens, Some(30)); - assert_eq!(after_2.total_tokens, Some(15)); + assert_eq!(after_2.accumulated_usage.total_tokens, Some(30)); + assert_eq!(after_2.usage.total_tokens, Some(15)); Ok(()) } diff --git a/crates/goose/tests/compaction.rs b/crates/goose/tests/compaction.rs index ad9dd39b1..dfb4ea780 100644 --- a/crates/goose/tests/compaction.rs +++ b/crates/goose/tests/compaction.rs @@ -235,12 +235,8 @@ async fn setup_test_session( .config .session_manager .update(&session.id) - .total_tokens(Some(1000)) - .input_tokens(Some(600)) - .output_tokens(Some(400)) - .accumulated_total_tokens(Some(1000)) - .accumulated_input_tokens(Some(600)) - .accumulated_output_tokens(Some(400)) + .usage(Usage::new(Some(600), Some(400), Some(1000))) + .accumulated_usage(Usage::new(Some(600), Some(400), Some(1000))) .apply() .await?; @@ -384,16 +380,16 @@ async fn test_manual_compaction_updates_token_counts_and_conversation() -> Resul // After compaction, the current context is ONLY the summary (200 tokens) // This is the new agent-visible input context assert_eq!( - updated_session.input_tokens, + updated_session.usage.input_tokens, Some(expected_summary_output), "Input tokens should be exactly the summary output (200 tokens)" ); assert_eq!( - updated_session.output_tokens, None, + updated_session.usage.output_tokens, None, "Output tokens should be None after compaction (no new assistant output)" ); assert_eq!( - updated_session.total_tokens, + updated_session.usage.total_tokens, Some(expected_summary_output), "Total should equal input (200 tokens) after compaction" ); @@ -403,7 +399,7 @@ async fn test_manual_compaction_updates_token_counts_and_conversation() -> Resul // Compaction input: ~6400 (system 6000 + 4 messages ~400) // Compaction output: 200 // Expected accumulated: 1000 + 6400 + 200 = 7600 - let accumulated = updated_session.accumulated_total_tokens.unwrap(); + let accumulated = updated_session.accumulated_usage.total_tokens.unwrap(); assert!( (7300..=7900).contains(&accumulated), "Accumulated should be ~7600 (1000 initial + 6400 input + 200 output). Got: {}", @@ -442,7 +438,7 @@ async fn test_auto_compaction_during_reply() -> Result<()> { .session_manager .get_session(&session.id, true) .await?; - let initial_input_tokens = initial_session.input_tokens.unwrap_or(0); + let initial_input_tokens = initial_session.usage.input_tokens.unwrap_or(0); // Setup mock provider (no context limit enforcement) let provider = Arc::new(MockCompactionProvider::new()); @@ -479,7 +475,7 @@ async fn test_auto_compaction_during_reply() -> Result<()> { .session_manager .get_session(&session.id, true) .await?; - input_tokens_after_compaction = session_after_compact.input_tokens; + input_tokens_after_compaction = session_after_compact.usage.input_tokens; } Ok(_) => {} Err(e) => return Err(e), @@ -515,9 +511,9 @@ async fn test_auto_compaction_during_reply() -> Result<()> { // After the subsequent reply, the current window includes: // - system (6000) + summary (200) + new user message (100) + reply (100) = 6400 - let final_input = updated_session.input_tokens.unwrap(); - let final_output = updated_session.output_tokens.unwrap(); - let final_total = updated_session.total_tokens.unwrap(); + let final_input = updated_session.usage.input_tokens.unwrap(); + let final_output = updated_session.usage.output_tokens.unwrap(); + let final_total = updated_session.usage.total_tokens.unwrap(); assert!( final_input >= 6000, @@ -540,7 +536,7 @@ async fn test_auto_compaction_during_reply() -> Result<()> { // - Compaction: ~10,400 input + 200 output = 10,600 // - Reply: ~6,300 input + 100 output = 6,400 // Total: 1000 + 10,600 + 6,400 = 18,000 - let accumulated = updated_session.accumulated_total_tokens.unwrap(); + let accumulated = updated_session.accumulated_usage.total_tokens.unwrap(); assert!( (17000..=19000).contains(&accumulated), "Accumulated should be ~18,000 (initial + compaction + reply). Got: {}", @@ -552,7 +548,7 @@ async fn test_auto_compaction_during_reply() -> Result<()> { // - Reply: system (6000) + 40 messages (4000) + new message (100) = 10,100 input // - Reply output: 100 // Total: 1000 + 10,100 + 100 = 11,200 - let accumulated = updated_session.accumulated_total_tokens.unwrap(); + let accumulated = updated_session.accumulated_usage.total_tokens.unwrap(); assert!( (11000..=11500).contains(&accumulated), "Accumulated should be ~11,200 (initial + reply). Got: {}", @@ -560,8 +556,8 @@ async fn test_auto_compaction_during_reply() -> Result<()> { ); // Current window should be: 10,100 input + 100 output = 10,200 - let final_input = updated_session.input_tokens.unwrap(); - let final_output = updated_session.output_tokens.unwrap(); + let final_input = updated_session.usage.input_tokens.unwrap(); + let final_output = updated_session.usage.output_tokens.unwrap(); assert!( (10000..=10500).contains(&final_input), @@ -637,7 +633,7 @@ async fn test_context_limit_recovery_compaction() -> Result<()> { .session_manager .get_session(&session.id, true) .await?; - input_tokens_after_compaction = session_after_compact.input_tokens; + input_tokens_after_compaction = session_after_compact.usage.input_tokens; } Ok(AgentEvent::Message(msg)) => { // Check if we got a real response (not just a notification) @@ -704,9 +700,9 @@ async fn test_context_limit_recovery_compaction() -> Result<()> { // Check the final token state after recovery // Note: The session state reflects the RETRY call (after compaction), // which only sees agent-visible messages (summary + continuation + user message) - let final_input = updated_session.input_tokens.unwrap(); - let final_output = updated_session.output_tokens; - let final_total = updated_session.total_tokens.unwrap(); + let final_input = updated_session.usage.input_tokens.unwrap(); + let final_output = updated_session.usage.output_tokens; + let final_total = updated_session.usage.total_tokens.unwrap(); // After compaction, the retry only sees agent-visible messages: // Input: system (6000) + summary (~100) + continuation (~100) + user message (~100) = ~6300 @@ -736,7 +732,7 @@ async fn test_context_limit_recovery_compaction() -> Result<()> { // - Compaction: ~6400 input (mock uses system_prompt.len()/4) + 200 output = ~6600 // - Reply: ~6500 input + 200 output = ~6700 // Total: 1000 + 6600 + 6700 = ~14300 - let accumulated = updated_session.accumulated_total_tokens.unwrap(); + let accumulated = updated_session.accumulated_usage.total_tokens.unwrap(); assert!( (13000..=16000).contains(&accumulated), "Accumulated should be ~14300 (initial + compaction + reply). Got: {}", diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 93466b85c..ec5898090 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -8037,20 +8037,8 @@ "format": "double", "nullable": true }, - "accumulated_input_tokens": { - "type": "integer", - "format": "int32", - "nullable": true - }, - "accumulated_output_tokens": { - "type": "integer", - "format": "int32", - "nullable": true - }, - "accumulated_total_tokens": { - "type": "integer", - "format": "int32", - "nullable": true + "accumulated_usage": { + "$ref": "#/components/schemas/Usage" }, "archived_at": { "type": "string", @@ -8078,11 +8066,6 @@ "id": { "type": "string" }, - "input_tokens": { - "type": "integer", - "format": "int32", - "nullable": true - }, "last_message_snippet": { "type": "string", "nullable": true @@ -8102,11 +8085,6 @@ "name": { "type": "string" }, - "output_tokens": { - "type": "integer", - "format": "int32", - "nullable": true - }, "project_id": { "type": "string", "nullable": true @@ -8130,15 +8108,13 @@ "session_type": { "$ref": "#/components/schemas/SessionType" }, - "total_tokens": { - "type": "integer", - "format": "int32", - "nullable": true - }, "updated_at": { "type": "string", "format": "date-time" }, + "usage": { + "$ref": "#/components/schemas/Usage" + }, "user_recipe_values": { "type": "object", "additionalProperties": { @@ -8718,6 +8694,14 @@ "accumulatedTotalTokens" ], "properties": { + "accumulatedCacheReadTokens": { + "type": "integer", + "format": "int32" + }, + "accumulatedCacheWriteTokens": { + "type": "integer", + "format": "int32" + }, "accumulatedCost": { "type": "number", "format": "double", @@ -8735,6 +8719,14 @@ "type": "integer", "format": "int32" }, + "cacheReadTokens": { + "type": "integer", + "format": "int32" + }, + "cacheWriteTokens": { + "type": "integer", + "format": "int32" + }, "inputTokens": { "type": "integer", "format": "int32" @@ -9258,6 +9250,38 @@ } } }, + "Usage": { + "type": "object", + "description": "`input_tokens` is the total input including cache read/write tokens;\nthe cache fields are breakdown subsets of it. Parsers for providers\nthat report cache tokens separately from input (e.g. Anthropic,\nBedrock) must fold them into `input_tokens`.", + "properties": { + "cache_read_input_tokens": { + "type": "integer", + "format": "int32", + "nullable": true + }, + "cache_write_input_tokens": { + "type": "integer", + "format": "int32", + "nullable": true + }, + "input_tokens": { + "type": "integer", + "format": "int32", + "description": "All prompt tokens, including any served from or written to cache.\n`cache_read_input_tokens` and `cache_write_input_tokens` are subsets of this.", + "nullable": true + }, + "output_tokens": { + "type": "integer", + "format": "int32", + "nullable": true + }, + "total_tokens": { + "type": "integer", + "format": "int32", + "nullable": true + } + } + }, "WhisperModelResponse": { "allOf": [ { diff --git a/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts b/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts index 3fc3dc630..e9bca9eb6 100644 --- a/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts +++ b/ui/desktop/src/acp/__tests__/chatSessionStore.test.ts @@ -34,12 +34,8 @@ function session(id: string, conversation: Message[] = []): Session { extension_data: {}, source: 'test', conversation, - input_tokens: 1, - output_tokens: 2, - total_tokens: 3, - accumulated_input_tokens: 4, - accumulated_output_tokens: 5, - accumulated_total_tokens: 9, + usage: { input_tokens: 1, output_tokens: 2, total_tokens: 3 }, + accumulated_usage: { input_tokens: 4, output_tokens: 5, total_tokens: 9 }, } as Session; } @@ -161,12 +157,12 @@ describe('acpChatSessionStore', () => { expect(snapshot.session?.id).toBe(currentSessionId); expect(snapshot.messages).toEqual([initialMessage]); expect(snapshot.tokenState).toMatchObject({ - inputTokens: 0, - outputTokens: 0, - totalTokens: 0, - accumulatedInputTokens: 0, - accumulatedOutputTokens: 0, - accumulatedTotalTokens: 0, + inputTokens: 1, + outputTokens: 2, + totalTokens: 3, + accumulatedInputTokens: 4, + accumulatedOutputTokens: 5, + accumulatedTotalTokens: 9, }); expect(snapshot.chatState).toBe(ChatState.Idle); expect(snapshot.sessionLoadError).toBeUndefined(); diff --git a/ui/desktop/src/acp/chatSessionStore.ts b/ui/desktop/src/acp/chatSessionStore.ts index 7dfc62608..70cee7f13 100644 --- a/ui/desktop/src/acp/chatSessionStore.ts +++ b/ui/desktop/src/acp/chatSessionStore.ts @@ -38,6 +38,22 @@ const initialTokenState: TokenState = { accumulatedTotalTokens: 0, }; +function tokenStateFromSession(session: Session | undefined): Partial { + return { + inputTokens: session?.usage?.input_tokens ?? 0, + outputTokens: session?.usage?.output_tokens ?? 0, + totalTokens: session?.usage?.total_tokens ?? 0, + cacheReadTokens: session?.usage?.cache_read_input_tokens ?? 0, + cacheWriteTokens: session?.usage?.cache_write_input_tokens ?? 0, + accumulatedInputTokens: session?.accumulated_usage?.input_tokens ?? 0, + accumulatedOutputTokens: session?.accumulated_usage?.output_tokens ?? 0, + accumulatedTotalTokens: session?.accumulated_usage?.total_tokens ?? 0, + accumulatedCacheReadTokens: session?.accumulated_usage?.cache_read_input_tokens ?? 0, + accumulatedCacheWriteTokens: session?.accumulated_usage?.cache_write_input_tokens ?? 0, + ...(session?.accumulated_cost != null ? { accumulatedCost: session.accumulated_cost } : {}), + }; +} + export interface AcpChatSessionStore { getSnapshot(sessionId: string): AcpChatSessionSnapshot | undefined; } @@ -164,6 +180,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const finishSessionLoad: AcpChatSessionActions['finishSessionLoad'] = (sessionId, session) => { const entry = getOrCreateEntry(sessionId); entry.session = session; + entry.tokenState = { ...entry.tokenState, ...tokenStateFromSession(session) }; entry.sessionLoadError = undefined; entry.chatState = entry.activePromptAttemptId ? ChatState.Streaming : ChatState.Idle; return notify(sessionId, entry); diff --git a/ui/desktop/src/api/index.ts b/ui/desktop/src/api/index.ts index 3eedb831d..f1de1c86a 100644 --- a/ui/desktop/src/api/index.ts +++ b/ui/desktop/src/api/index.ts @@ -1,4 +1,4 @@ // This file is auto-generated by @hey-api/openapi-ts export { addExtension, agentAddExtension, agentRemoveExtension, callTool, cancelDownload, cancelLocalModelDownload, checkProvider, cleanupProviderCache, configureProviderOauth, confirmToolAction, createCustomProvider, createRecipe, createSchedule, decodeRecipe, deleteLocalModel, deleteModel, deleteProviderSecret, deleteRecipe, deleteSchedule, diagnostics, downloadHfModel, downloadModel, encodeRecipe, exportApp, forkSession, getCanonicalModelInfo, getCustomProvider, getDictationConfig, getDownloadProgress, getExtensions, getFeatures, getLocalModelDownloadProgress, getModelSettings, getPrompt, getPrompts, getProviderCatalog, getProviderCatalogTemplate, getProviderModelInfo, getProviderModels, getRepoFiles, getSession, getSessionExtensions, getSlashCommands, getTools, getTunnelStatus, importApp, importSessionNostr, inspectRunningJob, killRunningJob, listApps, listBuiltinChatTemplates, listLocalModels, listModels, listProviderSecrets, listRecipes, listSchedules, mcpUiProxy, type Options, parseRecipe, pauseSchedule, providers, readAllConfig, readConfig, readResource, recipeToYaml, removeConfig, removeCustomProvider, removeExtension, reply, resetPrompt, restartAgent, resumeAgent, runNowHandler, savePrompt, saveRecipe, scanRecipe, scheduleRecipe, searchHfModels, sendTelemetryEvent, sessionCancel, sessionEvents, sessionReply, sessionsHandler, setConfigProvider, setRecipeSlashCommand, shareSessionNostr, startAgent, startNanogptSetup, startOpenrouterSetup, startTetrateSetup, startTunnel, status, stopAgent, stopTunnel, syncFeaturedModels, systemInfo, transcribeDictation, unpauseSchedule, updateAgentProvider, updateCustomProvider, updateFromSession, updateModelSettings, updateSchedule, updateSession, updateSessionName, updateSessionUserRecipeValues, updateWorkingDir, upsertConfig, upsertPermissions, validateConfig } from './sdk.gen'; -export type { ActionRequired, ActionRequiredData, AddExtensionData, AddExtensionErrors, AddExtensionRequest, AddExtensionResponse, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponse, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponse, AgentRemoveExtensionResponses, Annotations, Author, AuthorRequest, CallToolData, CallToolError, CallToolErrors, CallToolRequest, CallToolResponse, CallToolResponse2, CallToolResponses, CancelDownloadData, CancelDownloadErrors, CancelDownloadResponses, CancelLocalModelDownloadData, CancelLocalModelDownloadErrors, CancelLocalModelDownloadResponses, CancelRequest, ChatRequest, ChatTemplate, CheckProviderData, CheckProviderRequest, CleanupProviderCacheData, CleanupProviderCacheErrors, CleanupProviderCacheResponse, CleanupProviderCacheResponses, ClientOptions, CommandType, ConfigKey, ConfigKeyQuery, ConfigResponse, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionRequest, ConfirmToolActionResponses, Content, ContentBlock, Conversation, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponse, CreateCustomProviderResponse2, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeRequest, CreateRecipeResponse, CreateRecipeResponse2, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleRequest, CreateScheduleResponse, CreateScheduleResponses, CspMetadata, DeclarativeProviderConfig, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeRequest, DecodeRecipeResponse, DecodeRecipeResponse2, DecodeRecipeResponses, DeleteLocalModelData, DeleteLocalModelErrors, DeleteLocalModelResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteProviderSecretData, DeleteProviderSecretErrors, DeleteProviderSecretResponse, DeleteProviderSecretResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeRequest, DeleteRecipeResponse, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponse, DeleteScheduleResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponse, DiagnosticsResponses, DictationProvider, DictationProviderStatus, DownloadHfModelData, DownloadHfModelErrors, DownloadHfModelResponse, DownloadHfModelResponses, DownloadModelData, DownloadModelErrors, DownloadModelRequest, DownloadModelResponses, DownloadProgress, DownloadStatus, EmbeddedResource, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeRequest, EncodeRecipeResponse, EncodeRecipeResponse2, EncodeRecipeResponses, Envs, EnvVarConfig, ErrorResponse, ExportAppData, ExportAppError, ExportAppErrors, ExportAppResponse, ExportAppResponses, ExtensionConfig, ExtensionData, ExtensionEntry, ExtensionLoadResult, ExtensionQuery, ExtensionResponse, FeaturesResponse, ForkRequest, ForkResponse, ForkSessionData, ForkSessionErrors, ForkSessionResponse, ForkSessionResponses, FrontendToolRequest, GetCanonicalModelInfoData, GetCanonicalModelInfoResponse, GetCanonicalModelInfoResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponse, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponse, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponse, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponse, GetExtensionsResponses, GetFeaturesData, GetFeaturesResponse, GetFeaturesResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponse, GetLocalModelDownloadProgressResponses, GetModelSettingsData, GetModelSettingsErrors, GetModelSettingsResponse, GetModelSettingsResponses, GetPromptData, GetPromptErrors, GetPromptResponse, GetPromptResponses, GetPromptsData, GetPromptsResponse, GetPromptsResponses, GetProviderCatalogData, GetProviderCatalogErrors, GetProviderCatalogResponse, GetProviderCatalogResponses, GetProviderCatalogTemplateData, GetProviderCatalogTemplateErrors, GetProviderCatalogTemplateResponse, GetProviderCatalogTemplateResponses, GetProviderModelInfoData, GetProviderModelInfoErrors, GetProviderModelInfoResponse, GetProviderModelInfoResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponse, GetProviderModelsResponses, GetRepoFilesData, GetRepoFilesResponse, GetRepoFilesResponses, GetSessionData, GetSessionErrors, GetSessionExtensionsData, GetSessionExtensionsErrors, GetSessionExtensionsResponse, GetSessionExtensionsResponses, GetSessionResponse, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponse, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsQuery, GetToolsResponse, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponse, GetTunnelStatusResponses, GooseApp, GooseMode, HfGgufFile, HfModelInfo, HfModelVariant, HfQuantVariant, Icon, IconTheme, ImageContent, ImportAppData, ImportAppError, ImportAppErrors, ImportAppRequest, ImportAppResponse, ImportAppResponse2, ImportAppResponses, ImportSessionNostrData, ImportSessionNostrErrors, ImportSessionNostrRequest, ImportSessionNostrResponse, ImportSessionNostrResponses, InferenceMetadata, InspectJobResponse, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponse, InspectRunningJobResponses, JsonObject, KillJobResponse, KillRunningJobData, KillRunningJobResponses, ListAppsData, ListAppsError, ListAppsErrors, ListAppsRequest, ListAppsResponse, ListAppsResponse2, ListAppsResponses, ListBuiltinChatTemplatesData, ListBuiltinChatTemplatesResponse, ListBuiltinChatTemplatesResponses, ListLocalModelsData, ListLocalModelsResponse, ListLocalModelsResponses, ListModelsData, ListModelsResponse, ListModelsResponses, ListProviderSecretsData, ListProviderSecretsErrors, ListProviderSecretsResponse, ListProviderSecretsResponses, ListRecipeResponse, ListRecipesData, ListRecipesErrors, ListRecipesResponse, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponse, ListSchedulesResponse2, ListSchedulesResponses, LoadedProvider, LocalModelResponse, McpAppResource, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, Message, MessageContent, MessageEvent, MessageMetadata, ModelCapabilities, ModelConfig, ModelDownloadStatus, ModelInfo, ModelInfoData, ModelInfoQuery, ModelInfoResponse, ModelSettings, ModelTemplate, ParseRecipeData, ParseRecipeError, ParseRecipeErrors, ParseRecipeRequest, ParseRecipeResponse, ParseRecipeResponse2, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponse, PauseScheduleResponses, Permission, PermissionLevel, PermissionsMetadata, PrincipalType, PromptContentResponse, PromptsListResponse, ProviderCatalogEntry, ProviderDetails, ProviderEngine, ProviderMetadata, ProviderModelInfoQuery, ProvidersData, ProviderSecret, ProviderSecretsResponse, ProviderSecretStatus, ProviderSecretStorage, ProvidersResponse, ProvidersResponse2, ProvidersResponses, ProviderTemplate, ProviderType, RawAudioContent, RawEmbeddedResource, RawImageContent, RawResource, RawTextContent, ReadAllConfigData, ReadAllConfigResponse, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceRequest, ReadResourceResponse, ReadResourceResponse2, ReadResourceResponses, Recipe, RecipeManifest, RecipeParameter, RecipeParameterInputType, RecipeParameterRequirement, RecipeToYamlData, RecipeToYamlError, RecipeToYamlErrors, RecipeToYamlRequest, RecipeToYamlResponse, RecipeToYamlResponse2, RecipeToYamlResponses, RedactedThinkingContent, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponse, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponse, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionRequest, RemoveExtensionResponse, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponse, ReplyResponses, RepoVariantsResponse, ResetPromptData, ResetPromptErrors, ResetPromptResponse, ResetPromptResponses, ResourceContents, ResourceMetadata, Response, RestartAgentData, RestartAgentErrors, RestartAgentRequest, RestartAgentResponse, RestartAgentResponse2, RestartAgentResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentRequest, ResumeAgentResponse, ResumeAgentResponse2, ResumeAgentResponses, RetryConfig, Role, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponse, RunNowHandlerResponses, RunNowResponse, SamplingConfig, SavePromptData, SavePromptErrors, SavePromptRequest, SavePromptResponse, SavePromptResponses, SaveRecipeData, SaveRecipeError, SaveRecipeErrors, SaveRecipeRequest, SaveRecipeResponse, SaveRecipeResponse2, SaveRecipeResponses, ScanRecipeData, ScanRecipeRequest, ScanRecipeResponse, ScanRecipeResponse2, ScanRecipeResponses, ScheduledJob, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeRequest, ScheduleRecipeResponses, SearchHfModelsData, SearchHfModelsErrors, SearchHfModelsResponse, SearchHfModelsResponses, SendTelemetryEventData, SendTelemetryEventResponses, Session, SessionCancelData, SessionCancelResponses, SessionDisplayInfo, SessionEventsData, SessionEventsErrors, SessionEventsResponse, SessionEventsResponses, SessionExtensionsResponse, SessionReplyData, SessionReplyErrors, SessionReplyRequest, SessionReplyResponse, SessionReplyResponse2, SessionReplyResponses, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponse, SessionsHandlerResponses, SessionsQuery, SessionType, SetConfigProviderData, SetProviderRequest, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, SetSlashCommandRequest, Settings, SetupResponse, ShareSessionNostrData, ShareSessionNostrErrors, ShareSessionNostrRequest, ShareSessionNostrResponse, ShareSessionNostrResponse2, ShareSessionNostrResponses, SlashCommand, SlashCommandsResponse, StartAgentData, StartAgentError, StartAgentErrors, StartAgentRequest, StartAgentResponse, StartAgentResponses, StartNanogptSetupData, StartNanogptSetupResponse, StartNanogptSetupResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponse, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponse, StartTetrateSetupResponses, StartTunnelData, StartTunnelError, StartTunnelErrors, StartTunnelResponse, StartTunnelResponses, StatusData, StatusResponse, StatusResponses, StopAgentData, StopAgentErrors, StopAgentRequest, StopAgentResponse, StopAgentResponses, StopTunnelData, StopTunnelError, StopTunnelErrors, StopTunnelResponses, SubRecipe, SuccessCheck, SyncFeaturedModelsData, SyncFeaturedModelsResponses, SystemInfo, SystemInfoData, SystemInfoResponse, SystemInfoResponses, SystemNotificationContent, SystemNotificationType, TaskSupport, TelemetryEventRequest, Template, TextContent, ThinkingContent, ThinkingEffort, TokenState, Tool, ToolAnnotations, ToolCallingMode, ToolConfirmationRequest, ToolExecution, ToolInfo, ToolPermission, ToolRequest, ToolResponse, TranscribeDictationData, TranscribeDictationErrors, TranscribeDictationResponse, TranscribeDictationResponses, TranscribeRequest, TranscribeResponse, TunnelInfo, TunnelState, UiMetadata, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponse, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderRequest, UpdateCustomProviderResponse, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionRequest, UpdateFromSessionResponses, UpdateModelSettingsData, UpdateModelSettingsErrors, UpdateModelSettingsResponse, UpdateModelSettingsResponses, UpdateProviderRequest, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleRequest, UpdateScheduleResponse, UpdateScheduleResponses, UpdateSessionData, UpdateSessionErrors, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameRequest, UpdateSessionNameResponses, UpdateSessionRequest, UpdateSessionResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesError, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesRequest, UpdateSessionUserRecipeValuesResponse, UpdateSessionUserRecipeValuesResponse2, UpdateSessionUserRecipeValuesResponses, UpdateWorkingDirData, UpdateWorkingDirErrors, UpdateWorkingDirRequest, UpdateWorkingDirResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigQuery, UpsertConfigResponse, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsQuery, UpsertPermissionsResponse, UpsertPermissionsResponses, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponse, ValidateConfigResponses, WhisperModelResponse, WindowProps } from './types.gen'; +export type { ActionRequired, ActionRequiredData, AddExtensionData, AddExtensionErrors, AddExtensionRequest, AddExtensionResponse, AddExtensionResponses, AgentAddExtensionData, AgentAddExtensionErrors, AgentAddExtensionResponse, AgentAddExtensionResponses, AgentRemoveExtensionData, AgentRemoveExtensionErrors, AgentRemoveExtensionResponse, AgentRemoveExtensionResponses, Annotations, Author, AuthorRequest, CallToolData, CallToolError, CallToolErrors, CallToolRequest, CallToolResponse, CallToolResponse2, CallToolResponses, CancelDownloadData, CancelDownloadErrors, CancelDownloadResponses, CancelLocalModelDownloadData, CancelLocalModelDownloadErrors, CancelLocalModelDownloadResponses, CancelRequest, ChatRequest, ChatTemplate, CheckProviderData, CheckProviderRequest, CleanupProviderCacheData, CleanupProviderCacheErrors, CleanupProviderCacheResponse, CleanupProviderCacheResponses, ClientOptions, CommandType, ConfigKey, ConfigKeyQuery, ConfigResponse, ConfigureProviderOauthData, ConfigureProviderOauthErrors, ConfigureProviderOauthResponses, ConfirmToolActionData, ConfirmToolActionErrors, ConfirmToolActionRequest, ConfirmToolActionResponses, Content, ContentBlock, Conversation, CreateCustomProviderData, CreateCustomProviderErrors, CreateCustomProviderResponse, CreateCustomProviderResponse2, CreateCustomProviderResponses, CreateRecipeData, CreateRecipeErrors, CreateRecipeRequest, CreateRecipeResponse, CreateRecipeResponse2, CreateRecipeResponses, CreateScheduleData, CreateScheduleErrors, CreateScheduleRequest, CreateScheduleResponse, CreateScheduleResponses, CspMetadata, DeclarativeProviderConfig, DecodeRecipeData, DecodeRecipeErrors, DecodeRecipeRequest, DecodeRecipeResponse, DecodeRecipeResponse2, DecodeRecipeResponses, DeleteLocalModelData, DeleteLocalModelErrors, DeleteLocalModelResponses, DeleteModelData, DeleteModelErrors, DeleteModelResponses, DeleteProviderSecretData, DeleteProviderSecretErrors, DeleteProviderSecretResponse, DeleteProviderSecretResponses, DeleteRecipeData, DeleteRecipeErrors, DeleteRecipeRequest, DeleteRecipeResponse, DeleteRecipeResponses, DeleteScheduleData, DeleteScheduleErrors, DeleteScheduleResponse, DeleteScheduleResponses, DiagnosticsData, DiagnosticsErrors, DiagnosticsResponse, DiagnosticsResponses, DictationProvider, DictationProviderStatus, DownloadHfModelData, DownloadHfModelErrors, DownloadHfModelResponse, DownloadHfModelResponses, DownloadModelData, DownloadModelErrors, DownloadModelRequest, DownloadModelResponses, DownloadProgress, DownloadStatus, EmbeddedResource, EncodeRecipeData, EncodeRecipeErrors, EncodeRecipeRequest, EncodeRecipeResponse, EncodeRecipeResponse2, EncodeRecipeResponses, Envs, EnvVarConfig, ErrorResponse, ExportAppData, ExportAppError, ExportAppErrors, ExportAppResponse, ExportAppResponses, ExtensionConfig, ExtensionData, ExtensionEntry, ExtensionLoadResult, ExtensionQuery, ExtensionResponse, FeaturesResponse, ForkRequest, ForkResponse, ForkSessionData, ForkSessionErrors, ForkSessionResponse, ForkSessionResponses, FrontendToolRequest, GetCanonicalModelInfoData, GetCanonicalModelInfoResponse, GetCanonicalModelInfoResponses, GetCustomProviderData, GetCustomProviderErrors, GetCustomProviderResponse, GetCustomProviderResponses, GetDictationConfigData, GetDictationConfigResponse, GetDictationConfigResponses, GetDownloadProgressData, GetDownloadProgressErrors, GetDownloadProgressResponse, GetDownloadProgressResponses, GetExtensionsData, GetExtensionsErrors, GetExtensionsResponse, GetExtensionsResponses, GetFeaturesData, GetFeaturesResponse, GetFeaturesResponses, GetLocalModelDownloadProgressData, GetLocalModelDownloadProgressErrors, GetLocalModelDownloadProgressResponse, GetLocalModelDownloadProgressResponses, GetModelSettingsData, GetModelSettingsErrors, GetModelSettingsResponse, GetModelSettingsResponses, GetPromptData, GetPromptErrors, GetPromptResponse, GetPromptResponses, GetPromptsData, GetPromptsResponse, GetPromptsResponses, GetProviderCatalogData, GetProviderCatalogErrors, GetProviderCatalogResponse, GetProviderCatalogResponses, GetProviderCatalogTemplateData, GetProviderCatalogTemplateErrors, GetProviderCatalogTemplateResponse, GetProviderCatalogTemplateResponses, GetProviderModelInfoData, GetProviderModelInfoErrors, GetProviderModelInfoResponse, GetProviderModelInfoResponses, GetProviderModelsData, GetProviderModelsErrors, GetProviderModelsResponse, GetProviderModelsResponses, GetRepoFilesData, GetRepoFilesResponse, GetRepoFilesResponses, GetSessionData, GetSessionErrors, GetSessionExtensionsData, GetSessionExtensionsErrors, GetSessionExtensionsResponse, GetSessionExtensionsResponses, GetSessionResponse, GetSessionResponses, GetSlashCommandsData, GetSlashCommandsResponse, GetSlashCommandsResponses, GetToolsData, GetToolsErrors, GetToolsQuery, GetToolsResponse, GetToolsResponses, GetTunnelStatusData, GetTunnelStatusResponse, GetTunnelStatusResponses, GooseApp, GooseMode, HfGgufFile, HfModelInfo, HfModelVariant, HfQuantVariant, Icon, IconTheme, ImageContent, ImportAppData, ImportAppError, ImportAppErrors, ImportAppRequest, ImportAppResponse, ImportAppResponse2, ImportAppResponses, ImportSessionNostrData, ImportSessionNostrErrors, ImportSessionNostrRequest, ImportSessionNostrResponse, ImportSessionNostrResponses, InferenceMetadata, InspectJobResponse, InspectRunningJobData, InspectRunningJobErrors, InspectRunningJobResponse, InspectRunningJobResponses, JsonObject, KillJobResponse, KillRunningJobData, KillRunningJobResponses, ListAppsData, ListAppsError, ListAppsErrors, ListAppsRequest, ListAppsResponse, ListAppsResponse2, ListAppsResponses, ListBuiltinChatTemplatesData, ListBuiltinChatTemplatesResponse, ListBuiltinChatTemplatesResponses, ListLocalModelsData, ListLocalModelsResponse, ListLocalModelsResponses, ListModelsData, ListModelsResponse, ListModelsResponses, ListProviderSecretsData, ListProviderSecretsErrors, ListProviderSecretsResponse, ListProviderSecretsResponses, ListRecipeResponse, ListRecipesData, ListRecipesErrors, ListRecipesResponse, ListRecipesResponses, ListSchedulesData, ListSchedulesErrors, ListSchedulesResponse, ListSchedulesResponse2, ListSchedulesResponses, LoadedProvider, LocalModelResponse, McpAppResource, McpUiProxyData, McpUiProxyErrors, McpUiProxyResponses, Message, MessageContent, MessageEvent, MessageMetadata, ModelCapabilities, ModelConfig, ModelDownloadStatus, ModelInfo, ModelInfoData, ModelInfoQuery, ModelInfoResponse, ModelSettings, ModelTemplate, ParseRecipeData, ParseRecipeError, ParseRecipeErrors, ParseRecipeRequest, ParseRecipeResponse, ParseRecipeResponse2, ParseRecipeResponses, PauseScheduleData, PauseScheduleErrors, PauseScheduleResponse, PauseScheduleResponses, Permission, PermissionLevel, PermissionsMetadata, PrincipalType, PromptContentResponse, PromptsListResponse, ProviderCatalogEntry, ProviderDetails, ProviderEngine, ProviderMetadata, ProviderModelInfoQuery, ProvidersData, ProviderSecret, ProviderSecretsResponse, ProviderSecretStatus, ProviderSecretStorage, ProvidersResponse, ProvidersResponse2, ProvidersResponses, ProviderTemplate, ProviderType, RawAudioContent, RawEmbeddedResource, RawImageContent, RawResource, RawTextContent, ReadAllConfigData, ReadAllConfigResponse, ReadAllConfigResponses, ReadConfigData, ReadConfigErrors, ReadConfigResponses, ReadResourceData, ReadResourceErrors, ReadResourceRequest, ReadResourceResponse, ReadResourceResponse2, ReadResourceResponses, Recipe, RecipeManifest, RecipeParameter, RecipeParameterInputType, RecipeParameterRequirement, RecipeToYamlData, RecipeToYamlError, RecipeToYamlErrors, RecipeToYamlRequest, RecipeToYamlResponse, RecipeToYamlResponse2, RecipeToYamlResponses, RedactedThinkingContent, RemoveConfigData, RemoveConfigErrors, RemoveConfigResponse, RemoveConfigResponses, RemoveCustomProviderData, RemoveCustomProviderErrors, RemoveCustomProviderResponse, RemoveCustomProviderResponses, RemoveExtensionData, RemoveExtensionErrors, RemoveExtensionRequest, RemoveExtensionResponse, RemoveExtensionResponses, ReplyData, ReplyErrors, ReplyResponse, ReplyResponses, RepoVariantsResponse, ResetPromptData, ResetPromptErrors, ResetPromptResponse, ResetPromptResponses, ResourceContents, ResourceMetadata, Response, RestartAgentData, RestartAgentErrors, RestartAgentRequest, RestartAgentResponse, RestartAgentResponse2, RestartAgentResponses, ResumeAgentData, ResumeAgentErrors, ResumeAgentRequest, ResumeAgentResponse, ResumeAgentResponse2, ResumeAgentResponses, RetryConfig, Role, RunNowHandlerData, RunNowHandlerErrors, RunNowHandlerResponse, RunNowHandlerResponses, RunNowResponse, SamplingConfig, SavePromptData, SavePromptErrors, SavePromptRequest, SavePromptResponse, SavePromptResponses, SaveRecipeData, SaveRecipeError, SaveRecipeErrors, SaveRecipeRequest, SaveRecipeResponse, SaveRecipeResponse2, SaveRecipeResponses, ScanRecipeData, ScanRecipeRequest, ScanRecipeResponse, ScanRecipeResponse2, ScanRecipeResponses, ScheduledJob, ScheduleRecipeData, ScheduleRecipeErrors, ScheduleRecipeRequest, ScheduleRecipeResponses, SearchHfModelsData, SearchHfModelsErrors, SearchHfModelsResponse, SearchHfModelsResponses, SendTelemetryEventData, SendTelemetryEventResponses, Session, SessionCancelData, SessionCancelResponses, SessionDisplayInfo, SessionEventsData, SessionEventsErrors, SessionEventsResponse, SessionEventsResponses, SessionExtensionsResponse, SessionReplyData, SessionReplyErrors, SessionReplyRequest, SessionReplyResponse, SessionReplyResponse2, SessionReplyResponses, SessionsHandlerData, SessionsHandlerErrors, SessionsHandlerResponse, SessionsHandlerResponses, SessionsQuery, SessionType, SetConfigProviderData, SetProviderRequest, SetRecipeSlashCommandData, SetRecipeSlashCommandErrors, SetRecipeSlashCommandResponses, SetSlashCommandRequest, Settings, SetupResponse, ShareSessionNostrData, ShareSessionNostrErrors, ShareSessionNostrRequest, ShareSessionNostrResponse, ShareSessionNostrResponse2, ShareSessionNostrResponses, SlashCommand, SlashCommandsResponse, StartAgentData, StartAgentError, StartAgentErrors, StartAgentRequest, StartAgentResponse, StartAgentResponses, StartNanogptSetupData, StartNanogptSetupResponse, StartNanogptSetupResponses, StartOpenrouterSetupData, StartOpenrouterSetupResponse, StartOpenrouterSetupResponses, StartTetrateSetupData, StartTetrateSetupResponse, StartTetrateSetupResponses, StartTunnelData, StartTunnelError, StartTunnelErrors, StartTunnelResponse, StartTunnelResponses, StatusData, StatusResponse, StatusResponses, StopAgentData, StopAgentErrors, StopAgentRequest, StopAgentResponse, StopAgentResponses, StopTunnelData, StopTunnelError, StopTunnelErrors, StopTunnelResponses, SubRecipe, SuccessCheck, SyncFeaturedModelsData, SyncFeaturedModelsResponses, SystemInfo, SystemInfoData, SystemInfoResponse, SystemInfoResponses, SystemNotificationContent, SystemNotificationType, TaskSupport, TelemetryEventRequest, Template, TextContent, ThinkingContent, ThinkingEffort, TokenState, Tool, ToolAnnotations, ToolCallingMode, ToolConfirmationRequest, ToolExecution, ToolInfo, ToolPermission, ToolRequest, ToolResponse, TranscribeDictationData, TranscribeDictationErrors, TranscribeDictationResponse, TranscribeDictationResponses, TranscribeRequest, TranscribeResponse, TunnelInfo, TunnelState, UiMetadata, UnpauseScheduleData, UnpauseScheduleErrors, UnpauseScheduleResponse, UnpauseScheduleResponses, UpdateAgentProviderData, UpdateAgentProviderErrors, UpdateAgentProviderResponses, UpdateCustomProviderData, UpdateCustomProviderErrors, UpdateCustomProviderRequest, UpdateCustomProviderResponse, UpdateCustomProviderResponses, UpdateFromSessionData, UpdateFromSessionErrors, UpdateFromSessionRequest, UpdateFromSessionResponses, UpdateModelSettingsData, UpdateModelSettingsErrors, UpdateModelSettingsResponse, UpdateModelSettingsResponses, UpdateProviderRequest, UpdateScheduleData, UpdateScheduleErrors, UpdateScheduleRequest, UpdateScheduleResponse, UpdateScheduleResponses, UpdateSessionData, UpdateSessionErrors, UpdateSessionNameData, UpdateSessionNameErrors, UpdateSessionNameRequest, UpdateSessionNameResponses, UpdateSessionRequest, UpdateSessionResponses, UpdateSessionUserRecipeValuesData, UpdateSessionUserRecipeValuesError, UpdateSessionUserRecipeValuesErrors, UpdateSessionUserRecipeValuesRequest, UpdateSessionUserRecipeValuesResponse, UpdateSessionUserRecipeValuesResponse2, UpdateSessionUserRecipeValuesResponses, UpdateWorkingDirData, UpdateWorkingDirErrors, UpdateWorkingDirRequest, UpdateWorkingDirResponses, UpsertConfigData, UpsertConfigErrors, UpsertConfigQuery, UpsertConfigResponse, UpsertConfigResponses, UpsertPermissionsData, UpsertPermissionsErrors, UpsertPermissionsQuery, UpsertPermissionsResponse, UpsertPermissionsResponses, Usage, ValidateConfigData, ValidateConfigErrors, ValidateConfigResponse, ValidateConfigResponses, WhisperModelResponse, WindowProps } from './types.gen'; diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 1aeb78722..2bc1a817c 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -1364,28 +1364,24 @@ export type ScheduledJob = { export type Session = { accumulated_cost?: number | null; - accumulated_input_tokens?: number | null; - accumulated_output_tokens?: number | null; - accumulated_total_tokens?: number | null; + accumulated_usage?: Usage; archived_at?: string | null; conversation?: Conversation | null; created_at: string; extension_data: ExtensionData; goose_mode?: GooseMode; id: string; - input_tokens?: number | null; last_message_snippet?: string | null; message_count: number; model_config?: ModelConfig | null; name: string; - output_tokens?: number | null; project_id?: string | null; provider_name?: string | null; recipe?: Recipe | null; schedule_id?: string | null; session_type?: SessionType; - total_tokens?: number | null; updated_at: string; + usage?: Usage; user_recipe_values?: { [key: string]: string; } | null; @@ -1563,10 +1559,14 @@ export type ThinkingContent = { export type ThinkingEffort = 'off' | 'low' | 'medium' | 'high' | 'max'; export type TokenState = { + accumulatedCacheReadTokens?: number; + accumulatedCacheWriteTokens?: number; accumulatedCost?: number | null; accumulatedInputTokens: number; accumulatedOutputTokens: number; accumulatedTotalTokens: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; inputTokens: number; outputTokens: number; totalTokens: number; @@ -1776,6 +1776,24 @@ export type UpsertPermissionsQuery = { tool_permissions: Array; }; +/** + * `input_tokens` is the total input including cache read/write tokens; + * the cache fields are breakdown subsets of it. Parsers for providers + * that report cache tokens separately from input (e.g. Anthropic, + * Bedrock) must fold them into `input_tokens`. + */ +export type Usage = { + cache_read_input_tokens?: number | null; + cache_write_input_tokens?: number | null; + /** + * All prompt tokens, including any served from or written to cache. + * `cache_read_input_tokens` and `cache_write_input_tokens` are subsets of this. + */ + input_tokens?: number | null; + output_tokens?: number | null; + total_tokens?: number | null; +}; + export type WhisperModelResponse = { /** * Description diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index b89cbba24..97a6f0ebe 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -523,12 +523,16 @@ export default function BaseChat({ commandHistory={commandHistory} initialValue={initialPrompt} setView={setView} - totalTokens={tokenState?.totalTokens ?? session?.total_tokens ?? undefined} + totalTokens={tokenState?.totalTokens ?? session?.usage?.total_tokens ?? undefined} accumulatedInputTokens={ - tokenState?.accumulatedInputTokens ?? session?.accumulated_input_tokens ?? undefined + tokenState?.accumulatedInputTokens ?? + session?.accumulated_usage?.input_tokens ?? + undefined } accumulatedOutputTokens={ - tokenState?.accumulatedOutputTokens ?? session?.accumulated_output_tokens ?? undefined + tokenState?.accumulatedOutputTokens ?? + session?.accumulated_usage?.output_tokens ?? + undefined } accumulatedCost={tokenState?.accumulatedCost ?? session?.accumulated_cost ?? undefined} droppedFiles={droppedFiles} diff --git a/ui/desktop/src/components/sessions/SessionHistoryView.tsx b/ui/desktop/src/components/sessions/SessionHistoryView.tsx index 562af8a18..2806f5169 100644 --- a/ui/desktop/src/components/sessions/SessionHistoryView.tsx +++ b/ui/desktop/src/components/sessions/SessionHistoryView.tsx @@ -246,7 +246,7 @@ const SessionHistoryView: React.FC = ({ session.working_dir, messages, session.name || 'Shared Session', - session.total_tokens || 0 + session.usage?.total_tokens || 0 ); const shareableLink = `goose://sessions/${shareToken}`; @@ -343,10 +343,10 @@ const SessionHistoryView: React.FC = ({ {session.message_count} - {session.total_tokens !== null && ( + {session.usage?.total_tokens != null && ( - {(session.total_tokens || 0).toLocaleString()} + {session.usage.total_tokens.toLocaleString()} )} diff --git a/ui/desktop/src/components/sessions/SessionsView.tsx b/ui/desktop/src/components/sessions/SessionsView.tsx index a7eb51360..0cdc4a017 100644 --- a/ui/desktop/src/components/sessions/SessionsView.tsx +++ b/ui/desktop/src/components/sessions/SessionsView.tsx @@ -94,7 +94,7 @@ const SessionsView: React.FC = () => { name: intl.formatMessage(i18n.loading), working_dir: '', message_count: 0, - total_tokens: 0, + usage: {}, created_at: '', updated_at: '', extension_data: {}, diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index e8002ef4f..b027a93d8 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -722,12 +722,12 @@ export function useChatStream({ session: cached.session, messages: cached.messages, tokenState: { - inputTokens: cached.session?.input_tokens ?? 0, - outputTokens: cached.session?.output_tokens ?? 0, - totalTokens: cached.session?.total_tokens ?? 0, - accumulatedInputTokens: cached.session?.accumulated_input_tokens ?? 0, - accumulatedOutputTokens: cached.session?.accumulated_output_tokens ?? 0, - accumulatedTotalTokens: cached.session?.accumulated_total_tokens ?? 0, + inputTokens: cached.session?.usage?.input_tokens ?? 0, + outputTokens: cached.session?.usage?.output_tokens ?? 0, + totalTokens: cached.session?.usage?.total_tokens ?? 0, + accumulatedInputTokens: cached.session?.accumulated_usage?.input_tokens ?? 0, + accumulatedOutputTokens: cached.session?.accumulated_usage?.output_tokens ?? 0, + accumulatedTotalTokens: cached.session?.accumulated_usage?.total_tokens ?? 0, }, }, }); @@ -778,12 +778,12 @@ export function useChatStream({ session: loadedSession!, messages: loadedSession?.conversation || [], tokenState: { - inputTokens: loadedSession?.input_tokens ?? 0, - outputTokens: loadedSession?.output_tokens ?? 0, - totalTokens: loadedSession?.total_tokens ?? 0, - accumulatedInputTokens: loadedSession?.accumulated_input_tokens ?? 0, - accumulatedOutputTokens: loadedSession?.accumulated_output_tokens ?? 0, - accumulatedTotalTokens: loadedSession?.accumulated_total_tokens ?? 0, + inputTokens: loadedSession?.usage?.input_tokens ?? 0, + outputTokens: loadedSession?.usage?.output_tokens ?? 0, + totalTokens: loadedSession?.usage?.total_tokens ?? 0, + accumulatedInputTokens: loadedSession?.accumulated_usage?.input_tokens ?? 0, + accumulatedOutputTokens: loadedSession?.accumulated_usage?.output_tokens ?? 0, + accumulatedTotalTokens: loadedSession?.accumulated_usage?.total_tokens ?? 0, }, }, }); @@ -797,12 +797,12 @@ export function useChatStream({ dispatch({ type: 'SET_TOKEN_STATE', payload: { - inputTokens: loadedSession?.input_tokens ?? 0, - outputTokens: loadedSession?.output_tokens ?? 0, - totalTokens: loadedSession?.total_tokens ?? 0, - accumulatedInputTokens: loadedSession?.accumulated_input_tokens ?? 0, - accumulatedOutputTokens: loadedSession?.accumulated_output_tokens ?? 0, - accumulatedTotalTokens: loadedSession?.accumulated_total_tokens ?? 0, + inputTokens: loadedSession?.usage?.input_tokens ?? 0, + outputTokens: loadedSession?.usage?.output_tokens ?? 0, + totalTokens: loadedSession?.usage?.total_tokens ?? 0, + accumulatedInputTokens: loadedSession?.accumulated_usage?.input_tokens ?? 0, + accumulatedOutputTokens: loadedSession?.accumulated_usage?.output_tokens ?? 0, + accumulatedTotalTokens: loadedSession?.accumulated_usage?.total_tokens ?? 0, }, }); } else { @@ -812,12 +812,12 @@ export function useChatStream({ session: loadedSession!, messages: loadedSession?.conversation || [], tokenState: { - inputTokens: loadedSession?.input_tokens ?? 0, - outputTokens: loadedSession?.output_tokens ?? 0, - totalTokens: loadedSession?.total_tokens ?? 0, - accumulatedInputTokens: loadedSession?.accumulated_input_tokens ?? 0, - accumulatedOutputTokens: loadedSession?.accumulated_output_tokens ?? 0, - accumulatedTotalTokens: loadedSession?.accumulated_total_tokens ?? 0, + inputTokens: loadedSession?.usage?.input_tokens ?? 0, + outputTokens: loadedSession?.usage?.output_tokens ?? 0, + totalTokens: loadedSession?.usage?.total_tokens ?? 0, + accumulatedInputTokens: loadedSession?.accumulated_usage?.input_tokens ?? 0, + accumulatedOutputTokens: loadedSession?.accumulated_usage?.output_tokens ?? 0, + accumulatedTotalTokens: loadedSession?.accumulated_usage?.total_tokens ?? 0, }, }, });