Skip to content

Fix context window display issues #99

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ailinal wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix context window display issues #99

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions src/config/models.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
1 change: 1 addition & 0 deletions src/config/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,7 @@ pub type Usage = RawUsage;
#[derive(Deserialize)]
pub struct Message {
pub usage: Option<Usage>,
pub stop_reason: Option<String>,
}

#[derive(Deserialize)]
Expand Down
256 changes: 89 additions & 167 deletions src/core/segments/context_window.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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<SegmentData> {
// 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());

Expand All @@ -85,59 +74,32 @@ impl Segment for ContextWindowSegment {

fn parse_transcript_usage<P: AsRef<Path>>(transcript_path: P) -> Option<u32> {
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<u32> {
let file = fs::File::open(path).ok()?;
let reader = BufReader::new(file);
let lines: Vec<String> = reader
.lines()
.collect::<Result<Vec<_>, _>>()
.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::<TranscriptEntry>(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::<TranscriptEntry>(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);
}
}
}
Expand All @@ -148,125 +110,85 @@ fn try_parse_transcript_file(path: &Path) -> Option<u32> {
}

fn find_usage_by_leaf_uuid(leaf_uuid: &str, project_dir: &Path) -> Option<u32> {
// 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<u32> {
let file = fs::File::open(path).ok()?;
let reader = BufReader::new(file);
let lines: Vec<String> = reader
.lines()
.collect::<Result<Vec<_>, _>>()
.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::<TranscriptEntry>(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<u32> {
fn find_assistant_by_uuid(lines: &[String], target_uuid: &str) -> Option<u32> {
for line in lines {
let line = line.trim();
if line.is_empty() {
continue;
}

if let Ok(entry) = serde_json::from_str::<TranscriptEntry>(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<u32> {
let project_dir = transcript_path.parent()?;

// Find the most recent session file in the project directory
let mut session_files: Vec<PathBuf> = 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<Vec<String>> {
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<TranscriptEntry> {
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<u32> {
// 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
}
}