From 9d8e3374c7488b3301a0d3ed0ef7957e63f11e90 Mon Sep 17 00:00:00 2001 From: yale <997942451@qq.com> Date: Fri, 20 Mar 2026 14:09:18 +0800 Subject: [PATCH] fix(context-window): accurate token display with stop_reason check Fix multiple issues with context window display: 1. **Incorrect token calculation**: Only messages with `stop_reason` have complete usage data. Streaming intermediate states have incomplete/incorrect usage values. 2. **Wrong aggregation strategy**: Take the latest complete response instead of the maximum value, because context can decrease (e.g., after compact). 3. **New session showing cached data**: Removed project history fallback to prevent new sessions from displaying data from other sessions. 4. **Missing model configs**: Added `glm-5` (200k) and `kimi-k2.5` (256k). Changes: - Add `stop_reason` field to Message struct - Extract helper functions: `read_lines`, `parse_entry`, `extract_usage` - Remove `try_find_usage_from_project_history` function - Show "0 tokens" for new sessions instead of "-" - Refactored code reduced by 67 lines Co-Authored-By: Claude Opus 4.6 --- src/config/models.rs | 10 ++ src/config/types.rs | 1 + src/core/segments/context_window.rs | 256 ++++++++++------------------ 3 files changed, 100 insertions(+), 167 deletions(-) diff --git a/src/config/models.rs b/src/config/models.rs index 6fa33273..6e552feb 100644 --- a/src/config/models.rs +++ b/src/config/models.rs @@ -298,11 +298,21 @@ impl Default for ModelConfig { // Only third-party models need explicit entries. // Claude models (Sonnet, Opus, Haiku) are handled by built-in regex families. model_entries: vec![ + ModelEntry { + pattern: "glm-5".to_string(), + display_name: "GLM-5".to_string(), + context_limit: 200_000, + }, ModelEntry { pattern: "glm-4.5".to_string(), display_name: "GLM-4.5".to_string(), context_limit: 128_000, }, + ModelEntry { + pattern: "kimi-k2.5".to_string(), + display_name: "Kimi K2.5".to_string(), + context_limit: 256_000, + }, ModelEntry { pattern: "kimi-k2-turbo".to_string(), display_name: "Kimi K2 Turbo".to_string(), diff --git a/src/config/types.rs b/src/config/types.rs index e5a78dc1..e2913acf 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -401,6 +401,7 @@ pub type Usage = RawUsage; #[derive(Deserialize)] pub struct Message { pub usage: Option, + pub stop_reason: Option, } #[derive(Deserialize)] diff --git a/src/core/segments/context_window.rs b/src/core/segments/context_window.rs index d0c7eece..de76f1fb 100644 --- a/src/core/segments/context_window.rs +++ b/src/core/segments/context_window.rs @@ -1,9 +1,9 @@ use super::{Segment, SegmentData}; -use crate::config::{InputData, ModelConfig, SegmentId, TranscriptEntry}; +use crate::config::{InputData, Message, ModelConfig, SegmentId, TranscriptEntry}; use std::collections::HashMap; use std::fs; use std::io::{BufRead, BufReader}; -use std::path::{Path, PathBuf}; +use std::path::Path; #[derive(Default)] pub struct ContextWindowSegment; @@ -13,61 +13,50 @@ impl ContextWindowSegment { Self } - /// Get context limit for the specified model fn get_context_limit_for_model(model_id: &str) -> u32 { - let model_config = ModelConfig::load(); - model_config.get_context_limit(model_id) + ModelConfig::load().get_context_limit(model_id) } } impl Segment for ContextWindowSegment { fn collect(&self, input: &InputData) -> Option { - // Dynamically determine context limit based on current model ID let context_limit = Self::get_context_limit_for_model(&input.model.id); + let tokens = parse_transcript_usage(&input.transcript_path); - let context_used_token_opt = parse_transcript_usage(&input.transcript_path); - - let (percentage_display, tokens_display) = match context_used_token_opt { - Some(context_used_token) => { - let context_used_rate = (context_used_token as f64 / context_limit as f64) * 100.0; - - let percentage = if context_used_rate.fract() == 0.0 { - format!("{:.0}%", context_used_rate) + let (percentage_display, tokens_display) = match tokens { + Some(t) => { + let rate = (t as f64 / context_limit as f64) * 100.0; + let pct = if rate.fract() == 0.0 { + format!("{:.0}%", rate) } else { - format!("{:.1}%", context_used_rate) + format!("{:.1}%", rate) }; - - let tokens = if context_used_token >= 1000 { - let k_value = context_used_token as f64 / 1000.0; - if k_value.fract() == 0.0 { - format!("{}k", k_value as u32) + let tokens_str = if t >= 1000 { + let k = t as f64 / 1000.0; + if k.fract() == 0.0 { + format!("{}k", k as u32) } else { - format!("{:.1}k", k_value) + format!("{:.1}k", k) } } else { - context_used_token.to_string() + t.to_string() }; - - (percentage, tokens) - } - None => { - // No usage data available - ("-".to_string(), "-".to_string()) + (pct, tokens_str) } + None => ("0%".to_string(), "0".to_string()), }; let mut metadata = HashMap::new(); - match context_used_token_opt { - Some(context_used_token) => { - let context_used_rate = (context_used_token as f64 / context_limit as f64) * 100.0; - metadata.insert("tokens".to_string(), context_used_token.to_string()); - metadata.insert("percentage".to_string(), context_used_rate.to_string()); - } - None => { - metadata.insert("tokens".to_string(), "-".to_string()); - metadata.insert("percentage".to_string(), "-".to_string()); - } - } + metadata.insert( + "tokens".to_string(), + tokens.map_or("0".to_string(), |t| t.to_string()), + ); + metadata.insert( + "percentage".to_string(), + tokens.map_or("0".to_string(), |t| { + ((t as f64 / context_limit as f64) * 100.0).to_string() + }), + ); metadata.insert("limit".to_string(), context_limit.to_string()); metadata.insert("model".to_string(), input.model.id.clone()); @@ -85,59 +74,32 @@ impl Segment for ContextWindowSegment { fn parse_transcript_usage>(transcript_path: P) -> Option { let path = transcript_path.as_ref(); - - // Try to parse from current transcript file - if let Some(usage) = try_parse_transcript_file(path) { - return Some(usage); - } - - // If file doesn't exist, try to find usage from project history - if !path.exists() { - if let Some(usage) = try_find_usage_from_project_history(path) { - return Some(usage); - } + if path.exists() { + try_parse_transcript_file(path) + } else { + None } - - None } fn try_parse_transcript_file(path: &Path) -> Option { - let file = fs::File::open(path).ok()?; - let reader = BufReader::new(file); - let lines: Vec = reader - .lines() - .collect::, _>>() - .unwrap_or_default(); + let lines = read_lines(path)?; - if lines.is_empty() { - return None; - } - - // Check if the last line is a summary - let last_line = lines.last()?.trim(); - if let Ok(entry) = serde_json::from_str::(last_line) { + // Check if last line is a summary + if let Some(entry) = parse_entry(lines.last()?) { if entry.r#type.as_deref() == Some("summary") { - // Handle summary case: find usage by leafUuid if let Some(leaf_uuid) = &entry.leaf_uuid { - let project_dir = path.parent()?; - return find_usage_by_leaf_uuid(leaf_uuid, project_dir); + return find_usage_by_leaf_uuid(leaf_uuid, path.parent()?); } } } - // Normal case: find the last assistant message in current file + // Find last assistant message with stop_reason (complete response) for line in lines.iter().rev() { - let line = line.trim(); - if line.is_empty() { - continue; - } - - if let Ok(entry) = serde_json::from_str::(line) { + if let Some(entry) = parse_entry(line) { if entry.r#type.as_deref() == Some("assistant") { if let Some(message) = &entry.message { - if let Some(raw_usage) = &message.usage { - let normalized = raw_usage.clone().normalize(); - return Some(normalized.display_tokens()); + if let Some(tokens) = extract_usage(message) { + return Some(tokens); } } } @@ -148,125 +110,85 @@ fn try_parse_transcript_file(path: &Path) -> Option { } fn find_usage_by_leaf_uuid(leaf_uuid: &str, project_dir: &Path) -> Option { - // Search for the leafUuid across all session files in the project directory - let entries = fs::read_dir(project_dir).ok()?; - - for entry in entries { - let entry = entry.ok()?; - let path = entry.path(); - - if path.extension().and_then(|s| s.to_str()) != Some("jsonl") { - continue; - } - - if let Some(usage) = search_uuid_in_file(&path, leaf_uuid) { - return Some(usage); + for entry in fs::read_dir(project_dir).ok()? { + let path = entry.ok()?.path(); + if path.extension().and_then(|s| s.to_str()) == Some("jsonl") { + if let Some(usage) = search_uuid_in_file(&path, leaf_uuid) { + return Some(usage); + } } } - None } fn search_uuid_in_file(path: &Path, target_uuid: &str) -> Option { - let file = fs::File::open(path).ok()?; - let reader = BufReader::new(file); - let lines: Vec = reader - .lines() - .collect::, _>>() - .unwrap_or_default(); + let lines = read_lines(path)?; - // Find the message with target_uuid for line in &lines { - let line = line.trim(); - if line.is_empty() { - continue; - } - - if let Ok(entry) = serde_json::from_str::(line) { - if let Some(uuid) = &entry.uuid { - if uuid == target_uuid { - // Found the target message, check its type - if entry.r#type.as_deref() == Some("assistant") { - // Direct assistant message with usage - if let Some(message) = &entry.message { - if let Some(raw_usage) = &message.usage { - let normalized = raw_usage.clone().normalize(); - return Some(normalized.display_tokens()); - } - } - } else if entry.r#type.as_deref() == Some("user") { - // User message, need to find the parent assistant message - if let Some(parent_uuid) = &entry.parent_uuid { - return find_assistant_message_by_uuid(&lines, parent_uuid); - } + if let Some(entry) = parse_entry(line) { + if entry.uuid.as_deref() == Some(target_uuid) { + if let Some(message) = &entry.message { + if let Some(tokens) = extract_usage(message) { + return Some(tokens); } - break; } + // If user message, find parent assistant + if entry.r#type.as_deref() == Some("user") { + if let Some(parent_uuid) = &entry.parent_uuid { + return find_assistant_by_uuid(&lines, parent_uuid); + } + } + break; } } } - None } -fn find_assistant_message_by_uuid(lines: &[String], target_uuid: &str) -> Option { +fn find_assistant_by_uuid(lines: &[String], target_uuid: &str) -> Option { for line in lines { - let line = line.trim(); - if line.is_empty() { - continue; - } - - if let Ok(entry) = serde_json::from_str::(line) { - if let Some(uuid) = &entry.uuid { - if uuid == target_uuid && entry.r#type.as_deref() == Some("assistant") { - if let Some(message) = &entry.message { - if let Some(raw_usage) = &message.usage { - let normalized = raw_usage.clone().normalize(); - return Some(normalized.display_tokens()); - } + if let Some(entry) = parse_entry(line) { + if entry.uuid.as_deref() == Some(target_uuid) + && entry.r#type.as_deref() == Some("assistant") + { + if let Some(message) = &entry.message { + if let Some(tokens) = extract_usage(message) { + return Some(tokens); } } } } } - None } -fn try_find_usage_from_project_history(transcript_path: &Path) -> Option { - let project_dir = transcript_path.parent()?; - - // Find the most recent session file in the project directory - let mut session_files: Vec = Vec::new(); - let entries = fs::read_dir(project_dir).ok()?; +// Helper functions - for entry in entries { - let entry = entry.ok()?; - let path = entry.path(); - - if path.extension().and_then(|s| s.to_str()) == Some("jsonl") { - session_files.push(path); - } - } +fn read_lines(path: &Path) -> Option> { + let file = fs::File::open(path).ok()?; + let reader = BufReader::new(file); + Some( + reader + .lines() + .filter_map(|l| l.ok()) + .collect(), + ) +} - if session_files.is_empty() { +fn parse_entry(line: &str) -> Option { + let line = line.trim(); + if line.is_empty() { return None; } + serde_json::from_str(line).ok() +} - // Sort by modification time (most recent first) - session_files.sort_by_key(|path| { - fs::metadata(path) - .and_then(|m| m.modified()) - .unwrap_or(std::time::UNIX_EPOCH) - }); - session_files.reverse(); - - // Try to find usage from the most recent session - for session_path in &session_files { - if let Some(usage) = try_parse_transcript_file(session_path) { - return Some(usage); +fn extract_usage(message: &Message) -> Option { + // Only messages with stop_reason have complete usage data + if message.stop_reason.is_some() { + if let Some(usage) = &message.usage { + return Some(usage.clone().normalize().display_tokens()); } } - None -} +} \ No newline at end of file