use std::io::Write as _;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use tempfile::NamedTempFile;
#[derive(Debug, Default, Clone, Deserialize)]
pub(crate) struct ContextWindow {
#[serde(default)]
pub total_input_tokens: u64,
#[serde(default)]
pub context_window_size: u64,
#[serde(default)]
pub used_percentage: f64,
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub(crate) struct CompactionState {
#[serde(default)]
pub peak_input_tokens: u64,
#[serde(default)]
pub last_compaction: Option<CompactionRecord>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct CompactionRecord {
pub before: u64,
pub after: u64,
pub reclaimed_pct: f64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum CtxUsageColor {
Normal,
Red,
}
pub(crate) fn ctx_usage_color(used_fraction: f64) -> CtxUsageColor {
if used_fraction.clamp(0.0, 1.0) >= 0.5 {
CtxUsageColor::Red
} else {
CtxUsageColor::Normal
}
}
pub(crate) fn colorize_ctx_segment(text: &str, used_fraction: f64) -> String {
match ctx_usage_color(used_fraction) {
CtxUsageColor::Red => format!("\u{1b}[31m{text}\u{1b}[0m"),
CtxUsageColor::Normal => text.to_string(),
}
}
pub(crate) fn humanize_tokens(n: u64) -> String {
if n >= 1_000_000 {
format!("{}m", n / 1_000_000)
} else if n >= 1_000 {
format!("{}k", n / 1_000)
} else {
format!("{n}")
}
}
pub(crate) fn update_state(state: &mut CompactionState, cur: u64) -> bool {
if cur == 0 {
return false;
}
let peak = state.peak_input_tokens;
if peak > 0 {
let drop = peak.saturating_sub(cur);
let below_ratio = (cur as f64) < (peak as f64) * 0.6;
if below_ratio && drop > 10_000 {
let reclaimed_pct = 100.0 * drop as f64 / peak as f64;
state.last_compaction = Some(CompactionRecord {
before: peak,
after: cur,
reclaimed_pct,
});
state.peak_input_tokens = cur;
return true;
}
}
if cur > state.peak_input_tokens {
state.peak_input_tokens = cur;
return true;
}
false
}
fn is_valid_session_id(s: &str) -> bool {
!s.is_empty()
&& s.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
}
pub(crate) fn load_state_from(path: &Path) -> CompactionState {
(|| -> Option<CompactionState> {
let bytes = std::fs::read(path).ok()?;
serde_json::from_slice(&bytes).ok()
})()
.unwrap_or_default()
}
pub(crate) fn save_state_to(path: &Path, state: &CompactionState) {
let _ = (|| -> Option<()> {
let dir = path.parent()?;
std::fs::create_dir_all(dir).ok()?;
let bytes = serde_json::to_vec(state).ok()?;
let mut tmp = NamedTempFile::new_in(dir).ok()?;
tmp.write_all(&bytes).ok()?;
let _ = tmp.persist(path);
Some(())
})();
}
pub(crate) fn state_file_path(session_id: &str) -> Option<PathBuf> {
let home = dirs::home_dir()?;
Some(
home.join(".trusty-mpm")
.join("statusline")
.join(format!("{session_id}.json")),
)
}
pub(crate) fn compaction_segment(
session_id: &str,
context_window: Option<&ContextWindow>,
) -> Option<String> {
if !is_valid_session_id(session_id) {
return None;
}
let cw = context_window?;
let path = state_file_path(session_id)?;
let cur = cw.total_input_tokens;
let cw2 = cw.clone();
let (tx, rx) = mpsc::channel::<Option<String>>();
std::thread::spawn(move || {
let mut state = load_state_from(&path);
if cur > 0 {
if update_state(&mut state, cur) {
save_state_to(&path, &state);
}
}
let _ = tx.send(render_compaction_segment(&state, &cw2));
});
rx.recv_timeout(Duration::from_millis(100)).ok().flatten()
}
pub(crate) fn render_compaction_segment(
state: &CompactionState,
cw: &ContextWindow,
) -> Option<String> {
if let Some(ref rec) = state.last_compaction {
let before = humanize_tokens(rec.before);
let after = humanize_tokens(rec.after);
let pct = rec.reclaimed_pct.round() as u64;
return Some(format!("\u{2935} {before}\u{2192}{after} ({pct}%)"));
}
if cw.total_input_tokens > 0 && cw.context_window_size > 0 {
let fraction = cw.total_input_tokens as f64 / cw.context_window_size as f64;
let pct = (fraction * 100.0).round() as u64;
return Some(colorize_ctx_segment(&format!("ctx {pct}%"), fraction));
}
if cw.used_percentage > 0.0 {
let fraction = cw.used_percentage / 100.0;
let pct = cw.used_percentage.round() as u64;
return Some(colorize_ctx_segment(&format!("ctx {pct}%"), fraction));
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn humanize_tokens_examples() {
assert_eq!(humanize_tokens(182_000), "182k");
assert_eq!(humanize_tokens(58_000), "58k");
assert_eq!(humanize_tokens(1_200_000), "1m");
assert_eq!(humanize_tokens(999), "999");
assert_eq!(humanize_tokens(0), "0");
}
#[test]
fn humanize_tokens_boundaries() {
assert_eq!(humanize_tokens(1_000), "1k");
assert_eq!(humanize_tokens(1_999), "1k"); assert_eq!(humanize_tokens(2_000), "2k");
assert_eq!(humanize_tokens(999_999), "999k");
assert_eq!(humanize_tokens(1_000_000), "1m");
assert_eq!(humanize_tokens(1_500_000), "1m"); assert_eq!(humanize_tokens(2_000_000), "2m");
}
#[test]
fn compaction_detection_sequence() {
let mut state = CompactionState::default();
assert!(
update_state(&mut state, 50_000),
"peak advance must return true"
);
assert_eq!(state.peak_input_tokens, 50_000);
assert!(state.last_compaction.is_none());
assert!(update_state(&mut state, 100_000));
assert_eq!(state.peak_input_tokens, 100_000);
assert!(state.last_compaction.is_none());
assert!(update_state(&mut state, 182_000));
assert_eq!(state.peak_input_tokens, 182_000);
assert!(state.last_compaction.is_none());
assert!(
update_state(&mut state, 58_000),
"compaction must return true"
);
assert!(state.last_compaction.is_some());
let rec = state.last_compaction.as_ref().unwrap();
assert_eq!(rec.before, 182_000);
assert_eq!(rec.after, 58_000);
assert!(
(rec.reclaimed_pct - 68.13).abs() < 0.1,
"reclaimed_pct = {}",
rec.reclaimed_pct
);
assert_eq!(state.peak_input_tokens, 58_000);
assert!(update_state(&mut state, 60_000));
assert_eq!(state.peak_input_tokens, 60_000);
assert!(state.last_compaction.is_some());
}
#[test]
fn small_drop_not_detected_as_compaction() {
let mut state = CompactionState::default();
update_state(&mut state, 100_000);
let changed = update_state(&mut state, 95_000);
assert!(
!changed,
"sub-threshold drop must not be detected as compaction"
);
assert!(state.last_compaction.is_none());
}
#[test]
fn moderate_drop_ratio_above_threshold_not_detected() {
let mut state = CompactionState::default();
update_state(&mut state, 100_000);
let changed = update_state(&mut state, 70_000);
assert!(!changed, "above-ratio drop must not be detected");
assert!(state.last_compaction.is_none());
}
#[test]
fn zero_cur_returns_false() {
let mut state = CompactionState::default();
update_state(&mut state, 100_000);
let changed = update_state(&mut state, 0); assert!(!changed, "cur=0 must return false (no save needed)");
assert_eq!(state.peak_input_tokens, 100_000, "peak must be unchanged");
assert!(state.last_compaction.is_none());
}
#[test]
fn unchanged_state_returns_false() {
let mut state = CompactionState::default();
update_state(&mut state, 100_000);
let changed = update_state(&mut state, 100_000);
assert!(!changed, "cur == peak must return false");
let changed2 = update_state(&mut state, 95_000);
assert!(!changed2, "small drop must return false");
}
#[test]
fn valid_session_id_accepted() {
assert!(is_valid_session_id("abc-123_XYZ"));
assert!(is_valid_session_id("a"));
assert!(is_valid_session_id("session-id-with-dashes"));
assert!(is_valid_session_id("ABC0123456789"));
}
#[test]
fn rejects_path_traversal_session_id() {
for bad in &[
"../../evil",
"../evil",
"/tmp/evil",
"evil/subdir",
"evil.json",
"",
"has space",
"has\nnewline",
] {
assert!(
!is_valid_session_id(bad),
"expected validation failure for: {bad:?}"
);
let cw = ContextWindow {
total_input_tokens: 100_000,
context_window_size: 200_000,
..Default::default()
};
assert!(
compaction_segment(bad, Some(&cw)).is_none(),
"compaction_segment must return None for bad id: {bad:?}"
);
}
}
#[test]
fn state_round_trip() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("state.json");
let loaded = load_state_from(&path);
assert_eq!(loaded.peak_input_tokens, 0);
assert!(loaded.last_compaction.is_none());
let state = CompactionState {
peak_input_tokens: 182_000,
last_compaction: Some(CompactionRecord {
before: 182_000,
after: 58_000,
reclaimed_pct: 68.13,
}),
};
save_state_to(&path, &state);
let loaded2 = load_state_from(&path);
assert_eq!(loaded2.peak_input_tokens, 182_000);
let rec = loaded2.last_compaction.unwrap();
assert_eq!(rec.before, 182_000);
assert_eq!(rec.after, 58_000);
assert!((rec.reclaimed_pct - 68.13).abs() < 0.001);
let nested = dir.path().join("a").join("b").join("state.json");
let state2 = CompactionState {
peak_input_tokens: 42_000,
..Default::default()
};
save_state_to(&nested, &state2);
let loaded3 = load_state_from(&nested);
assert_eq!(loaded3.peak_input_tokens, 42_000);
}
#[test]
fn render_segment_after_compaction() {
let state = CompactionState {
peak_input_tokens: 58_000,
last_compaction: Some(CompactionRecord {
before: 182_000,
after: 58_000,
reclaimed_pct: 68.13,
}),
};
let cw = ContextWindow {
total_input_tokens: 60_000,
context_window_size: 200_000,
used_percentage: 30.0,
};
let seg = render_compaction_segment(&state, &cw).unwrap();
assert!(seg.contains("182k"), "got: {seg}");
assert!(seg.contains("58k"), "got: {seg}");
assert!(seg.contains("68%"), "got: {seg}");
assert!(seg.contains('\u{2935}'), "must contain ⤵: {seg}");
}
#[test]
fn render_segment_live_fill_from_tokens() {
let state = CompactionState::default();
let cw = ContextWindow {
total_input_tokens: 82_000,
context_window_size: 200_000,
used_percentage: 20.0, };
let seg = render_compaction_segment(&state, &cw).unwrap();
assert_eq!(seg, "ctx 41%", "got: {seg}");
}
#[test]
fn render_segment_live_fill_fallback_to_percentage() {
let state = CompactionState::default();
let cw = ContextWindow {
total_input_tokens: 0,
context_window_size: 0,
used_percentage: 41.0,
};
let seg = render_compaction_segment(&state, &cw).unwrap();
assert_eq!(seg, "ctx 41%");
}
#[test]
fn render_segment_none_when_no_info() {
let state = CompactionState::default();
let cw = ContextWindow::default(); assert!(render_compaction_segment(&state, &cw).is_none());
}
#[test]
fn ctx_usage_color_boundaries() {
assert_eq!(ctx_usage_color(0.0), CtxUsageColor::Normal);
assert_eq!(ctx_usage_color(0.49), CtxUsageColor::Normal);
assert_eq!(ctx_usage_color(0.4999), CtxUsageColor::Normal);
assert_eq!(ctx_usage_color(0.5), CtxUsageColor::Red);
assert_eq!(ctx_usage_color(0.51), CtxUsageColor::Red);
assert_eq!(ctx_usage_color(1.0), CtxUsageColor::Red);
}
#[test]
fn ctx_usage_color_clamps_out_of_range() {
assert_eq!(ctx_usage_color(-0.3), CtxUsageColor::Normal);
assert_eq!(ctx_usage_color(-100.0), CtxUsageColor::Normal);
assert_eq!(ctx_usage_color(1.5), CtxUsageColor::Red);
assert_eq!(ctx_usage_color(1_000.0), CtxUsageColor::Red);
}
#[test]
fn render_ctx_segment_colors_high_usage_red() {
let seg = colorize_ctx_segment("ctx 62%", 0.62);
assert!(
seg.contains('\u{1b}'),
"expected an ANSI escape in high-usage segment: {seg:?}"
);
assert!(
seg.contains("ctx 62%"),
"original text must survive the ANSI wrapping: {seg:?}"
);
}
#[test]
fn render_ctx_segment_leaves_low_usage_plain() {
let seg = colorize_ctx_segment("ctx 30%", 0.30);
assert_eq!(seg, "ctx 30%", "low usage must not be colorized");
assert!(!seg.contains('\u{1b}'), "must contain no ANSI escape");
}
#[test]
fn render_compaction_segment_colors_high_live_fill_red() {
let state = CompactionState::default();
let cw = ContextWindow {
total_input_tokens: 164_000,
context_window_size: 200_000,
used_percentage: 0.0,
};
let seg = render_compaction_segment(&state, &cw).unwrap();
assert!(seg.contains("ctx 82%"), "got: {seg}");
assert!(seg.contains('\u{1b}'), "expected ANSI escape: {seg:?}");
}
#[test]
fn graceful_on_missing_fields() {
assert!(compaction_segment("", None).is_none());
assert!(compaction_segment("../bad", None).is_none());
assert!(compaction_segment("valid-session-id", None).is_none());
}
}