diff --git a/src/app.rs b/src/app.rs index 5e6d868..4b023ea 100644 --- a/src/app.rs +++ b/src/app.rs @@ -6537,6 +6537,9 @@ impl App { let agent = self.agent.clone(); let original_messages = messages; let task_session_id = session_id.to_string(); + let compaction_pricing = self.discovery.as_ref().and_then(|discovery| { + discovery.get_model_pricing(&self.provider_name.to_lowercase(), &self.model) + }); tokio::spawn(async move { let result = crate::llm::client::summarize_for_compaction( @@ -6557,7 +6560,7 @@ impl App { let mut messages = crate::session::compaction::apply_soft_compaction( &original_messages, &selection, - &summary, + &summary.text, Some(model), Some(provider_name), Some(agent), @@ -6568,6 +6571,27 @@ impl App { after_messages: 0, }, ); + let cost = compaction_pricing + .as_ref() + .map(|pricing| { + pricing.estimate_tokens( + summary.usage.input, + summary.usage.output, + summary.usage.cache_read, + summary.usage.cache_write, + ) + }) + .unwrap_or(0.0); + crate::session::compaction::attach_summary_usage( + &mut messages, + crate::session::types::RecordedUsage { + input: summary.usage.input, + output: summary.usage.output, + cache_read: summary.usage.cache_read, + cache_write: summary.usage.cache_write, + cost, + }, + ); // Count post-boundary context only (new layout: // [history][summary][tail…][marker] — marker excluded). let after_tokens = crate::session::compaction::total_context_tokens(&messages); diff --git a/src/llm/client.rs b/src/llm/client.rs index 0273806..350d36a 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -914,13 +914,57 @@ fn resolve_api_key( configured_api_key(auth_config).or(custom_provider_api_key) } +#[derive(Debug, Clone, Default, PartialEq)] +pub struct CompactionSummary { + pub text: String, + pub usage: crate::aisdk::chunk::TokenUsage, +} + +fn apply_compaction_stream_chunk( + summary: &mut String, + usage: &mut crate::aisdk::chunk::TokenUsage, + chunk: ChunkType, +) -> Result<(), DynError> { + match chunk { + ChunkType::Text(text) => summary.push_str(&text), + ChunkType::Failed(err) => { + return Err(anyhow::anyhow!("Compaction failed: {}", err).into()); + } + ChunkType::NotSupported(msg) => { + return Err(anyhow::anyhow!("Compaction unsupported: {}", msg).into()); + } + ChunkType::Usage(chunk_usage) => { + *usage = usage.saturating_add(chunk_usage); + } + ChunkType::StreamRollback { text, .. } => { + if summary.ends_with(&text) { + summary.truncate(summary.len() - text.len()); + } + } + ChunkType::Reasoning(_) + | ChunkType::ReasoningItem(_) + | ChunkType::ToolCall(_) + | ChunkType::ProviderToolCall(_) + | ChunkType::End { .. } + | ChunkType::AssistantMessagePhase { .. } + | ChunkType::ResponseCompleted { .. } + | ChunkType::Retry(_) + | ChunkType::RetryableFailure(_) + | ChunkType::Warning(_) + | ChunkType::Metadata(_) + | ChunkType::Start + | ChunkType::Incomplete(_) => {} + } + Ok(()) +} + pub async fn summarize_for_compaction( provider_name: String, model: String, reasoning_effort: Option, prompt: String, cancel_token: CancellationToken, -) -> Result { +) -> Result { if cancel_token.is_cancelled() { return Err(anyhow::anyhow!("Compaction cancelled by user").into()); } @@ -939,6 +983,7 @@ pub async fn summarize_for_compaction( .await?; let mut summary = String::new(); + let mut usage = crate::aisdk::chunk::TokenUsage::default(); loop { let chunk = tokio::select! { _ = cancel_token.cancelled() => { @@ -951,34 +996,7 @@ pub async fn summarize_for_compaction( break; }; - match chunk { - ChunkType::Text(text) => summary.push_str(&text), - ChunkType::Failed(err) => { - return Err(anyhow::anyhow!("Compaction failed: {}", err).into()); - } - ChunkType::NotSupported(msg) => { - return Err(anyhow::anyhow!("Compaction unsupported: {}", msg).into()); - } - ChunkType::Reasoning(_) - | ChunkType::ReasoningItem(_) - | ChunkType::ToolCall(_) - | ChunkType::ProviderToolCall(_) - | ChunkType::End { .. } - | ChunkType::AssistantMessagePhase { .. } - | ChunkType::ResponseCompleted { .. } - | ChunkType::Retry(_) - | ChunkType::RetryableFailure(_) - | ChunkType::Warning(_) - | ChunkType::Metadata(_) - | ChunkType::Usage(_) - | ChunkType::Start - | ChunkType::Incomplete(_) => {} - ChunkType::StreamRollback { text, .. } => { - if summary.ends_with(&text) { - summary.truncate(summary.len() - text.len()); - } - } - } + apply_compaction_stream_chunk(&mut summary, &mut usage, chunk)?; } if cancel_token.is_cancelled() { @@ -990,7 +1008,10 @@ pub async fn summarize_for_compaction( return Err(anyhow::anyhow!("Compaction returned an empty summary").into()); } - Ok(summary) + Ok(CompactionSummary { + text: summary, + usage, + }) } pub async fn generate_session_title( @@ -2602,16 +2623,71 @@ fn normalize_anthropic_base_url(base_url: &str) -> String { #[cfg(test)] mod tests { use super::{ - apply_provider_request_defaults, convert_messages, convert_messages_for_model, - is_openai_oauth_model_allowed, maybe_apply_unauthenticated_free_provider_key, - model_supports_image_input, openai_oauth_default_originator, - openai_oauth_model_uses_responses_lite, openai_request_instructions, resolve_api_key, - resolve_model_route, ui_vs_request_model_mismatch_warning, vlm_agent_has_model, - AisdkMessage, OpenAIRequestOptions, ProviderKind, ProviderRequestConfig, + apply_compaction_stream_chunk, apply_provider_request_defaults, convert_messages, + convert_messages_for_model, is_openai_oauth_model_allowed, + maybe_apply_unauthenticated_free_provider_key, model_supports_image_input, + openai_oauth_default_originator, openai_oauth_model_uses_responses_lite, + openai_request_instructions, resolve_api_key, resolve_model_route, + ui_vs_request_model_mismatch_warning, vlm_agent_has_model, AisdkMessage, + OpenAIRequestOptions, ProviderKind, ProviderRequestConfig, }; + use crate::aisdk::core::chunk::ChunkType; use crate::persistence::AuthConfig; + #[test] + fn compaction_stream_accumulates_usage_and_text() { + let mut summary = String::new(); + let mut usage = crate::aisdk::chunk::TokenUsage::default(); + + apply_compaction_stream_chunk(&mut summary, &mut usage, ChunkType::Text("hello ".into())) + .unwrap(); + apply_compaction_stream_chunk( + &mut summary, + &mut usage, + ChunkType::Usage(crate::aisdk::chunk::TokenUsage { + input: 1_000, + output: 40, + cache_read: 200, + cache_write: 10, + }), + ) + .unwrap(); + apply_compaction_stream_chunk( + &mut summary, + &mut usage, + ChunkType::Usage(crate::aisdk::chunk::TokenUsage { + input: 50, + output: 10, + cache_read: 0, + cache_write: 0, + }), + ) + .unwrap(); + apply_compaction_stream_chunk(&mut summary, &mut usage, ChunkType::Text("world".into())) + .unwrap(); + apply_compaction_stream_chunk( + &mut summary, + &mut usage, + ChunkType::StreamRollback { + text: "world".into(), + reasoning: String::new(), + }, + ) + .unwrap(); + + assert_eq!(summary, "hello "); + assert_eq!( + usage, + crate::aisdk::chunk::TokenUsage { + input: 1_050, + output: 50, + cache_read: 200, + cache_write: 10, + } + ); + } + #[test] fn stored_auth_takes_precedence_over_custom_provider_api_key() { assert_eq!( diff --git a/src/persistence/conversions.rs b/src/persistence/conversions.rs index b1f2f35..b5076eb 100644 --- a/src/persistence/conversions.rs +++ b/src/persistence/conversions.rs @@ -11,6 +11,7 @@ impl From for Message { // Move the owned parts instead of cloning them: this conversion runs // for the whole transcript on every streaming snapshot. let usage = msg.recorded_usage(); + let is_compaction_summary = crate::session::compaction::is_compaction_summary(&msg); let mut parts: Vec = if msg.parts.is_empty() { let mut parts = Vec::new(); if !msg.content.is_empty() { @@ -71,6 +72,20 @@ impl From for Message { }); } + // Compaction summaries store billed prompt/completion on a usage part + // for stats/cost. `tokens_used` is the context estimate (summary text), + // not billed buckets — otherwise reload inflates the model window. + let tokens_used = if is_compaction_summary { + msg.token_count + .map(|count| count.min(i32::MAX as usize) as i32) + .unwrap_or(0) + } else { + usage + .map(|usage| usage.tokens().min(i32::MAX as u64) as i32) + .or(msg.token_count.map(|c| c as i32)) + .unwrap_or(0) + }; + Message { id: cuid2::create_id(), session_id: 0, @@ -86,10 +101,7 @@ impl From for Message { .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_secs() as i64, - tokens_used: usage - .map(|usage| usage.tokens().min(i32::MAX as u64) as i32) - .or(msg.token_count.map(|c| c as i32)) - .unwrap_or(0), + tokens_used, model: msg.model.clone(), provider: msg.provider.clone(), agent_mode: msg.agent_mode.clone(), @@ -320,4 +332,30 @@ mod tests { .iter() .any(|part| part.part_type == "usage")); } + + #[test] + fn compaction_summary_keeps_context_token_count_not_billed_usage() { + let mut summary = SessionMessage::user(format!( + "{}\n{}", + crate::session::compaction::SUMMARY_PREFIX, + "handoff summary" + )); + summary.token_count = Some(40); + summary + .parts + .push(SessionMessagePart::usage(80_000, 400, 12_000, 1_000, 0.42)); + + let persistence_message: Message = summary.into(); + assert_eq!(persistence_message.tokens_used, 40); + assert!(persistence_message + .parts + .iter() + .any(|part| part.part_type == "usage")); + + let restored = SessionMessage::try_from(persistence_message).unwrap(); + assert_eq!(restored.token_count, Some(40)); + let usage = restored.recorded_usage().unwrap(); + assert_eq!(usage.input, 80_000); + assert_eq!(usage.output, 400); + } } diff --git a/src/session/compaction.rs b/src/session/compaction.rs index 2acd847..2927eb2 100644 --- a/src/session/compaction.rs +++ b/src/session/compaction.rs @@ -1,4 +1,4 @@ -use crate::session::types::{CompactionStats, Message, MessageRole}; +use crate::session::types::{CompactionStats, Message, MessagePart, MessageRole, RecordedUsage}; /// Max recent user turns considered for the preserved tail. /// Actual tail is also capped by [`DEFAULT_PRESERVE_RECENT_TOKENS`]. @@ -373,6 +373,31 @@ pub fn apply_soft_compaction( result } +/// Persist billed compaction tokens/cost on the synthetic summary message. +/// +/// OpenCode stores this on the `summary:true` assistant message. Crabcode's +/// summary is a user message; attaching a usage part here is what `stats` and +/// the session footer already sum. +pub fn attach_summary_usage(messages: &mut [Message], usage: RecordedUsage) { + if usage.tokens() == 0 { + return; + } + + if let Some(summary) = messages + .iter_mut() + .rev() + .find(|message| is_compaction_summary(message)) + { + summary.parts.push(MessagePart::usage( + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + usage.cost, + )); + } +} + /// Token count for the active model context (post-boundary), not full UI history. pub fn total_context_tokens(messages: &[Message]) -> usize { filter_messages_for_context(messages) @@ -386,6 +411,12 @@ pub fn message_context_tokens(message: &Message) -> usize { return 0; } + // Billed compaction usage can be persisted on the summary. Context is the + // summary text, never those billed prompt tokens. + if is_compaction_summary(message) { + return estimate_tokens(&message.content); + } + let part_tokens = message_parts_context_tokens(message); if part_tokens > 0 { return message @@ -972,4 +1003,84 @@ mod tests { assert!(!prompt.contains(" = soft + .iter() + .cloned() + .map(crate::persistence::Message::from) + .collect(); + let restored: Vec = persisted + .into_iter() + .map(|message| Message::try_from(message).expect("restore")) + .collect(); + let restored_summary = restored + .iter() + .find(|message| is_compaction_summary(message)) + .expect("restored summary"); + let restored_usage = restored_summary.recorded_usage().expect("usage part"); + assert_eq!(restored_usage.input, 80_000); + assert!((restored_usage.cost - 0.42).abs() < f64::EPSILON); + assert_eq!(total_context_tokens(&restored), before); + assert!(message_context_tokens(restored_summary) < 80_000); + } + + #[test] + fn empty_summary_usage_is_not_attached() { + let compacted = build_compacted_messages("summary", Vec::new(), None, None, None, None); + let mut compacted = compacted; + attach_summary_usage(&mut compacted, RecordedUsage::default()); + assert!(compacted[0].recorded_usage().is_none()); + } } diff --git a/src/stats.rs b/src/stats.rs index cbc6e0b..adf72b5 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -610,6 +610,38 @@ mod tests { assert_eq!(report.models[0].1.usage.output, 2_000); } + #[test] + fn compaction_summary_usage_counts_toward_totals() { + let conn = test_db(); + conn.execute( + "INSERT INTO messages + (id, session_id, role, parts, timestamp, tokens_used, output_tokens, model, provider) + VALUES + ('m4', 1, 'user', ?1, 1002, 92400, 400, 'gpt-test', 'openai')", + params![ + r#"[{"type":"text","text":"Another language model started to solve this problem"},{"type":"usage","input":80000,"output":400,"cache_read":12000,"cache_write":0,"cost":0.42}]"# + ], + ) + .unwrap(); + + let report = collect( + &conn, + &StatsOptions { + models: Some(None), + ..StatsOptions::default() + }, + 100_000, + ) + .unwrap(); + + assert_eq!(report.usage.input, 84_000); + assert_eq!(report.usage.output, 2_400); + assert_eq!(report.usage.cache_read, 15_000); + assert!((report.usage.cost - 0.545).abs() < 1e-9); + assert_eq!(report.models[0].1.messages, 3); + assert_eq!(report.models[0].1.usage.input, 84_000); + } + #[test] fn days_filter_counts_only_active_sessions_and_uses_requested_days() { let report = collect(