use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use super::context_field::{
ContextItemId, ContextKind, ContextState, Provenance, ViewCosts, ViewKind,
};
const DEFAULT_CONTEXT_WINDOW: usize = 128_000;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextLedger {
pub window_size: usize,
pub entries: Vec<LedgerEntry>,
pub total_tokens_sent: usize,
pub total_tokens_saved: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LedgerEntry {
pub path: String,
pub mode: String,
pub original_tokens: usize,
pub sent_tokens: usize,
pub timestamp: i64,
#[serde(default)]
pub id: Option<ContextItemId>,
#[serde(default)]
pub kind: Option<ContextKind>,
#[serde(default)]
pub source_hash: Option<String>,
#[serde(default)]
pub state: Option<ContextState>,
#[serde(default)]
pub phi: Option<f64>,
#[serde(default)]
pub view_costs: Option<ViewCosts>,
#[serde(default)]
pub active_view: Option<ViewKind>,
#[serde(default)]
pub provenance: Option<Provenance>,
}
#[derive(Debug, Clone)]
pub struct ContextPressure {
pub utilization: f64,
pub remaining_tokens: usize,
pub entries_count: usize,
pub recommendation: PressureAction,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PressureAction {
NoAction,
SuggestCompression,
ForceCompression,
EvictLeastRelevant,
}
impl ContextLedger {
pub fn new() -> Self {
Self {
window_size: DEFAULT_CONTEXT_WINDOW,
entries: Vec::new(),
total_tokens_sent: 0,
total_tokens_saved: 0,
}
}
pub fn with_window_size(size: usize) -> Self {
Self {
window_size: size,
entries: Vec::new(),
total_tokens_sent: 0,
total_tokens_saved: 0,
}
}
pub fn record(&mut self, path: &str, mode: &str, original_tokens: usize, sent_tokens: usize) {
let item_id = ContextItemId::from_file(path);
if let Some(existing) = self.entries.iter_mut().find(|e| e.path == path) {
self.total_tokens_sent -= existing.sent_tokens;
self.total_tokens_saved -= existing
.original_tokens
.saturating_sub(existing.sent_tokens);
existing.mode = mode.to_string();
existing.original_tokens = original_tokens;
existing.sent_tokens = sent_tokens;
existing.timestamp = chrono::Utc::now().timestamp();
existing.active_view = Some(ViewKind::parse(mode));
if existing.id.is_none() {
existing.id = Some(item_id);
}
if existing.state.is_none() || existing.state == Some(ContextState::Candidate) {
existing.state = Some(ContextState::Included);
}
} else {
self.entries.push(LedgerEntry {
path: path.to_string(),
mode: mode.to_string(),
original_tokens,
sent_tokens,
timestamp: chrono::Utc::now().timestamp(),
id: Some(item_id),
kind: Some(ContextKind::File),
source_hash: None,
state: Some(ContextState::Included),
phi: None,
view_costs: Some(ViewCosts::from_full_tokens(original_tokens)),
active_view: Some(ViewKind::parse(mode)),
provenance: None,
});
}
self.total_tokens_sent += sent_tokens;
self.total_tokens_saved += original_tokens.saturating_sub(sent_tokens);
}
pub fn upsert(
&mut self,
path: &str,
mode: &str,
original_tokens: usize,
sent_tokens: usize,
source_hash: Option<&str>,
kind: ContextKind,
provenance: Option<Provenance>,
) {
self.record(path, mode, original_tokens, sent_tokens);
if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
entry.kind = Some(kind);
if let Some(h) = source_hash {
if entry.source_hash.as_deref() != Some(h) {
if entry.source_hash.is_some() {
entry.state = Some(ContextState::Stale);
}
entry.source_hash = Some(h.to_string());
}
}
if let Some(prov) = provenance {
entry.provenance = Some(prov);
}
}
}
pub fn update_phi(&mut self, path: &str, phi: f64) {
if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
entry.phi = Some(phi);
}
}
pub fn set_state(&mut self, path: &str, state: ContextState) {
if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
entry.state = Some(state);
}
}
pub fn find_by_id(&self, id: &ContextItemId) -> Option<&LedgerEntry> {
self.entries.iter().find(|e| e.id.as_ref() == Some(id))
}
pub fn items_by_state(&self, state: ContextState) -> Vec<&LedgerEntry> {
self.entries
.iter()
.filter(|e| e.state == Some(state))
.collect()
}
pub fn eviction_candidates_by_phi(&self, keep_count: usize) -> Vec<String> {
if self.entries.len() <= keep_count {
return Vec::new();
}
let mut sorted = self.entries.clone();
sorted.sort_by(|a, b| {
let a_phi = a.phi.unwrap_or(0.0);
let b_phi = b.phi.unwrap_or(0.0);
a_phi
.partial_cmp(&b_phi)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.timestamp.cmp(&b.timestamp))
});
sorted
.iter()
.filter(|e| e.state != Some(ContextState::Pinned))
.take(self.entries.len() - keep_count)
.map(|e| e.path.clone())
.collect()
}
pub fn mark_stale_by_hash(&mut self, path: &str, new_hash: &str) {
if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
if let Some(ref old_hash) = entry.source_hash {
if old_hash != new_hash {
entry.state = Some(ContextState::Stale);
entry.source_hash = Some(new_hash.to_string());
}
}
}
}
pub fn pressure(&self) -> ContextPressure {
let utilization = self.total_tokens_sent as f64 / self.window_size as f64;
let pinned_count = self
.entries
.iter()
.filter(|e| e.state == Some(ContextState::Pinned))
.count();
let stale_count = self
.entries
.iter()
.filter(|e| e.state == Some(ContextState::Stale))
.count();
let pinned_pressure = pinned_count as f64 * 0.02;
let stale_penalty = stale_count as f64 * 0.01;
let effective_utilization = (utilization + pinned_pressure + stale_penalty).min(1.0);
let effective_used = (effective_utilization * self.window_size as f64).round() as usize;
let remaining = self.window_size.saturating_sub(effective_used);
let recommendation = if effective_utilization > 0.9 {
PressureAction::EvictLeastRelevant
} else if effective_utilization > 0.75 {
PressureAction::ForceCompression
} else if effective_utilization > 0.5 {
PressureAction::SuggestCompression
} else {
PressureAction::NoAction
};
ContextPressure {
utilization: effective_utilization,
remaining_tokens: remaining,
entries_count: self.entries.len(),
recommendation,
}
}
pub fn compression_ratio(&self) -> f64 {
let total_original: usize = self.entries.iter().map(|e| e.original_tokens).sum();
if total_original == 0 {
return 1.0;
}
self.total_tokens_sent as f64 / total_original as f64
}
pub fn files_by_token_cost(&self) -> Vec<(String, usize)> {
let mut costs: Vec<(String, usize)> = self
.entries
.iter()
.map(|e| (e.path.clone(), e.sent_tokens))
.collect();
costs.sort_by_key(|b| std::cmp::Reverse(b.1));
costs
}
pub fn mode_distribution(&self) -> HashMap<String, usize> {
let mut dist: HashMap<String, usize> = HashMap::new();
for entry in &self.entries {
*dist.entry(entry.mode.clone()).or_insert(0) += 1;
}
dist
}
pub fn eviction_candidates(&self, keep_count: usize) -> Vec<String> {
if self.entries.len() <= keep_count {
return Vec::new();
}
let mut sorted = self.entries.clone();
sorted.sort_by_key(|e| e.timestamp);
sorted
.iter()
.take(self.entries.len() - keep_count)
.map(|e| e.path.clone())
.collect()
}
pub fn remove(&mut self, path: &str) {
if let Some(idx) = self.entries.iter().position(|e| e.path == path) {
let entry = &self.entries[idx];
self.total_tokens_sent -= entry.sent_tokens;
self.total_tokens_saved -= entry.original_tokens.saturating_sub(entry.sent_tokens);
self.entries.remove(idx);
}
}
pub fn save(&self) {
if let Ok(dir) = crate::core::data_dir::lean_ctx_data_dir() {
let path = dir.join("context_ledger.json");
if let Ok(json) = serde_json::to_string(self) {
let _ = std::fs::write(path, json);
}
}
}
pub fn load() -> Self {
crate::core::data_dir::lean_ctx_data_dir()
.ok()
.map(|d| d.join("context_ledger.json"))
.and_then(|p| std::fs::read_to_string(p).ok())
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default()
}
pub fn format_summary(&self) -> String {
let pressure = self.pressure();
format!(
"CTX: {}/{} tokens ({:.0}%), {} files, ratio {:.2}, action: {:?}",
self.total_tokens_sent,
self.window_size,
pressure.utilization * 100.0,
self.entries.len(),
self.compression_ratio(),
pressure.recommendation,
)
}
}
#[derive(Debug, Clone)]
pub struct ReinjectionAction {
pub path: String,
pub current_mode: String,
pub new_mode: String,
pub tokens_freed: usize,
}
#[derive(Debug, Clone)]
pub struct ReinjectionPlan {
pub actions: Vec<ReinjectionAction>,
pub total_tokens_freed: usize,
pub new_utilization: f64,
}
impl ContextLedger {
pub fn reinjection_plan(
&self,
intent: &super::intent_engine::StructuredIntent,
target_utilization: f64,
) -> ReinjectionPlan {
let current_util = self.total_tokens_sent as f64 / self.window_size as f64;
if current_util <= target_utilization {
return ReinjectionPlan {
actions: Vec::new(),
total_tokens_freed: 0,
new_utilization: current_util,
};
}
let tokens_to_free =
self.total_tokens_sent - (self.window_size as f64 * target_utilization) as usize;
let target_set: std::collections::HashSet<&str> = intent
.targets
.iter()
.map(std::string::String::as_str)
.collect();
let mut candidates: Vec<(usize, &LedgerEntry)> = self
.entries
.iter()
.enumerate()
.filter(|(_, e)| !target_set.iter().any(|t| e.path.contains(t)))
.collect();
candidates.sort_by(|a, b| {
let a_phi = a.1.phi.unwrap_or(0.0);
let b_phi = b.1.phi.unwrap_or(0.0);
a_phi
.partial_cmp(&b_phi)
.unwrap_or_else(|| a.1.timestamp.cmp(&b.1.timestamp))
});
let mut actions = Vec::new();
let mut freed = 0usize;
for (_, entry) in &candidates {
if freed >= tokens_to_free {
break;
}
if let Some((new_mode, new_tokens)) = downgrade_mode(&entry.mode, entry.sent_tokens) {
let saving = entry.sent_tokens.saturating_sub(new_tokens);
if saving > 0 {
actions.push(ReinjectionAction {
path: entry.path.clone(),
current_mode: entry.mode.clone(),
new_mode,
tokens_freed: saving,
});
freed += saving;
}
}
}
let new_sent = self.total_tokens_sent.saturating_sub(freed);
let new_utilization = new_sent as f64 / self.window_size as f64;
ReinjectionPlan {
actions,
total_tokens_freed: freed,
new_utilization,
}
}
}
fn downgrade_mode(current_mode: &str, current_tokens: usize) -> Option<(String, usize)> {
match current_mode {
"full" => Some(("signatures".to_string(), current_tokens / 5)),
"aggressive" => Some(("signatures".to_string(), current_tokens / 3)),
"signatures" => Some(("map".to_string(), current_tokens / 2)),
"map" => Some(("reference".to_string(), current_tokens / 4)),
_ => None,
}
}
impl Default for ContextLedger {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn new_ledger_is_empty() {
let ledger = ContextLedger::new();
assert_eq!(ledger.total_tokens_sent, 0);
assert_eq!(ledger.entries.len(), 0);
assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
}
#[test]
fn record_tracks_tokens() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("src/main.rs", "full", 500, 500);
ledger.record("src/lib.rs", "signatures", 1000, 200);
assert_eq!(ledger.total_tokens_sent, 700);
assert_eq!(ledger.total_tokens_saved, 800);
assert_eq!(ledger.entries.len(), 2);
}
#[test]
fn record_updates_existing_entry() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("src/main.rs", "full", 500, 500);
ledger.record("src/main.rs", "signatures", 500, 100);
assert_eq!(ledger.entries.len(), 1);
assert_eq!(ledger.total_tokens_sent, 100);
assert_eq!(ledger.total_tokens_saved, 400);
}
#[test]
fn pressure_escalates() {
let mut ledger = ContextLedger::with_window_size(1000);
ledger.record("a.rs", "full", 600, 600);
assert_eq!(
ledger.pressure().recommendation,
PressureAction::SuggestCompression
);
ledger.record("b.rs", "full", 200, 200);
assert_eq!(
ledger.pressure().recommendation,
PressureAction::ForceCompression
);
ledger.record("c.rs", "full", 150, 150);
assert_eq!(
ledger.pressure().recommendation,
PressureAction::EvictLeastRelevant
);
}
#[test]
fn compression_ratio_accurate() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("a.rs", "full", 1000, 1000);
ledger.record("b.rs", "signatures", 1000, 200);
let ratio = ledger.compression_ratio();
assert!((ratio - 0.6).abs() < 0.01);
}
#[test]
fn eviction_returns_oldest() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("old.rs", "full", 100, 100);
std::thread::sleep(std::time::Duration::from_millis(10));
ledger.record("new.rs", "full", 100, 100);
let candidates = ledger.eviction_candidates(1);
assert_eq!(candidates, vec!["old.rs"]);
}
#[test]
fn remove_updates_totals() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("a.rs", "full", 500, 500);
ledger.record("b.rs", "full", 300, 300);
ledger.remove("a.rs");
assert_eq!(ledger.total_tokens_sent, 300);
assert_eq!(ledger.entries.len(), 1);
}
#[test]
fn mode_distribution_counts() {
let mut ledger = ContextLedger::new();
ledger.record("a.rs", "full", 100, 100);
ledger.record("b.rs", "signatures", 100, 50);
ledger.record("c.rs", "full", 100, 100);
let dist = ledger.mode_distribution();
assert_eq!(dist.get("full"), Some(&2));
assert_eq!(dist.get("signatures"), Some(&1));
}
#[test]
fn format_summary_includes_key_info() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("a.rs", "full", 500, 500);
let summary = ledger.format_summary();
assert!(summary.contains("500/10000"));
assert!(summary.contains("1 files"));
}
#[test]
fn reinjection_no_action_when_low_pressure() {
use crate::core::intent_engine::StructuredIntent;
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("a.rs", "full", 100, 100);
let intent = StructuredIntent::from_query("fix bug in a.rs");
let plan = ledger.reinjection_plan(&intent, 0.7);
assert!(plan.actions.is_empty());
assert_eq!(plan.total_tokens_freed, 0);
}
#[test]
fn reinjection_downgrades_non_target_files() {
use crate::core::intent_engine::StructuredIntent;
let mut ledger = ContextLedger::with_window_size(1000);
ledger.record("src/target.rs", "full", 400, 400);
std::thread::sleep(std::time::Duration::from_millis(10));
ledger.record("src/other.rs", "full", 400, 400);
std::thread::sleep(std::time::Duration::from_millis(10));
ledger.record("src/utils.rs", "full", 200, 200);
let intent = StructuredIntent::from_query("fix bug in target.rs");
let plan = ledger.reinjection_plan(&intent, 0.5);
assert!(!plan.actions.is_empty());
assert!(
plan.actions.iter().all(|a| !a.path.contains("target")),
"should not downgrade target file"
);
assert!(plan.total_tokens_freed > 0);
}
#[test]
fn reinjection_preserves_targets() {
use crate::core::intent_engine::StructuredIntent;
let mut ledger = ContextLedger::with_window_size(1000);
ledger.record("src/auth.rs", "full", 900, 900);
let intent = StructuredIntent::from_query("fix bug in auth.rs");
let plan = ledger.reinjection_plan(&intent, 0.5);
assert!(
plan.actions.is_empty(),
"should not downgrade target files even under pressure"
);
}
#[test]
fn downgrade_mode_chain() {
assert_eq!(
downgrade_mode("full", 1000),
Some(("signatures".to_string(), 200))
);
assert_eq!(
downgrade_mode("signatures", 200),
Some(("map".to_string(), 100))
);
assert_eq!(
downgrade_mode("map", 100),
Some(("reference".to_string(), 25))
);
assert_eq!(downgrade_mode("reference", 25), None);
}
#[test]
fn record_assigns_item_id() {
let mut ledger = ContextLedger::new();
ledger.record("src/main.rs", "full", 500, 500);
let entry = &ledger.entries[0];
assert!(entry.id.is_some());
assert_eq!(entry.id.as_ref().unwrap().as_str(), "file:src/main.rs");
}
#[test]
fn record_sets_state_to_included() {
let mut ledger = ContextLedger::new();
ledger.record("src/main.rs", "full", 500, 500);
assert_eq!(
ledger.entries[0].state,
Some(crate::core::context_field::ContextState::Included)
);
}
#[test]
fn record_generates_view_costs() {
let mut ledger = ContextLedger::new();
ledger.record("src/main.rs", "full", 5000, 5000);
let vc = ledger.entries[0].view_costs.as_ref().unwrap();
assert_eq!(vc.get(&crate::core::context_field::ViewKind::Full), 5000);
assert_eq!(
vc.get(&crate::core::context_field::ViewKind::Signatures),
1000
);
}
#[test]
fn update_phi_works() {
let mut ledger = ContextLedger::new();
ledger.record("a.rs", "full", 100, 100);
ledger.update_phi("a.rs", 0.85);
assert_eq!(ledger.entries[0].phi, Some(0.85));
}
#[test]
fn set_state_works() {
let mut ledger = ContextLedger::new();
ledger.record("a.rs", "full", 100, 100);
ledger.set_state("a.rs", crate::core::context_field::ContextState::Pinned);
assert_eq!(
ledger.entries[0].state,
Some(crate::core::context_field::ContextState::Pinned)
);
}
#[test]
fn items_by_state_filters() {
let mut ledger = ContextLedger::new();
ledger.record("a.rs", "full", 100, 100);
ledger.record("b.rs", "full", 100, 100);
ledger.set_state("b.rs", crate::core::context_field::ContextState::Excluded);
let included = ledger.items_by_state(crate::core::context_field::ContextState::Included);
assert_eq!(included.len(), 1);
assert_eq!(included[0].path, "a.rs");
}
#[test]
fn eviction_by_phi_prefers_low_phi() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("high.rs", "full", 100, 100);
ledger.update_phi("high.rs", 0.9);
ledger.record("low.rs", "full", 100, 100);
ledger.update_phi("low.rs", 0.1);
let candidates = ledger.eviction_candidates_by_phi(1);
assert_eq!(candidates, vec!["low.rs"]);
}
#[test]
fn eviction_by_phi_skips_pinned() {
let mut ledger = ContextLedger::with_window_size(10000);
ledger.record("pinned.rs", "full", 100, 100);
ledger.update_phi("pinned.rs", 0.01);
ledger.set_state(
"pinned.rs",
crate::core::context_field::ContextState::Pinned,
);
ledger.record("normal.rs", "full", 100, 100);
ledger.update_phi("normal.rs", 0.5);
let candidates = ledger.eviction_candidates_by_phi(1);
assert_eq!(candidates, vec!["normal.rs"]);
}
#[test]
fn mark_stale_by_hash_detects_change() {
let mut ledger = ContextLedger::new();
ledger.record("a.rs", "full", 100, 100);
ledger.entries[0].source_hash = Some("hash_v1".to_string());
ledger.mark_stale_by_hash("a.rs", "hash_v2");
assert_eq!(
ledger.entries[0].state,
Some(crate::core::context_field::ContextState::Stale)
);
}
#[test]
fn find_by_id_works() {
let mut ledger = ContextLedger::new();
ledger.record("src/lib.rs", "full", 100, 100);
let id = crate::core::context_field::ContextItemId::from_file("src/lib.rs");
assert!(ledger.find_by_id(&id).is_some());
}
#[test]
fn upsert_sets_source_hash_and_kind() {
let mut ledger = ContextLedger::new();
ledger.upsert(
"src/main.rs",
"full",
500,
500,
Some("sha256_abc"),
crate::core::context_field::ContextKind::File,
None,
);
let entry = &ledger.entries[0];
assert_eq!(entry.source_hash.as_deref(), Some("sha256_abc"));
assert_eq!(
entry.kind,
Some(crate::core::context_field::ContextKind::File)
);
}
}