Merge pull request #99 from 7df-lab/dev/bugfix0613

Fix Anthropic usage and context accounting

Author: wangtsiao

crates/core/src/context/mod.rs

@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 
 use crate::{ItemId, ResponseItem, SessionId, SummaryModelSelection, TurnId};
-use devo_protocol::{ContentBlock, Message, Role};
+use devo_protocol::{ContentBlock, Message, Model, Role};
 
 // ---------------------------------------------------------------------------
 // Contextual user fragment traits and registration
@@ -120,6 +120,10 @@ pub struct TokenBudget {
     pub max_output_tokens: usize,
     /// The threshold at which automatic compaction should trigger.
     pub compact_threshold: f64,
+    /// Absolute token limit for automatic compaction, when model metadata
+    /// provides a more precise context boundary than the default ratio.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub auto_compact_token_limit: Option<usize>,
 }
 
 impl TokenBudget {
@@ -129,6 +133,22 @@ impl TokenBudget {
             context_window,
             max_output_tokens,
             compact_threshold: 0.9,
+            auto_compact_token_limit: None,
+        }
+    }
+
+    /// Creates a token budget aligned with the active model's effective context window.
+    pub fn for_model(model: &Model) -> Self {
+        let default_budget = Self::default();
+        let context_window = model.effective_context_window() as usize;
+        let max_output_tokens = model
+            .max_tokens
+            .map_or(default_budget.max_output_tokens, |value| value as usize);
+        Self {
+            context_window,
+            max_output_tokens,
+            compact_threshold: 1.0,
+            auto_compact_token_limit: Some(context_window),
         }
     }
 
@@ -139,6 +159,9 @@ impl TokenBudget {
 
     /// Returns whether compaction should run for the supplied prompt token count.
     pub fn should_compact(&self, current_tokens: usize) -> bool {
+        if let Some(limit) = self.auto_compact_token_limit {
+            return current_tokens > limit;
+        }
         current_tokens as f64 > self.input_budget() as f64 * self.compact_threshold
     }
 }
@@ -383,6 +406,8 @@ mod tests {
         ByteTokenEstimator, ContextualUserFragment, PromptAssemblyInput, SnapshotPersistFailure,
         TokenBudget, TokenEstimator,
     };
+    use devo_protocol::Model;
+    use pretty_assertions::assert_eq;
     use std::hint::black_box;
     use std::time::Instant;
 
@@ -403,11 +428,52 @@ mod tests {
     #[test]
     fn token_budget_default_values() {
         let budget = TokenBudget::default();
-        assert_eq!(budget.context_window, 200_000);
-        assert_eq!(budget.max_output_tokens, 8192);
+        assert_eq!(
+            budget,
+            TokenBudget {
+                context_window: 200_000,
+                max_output_tokens: 8192,
+                compact_threshold: 0.9,
+                auto_compact_token_limit: None,
+            }
+        );
         assert!((budget.compact_threshold - 0.9).abs() < f64::EPSILON);
     }
 
+    #[test]
+    fn token_budget_for_model_uses_effective_context_as_auto_compact_limit() {
+        let model = Model {
+            context_window: 1_000_000,
+            effective_context_window_percent: Some(95),
+            max_tokens: Some(384_000),
+            ..Model::default()
+        };
+
+        assert_eq!(
+            TokenBudget::for_model(&model),
+            TokenBudget {
+                context_window: 950_000,
+                max_output_tokens: 384_000,
+                compact_threshold: 1.0,
+                auto_compact_token_limit: Some(950_000),
+            }
+        );
+    }
+
+    #[test]
+    fn model_token_budget_does_not_compact_before_effective_context_limit() {
+        let model = Model {
+            context_window: 1_000_000,
+            effective_context_window_percent: Some(95),
+            max_tokens: Some(384_000),
+            ..Model::default()
+        };
+        let budget = TokenBudget::for_model(&model);
+
+        assert_eq!(budget.should_compact(950_000), false);
+        assert_eq!(budget.should_compact(950_001), true);
+    }
+
     #[test]
     fn token_budget_input_budget_saturates() {
         let budget = TokenBudget::new(100, 200);

crates/core/src/query.rs

@@ -691,11 +691,11 @@ pub async fn query(
         }
 
         // 1.3 + 1.7: Check token budget and compact before building the request
-        if session.last_input_tokens > 0
+        if session.last_turn_tokens > 0
             && session
                 .config
                 .token_budget
-                .should_compact(session.last_input_tokens)
+                .should_compact(session.last_turn_tokens)
         {
             if !budget_steer_injected {
                 if let Some(turn) = session.turn_state.as_mut() {
@@ -982,6 +982,10 @@ pub async fn query(
                     session.total_cache_read_tokens +=
                         response.usage.cache_read_input_tokens.unwrap_or(0);
                     session.last_input_tokens = response.usage.input_tokens;
+                    session.last_turn_tokens = response
+                        .usage
+                        .input_tokens
+                        .saturating_add(response.usage.output_tokens);
 
                     emit(QueryEvent::Usage {
                         input_tokens: response.usage.input_tokens,

crates/core/src/session.rs

@@ -116,6 +116,10 @@ impl From<HashMap<String, String>> for ProviderRequestModelMap {
 }
 
 impl TurnConfig {
+    pub fn token_budget(&self) -> TokenBudget {
+        TokenBudget::for_model(&self.model)
+    }
+
     pub fn new(model: Model, thinking_selection: Option<String>) -> Self {
         let request_model = model.slug.clone();
         let thinking_selection = model.normalize_thinking_selection(thinking_selection.as_deref());
@@ -239,8 +243,10 @@ pub struct SessionState {
     pub total_cache_read_tokens: usize,     // TODO: same with `total_input_cached_tokens`.
     pub prompt_token_estimate: usize,
     /// Input tokens reported by the model for the most recent turn.
-    /// Used by `TokenBudget::should_compact()` to decide when to compact.
     pub last_input_tokens: usize,
+    /// Total context tokens reported by the model for the most recent turn.
+    /// This includes input plus output and drives automatic compaction.
+    pub last_turn_tokens: usize,
     /// Thread-safe queue for pending turn inputs.
     /// - Source: user sends `turn/start` while a turn is active.
     /// - Lifecycle: preserved across turns; unconsumed items are pushed back
@@ -273,6 +279,7 @@ impl SessionState {
             total_cache_read_tokens: 0,
             prompt_token_estimate: 0,
             last_input_tokens: 0,
+            last_turn_tokens: 0,
             pending_turn_queue: Arc::new(Mutex::new(VecDeque::new())),
             btw_input_queue: Arc::new(Mutex::new(VecDeque::new())),
             turn_state: None,
@@ -480,6 +487,29 @@ mod tests {
         assert_eq!(provider_bound.thinking_selection, Some("high".to_string()));
     }
 
+    #[test]
+    fn turn_config_token_budget_uses_model_effective_context() {
+        let model = Model {
+            slug: "deepseek-v4-pro".to_string(),
+            display_name: "deepseek-v4-pro".to_string(),
+            context_window: 1_000_000,
+            effective_context_window_percent: Some(95),
+            max_tokens: Some(384_000),
+            ..Model::default()
+        };
+        let turn_config = TurnConfig::new(model, None);
+
+        assert_eq!(
+            turn_config.token_budget(),
+            TokenBudget {
+                context_window: 950_000,
+                max_output_tokens: 384_000,
+                compact_threshold: 1.0,
+                auto_compact_token_limit: Some(950_000),
+            }
+        );
+    }
+
     #[test]
     fn session_config_default_values() {
         let config = SessionConfig::default();
@@ -505,6 +535,7 @@ mod tests {
         assert_eq!(state.turn_count, 0);
         assert_eq!(state.total_input_tokens, 0);
         assert_eq!(state.total_output_tokens, 0);
+        assert_eq!(state.last_turn_tokens, 0);
     }
 
     #[test]

crates/provider/src/anthropic/messages.rs

@@ -37,6 +37,7 @@ use serde_json::json;
 use tracing::debug;
 
 use super::AnthropicAIRole;
+use super::stream_usage::AnthropicStreamUsage;
 use crate::ModelProviderSDK;
 use crate::ProviderAdapter;
 use crate::ProviderCapabilities;
@@ -363,8 +364,7 @@ impl ModelProviderSDK for AnthropicProvider {
             .context("failed to create anthropic event source")?;
         let stream = async_stream::try_stream! {
             let mut message_id = String::new();
-            let mut input_tokens = 0usize;
-            let mut output_tokens = 0usize;
+            let mut stream_usage = AnthropicStreamUsage::default();
             let mut stop_reason: Option<StopReason> = None;
             let mut content_blocks: BTreeMap<usize, ResponseContent> = BTreeMap::new();
             let mut reasoning_blocks: BTreeMap<usize, String> = BTreeMap::new();
@@ -412,12 +412,9 @@ impl ModelProviderSDK for AnthropicProvider {
                                 {
                                     message_id = id.to_string();
                                 }
-                                if let Some(usage) = data.get("usage")
-                                    && let Some(input) =
-                                        usage.get("input_tokens").and_then(Value::as_u64)
-                                    {
-                                        input_tokens = input as usize;
-                                    }
+                                if let Some(usage) = stream_usage.update_from_message_start(&data) {
+                                    yield StreamEvent::UsageDelta(usage);
+                                }
                             }
                             "content_block_start" => {
                                 let Some(index) = data.get("index").and_then(Value::as_u64) else {
@@ -696,17 +693,8 @@ impl ModelProviderSDK for AnthropicProvider {
                                     {
                                         stop_reason = Some(parse_stop_reason(reason));
                                     }
-                                if let Some(usage) = data.get("usage") {
-                                    if let Some(output) = usage.get("output_tokens").and_then(Value::as_u64)
-                                    {
-                                        output_tokens = output as usize;
-                                    }
-                                    yield StreamEvent::UsageDelta(Usage {
-                                        input_tokens,
-                                        output_tokens,
-                                        cache_creation_input_tokens: None,
-                                        cache_read_input_tokens: None,
-                                    });
+                                if let Some(usage) = stream_usage.update_from_message_delta(&data) {
+                                    yield StreamEvent::UsageDelta(usage);
                                 }
                             }
                             "message_stop" => {
@@ -732,12 +720,7 @@ impl ModelProviderSDK for AnthropicProvider {
                                     content: dsml_healer
                                         .heal_response_content(content_blocks.into_values().collect()),
                                     stop_reason: stop_reason.clone(),
-                                    usage: Usage {
-                                        input_tokens,
-                                        output_tokens,
-                                        cache_creation_input_tokens: None,
-                                        cache_read_input_tokens: None,
-                                    },
+                                    usage: stream_usage.snapshot(),
                                     metadata: ResponseMetadata {
                                         extras: reasoning_blocks
                                             .values()
@@ -777,12 +760,7 @@ impl ModelProviderSDK for AnthropicProvider {
                 id: message_id,
                 content: dsml_healer.heal_response_content(content_blocks.into_values().collect()),
                 stop_reason,
-                usage: Usage {
-                    input_tokens,
-                    output_tokens,
-                    cache_creation_input_tokens: None,
-                    cache_read_input_tokens: None,
-                },
+                usage: stream_usage.snapshot(),
                 metadata: ResponseMetadata {
                     extras: reasoning_blocks
                         .into_values()
@@ -1302,8 +1280,13 @@ fn insert_provider_reasoning_blocks(
 }
 
 fn map_usage(usage: &AnthropicUsage) -> Usage {
+    let cache_creation_input_tokens = usage.cache_creation_input_tokens.unwrap_or(0);
+    let cache_read_input_tokens = usage.cache_read_input_tokens.unwrap_or(0);
     Usage {
-        input_tokens: usage.input_tokens,
+        input_tokens: usage
+            .input_tokens
+            .saturating_add(cache_creation_input_tokens)
+            .saturating_add(cache_read_input_tokens),
         output_tokens: usage.output_tokens,
         cache_creation_input_tokens: usage.cache_creation_input_tokens,
         cache_read_input_tokens: usage.cache_read_input_tokens,
@@ -1772,7 +1755,7 @@ mod tests {
 
         assert_eq!(response.id, "msg_123");
         assert_eq!(response.stop_reason, Some(StopReason::ToolUse));
-        assert_eq!(response.usage.input_tokens, 11);
+        assert_eq!(response.usage.input_tokens, 19);
         assert_eq!(response.usage.output_tokens, 7);
         assert_eq!(response.usage.cache_creation_input_tokens, Some(3));
         assert_eq!(response.usage.cache_read_input_tokens, Some(5));

crates/provider/src/anthropic/mod.rs

@@ -1,5 +1,6 @@
 pub mod messages;
 pub mod role;
+mod stream_usage;
 
 pub use messages::AnthropicProvider;
 pub use role::AnthropicAIRole;

crates/provider/src/anthropic/stream_usage.rs

@@ -0,0 +1,65 @@
+use devo_protocol::Usage;
+use serde_json::Value;
+
+#[derive(Debug, Default)]
+pub(super) struct AnthropicStreamUsage {
+    uncached_input_tokens: usize,
+    output_tokens: usize,
+    cache_creation_input_tokens: Option<usize>,
+    cache_read_input_tokens: Option<usize>,
+}
+
+impl AnthropicStreamUsage {
+    pub(super) fn update_from_message_start(&mut self, data: &Value) -> Option<Usage> {
+        let usage = data
+            .get("message")
+            .and_then(|message| message.get("usage"))
+            .or_else(|| data.get("usage"))?;
+        self.update_from_usage(usage)
+    }
+
+    pub(super) fn update_from_message_delta(&mut self, data: &Value) -> Option<Usage> {
+        self.update_from_usage(data.get("usage")?)
+    }
+
+    pub(super) fn snapshot(&self) -> Usage {
+        let cache_creation_input_tokens = self.cache_creation_input_tokens.unwrap_or(0);
+        let cache_read_input_tokens = self.cache_read_input_tokens.unwrap_or(0);
+        Usage {
+            input_tokens: self
+                .uncached_input_tokens
+                .saturating_add(cache_creation_input_tokens)
+                .saturating_add(cache_read_input_tokens),
+            output_tokens: self.output_tokens,
+            cache_creation_input_tokens: self.cache_creation_input_tokens,
+            cache_read_input_tokens: self.cache_read_input_tokens,
+        }
+    }
+
+    fn update_from_usage(&mut self, usage: &Value) -> Option<Usage> {
+        let mut updated = false;
+        if let Some(input_tokens) = usage.get("input_tokens").and_then(Value::as_u64) {
+            self.uncached_input_tokens = input_tokens as usize;
+            updated = true;
+        }
+        if let Some(output_tokens) = usage.get("output_tokens").and_then(Value::as_u64) {
+            self.output_tokens = output_tokens as usize;
+            updated = true;
+        }
+        if let Some(cache_creation_input_tokens) = usage
+            .get("cache_creation_input_tokens")
+            .and_then(Value::as_u64)
+        {
+            self.cache_creation_input_tokens = Some(cache_creation_input_tokens as usize);
+            updated = true;
+        }
+        if let Some(cache_read_input_tokens) =
+            usage.get("cache_read_input_tokens").and_then(Value::as_u64)
+        {
+            self.cache_read_input_tokens = Some(cache_read_input_tokens as usize);
+            updated = true;
+        }
+
+        updated.then(|| self.snapshot())
+    }
+}