From fe254aec27420c9532dc8603a7e3668328fa72da Mon Sep 17 00:00:00 2001 From: Blankeos Date: Thu, 3 Sep 2026 20:09:46 +0800 Subject: [PATCH 1/2] feat(cost): persist billed compaction usage on summaries Capture provider TokenUsage from the compaction stream, price it with the models.dev catalog, and store it on the summary message so stats and the session footer include compaction cost without inflating context tokens. --- src/app.rs | 26 ++++++- src/llm/client.rs | 148 ++++++++++++++++++++++++++++---------- src/session/compaction.rs | 88 ++++++++++++++++++++++- src/stats.rs | 32 +++++++++ 4 files changed, 256 insertions(+), 38 deletions(-) diff --git a/src/app.rs b/src/app.rs index 5e6d8689..4b023eaf 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 02738064..350d36a1 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/session/compaction.rs b/src/session/compaction.rs index 2acd8474..4a9c40b4 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) @@ -972,4 +997,65 @@ mod tests { assert!(!prompt.contains(" Date: Thu, 3 Sep 2026 20:22:07 +0800 Subject: [PATCH 2/2] fix(cost): keep billed compaction usage out of context tokens Persisting usage parts was overwriting tokens_used with billed prompt buckets, so reloaded summaries inflated the model window. Store the summary text estimate instead and count context from content. --- src/persistence/conversions.rs | 46 +++++++++++++++++++++++++++++++--- src/session/compaction.rs | 25 ++++++++++++++++++ 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/src/persistence/conversions.rs b/src/persistence/conversions.rs index b1f2f353..b5076ebe 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 4a9c40b4..2927eb28 100644 --- a/src/session/compaction.rs +++ b/src/session/compaction.rs @@ -411,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 @@ -1049,6 +1055,25 @@ mod tests { assert!((usage.cost - 0.42).abs() < f64::EPSILON); assert_eq!(total_context_tokens(&soft), before); assert!(message_context_tokens(summary) < 80_000); + + let persisted: Vec = 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]