use serde::Serialize;
use std::cmp::Ordering;
use super::memory_archive::{ArchiveConfig, MemoryStore, archive_items};
pub(crate) fn reclaim_target(max: usize, headroom_pct: f32) -> usize {
if max == 0 {
return 0;
}
let pct = headroom_pct.clamp(0.0, 0.95);
let drop = ((max as f32) * pct).ceil() as usize;
max.saturating_sub(drop)
}
pub(crate) fn should_reclaim(len: usize, max: usize, enabled: bool) -> bool {
enabled && max > 0 && len >= max
}
pub(crate) fn reclaim_preview(len: usize, max: usize, headroom_pct: f32, enabled: bool) -> usize {
if !should_reclaim(len, max, enabled) {
return 0;
}
len.saturating_sub(reclaim_target(max, headroom_pct))
}
pub(crate) fn reclaim_store<T, F>(
store: MemoryStore,
scope: Option<&str>,
items: &mut Vec<T>,
max: usize,
headroom_pct: f32,
enabled: bool,
mut retention_cmp: F,
) -> Result<Vec<T>, String>
where
T: Serialize,
F: FnMut(&T, &T) -> Ordering,
{
if !should_reclaim(items.len(), max, enabled) {
return Ok(Vec::new());
}
let target = reclaim_target(max, headroom_pct);
let drop_count = items.len().saturating_sub(target);
if drop_count == 0 {
return Ok(Vec::new());
}
let mut ranked: Vec<usize> = (0..items.len()).collect();
ranked.sort_by(|&a, &b| retention_cmp(&items[a], &items[b]));
let mut evict: Vec<usize> = ranked[target..].to_vec();
evict.sort_unstable();
let candidates: Vec<&T> = evict.iter().map(|&idx| &items[idx]).collect();
archive_items(store, scope, &candidates, &ArchiveConfig::from_env())?;
let mut archived: Vec<T> = Vec::with_capacity(evict.len());
for &idx in evict.iter().rev() {
archived.push(items.remove(idx));
}
archived.reverse();
Ok(archived)
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Utc;
use serde::Deserialize;
fn with_temp_data_dir<T>(f: impl FnOnce() -> T) -> T {
let _lock = crate::core::data_dir::test_env_lock();
let dir = std::env::temp_dir().join(format!(
"lctx-capacity-{}-{}",
std::process::id(),
Utc::now().timestamp_nanos_opt().unwrap_or(0)
));
let _ = std::fs::create_dir_all(&dir);
crate::test_env::set_var("LEAN_CTX_DATA_DIR", dir.to_str().unwrap());
let out = f();
crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
let _ = std::fs::remove_dir_all(&dir);
out
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
struct Item {
rank: u32,
}
#[test]
fn reclaim_target_matches_legacy_quarter_reclaim() {
let legacy = |max: usize| max.saturating_sub(max.div_ceil(4));
for max in [1usize, 2, 3, 4, 5, 6, 7, 8, 10, 100, 200, 1000] {
assert_eq!(
reclaim_target(max, 0.25),
legacy(max),
"mismatch at max={max}"
);
}
}
#[test]
fn reclaim_target_zero_headroom_keeps_all() {
assert_eq!(reclaim_target(200, 0.0), 200);
}
#[test]
fn reclaim_target_is_clamped() {
assert!(reclaim_target(100, 9.9) >= 5);
}
#[test]
fn should_reclaim_hysteresis() {
assert!(!should_reclaim(99, 100, true), "under cap: no reclaim");
assert!(should_reclaim(100, 100, true), "at cap: reclaim");
assert!(should_reclaim(150, 100, true), "over cap: reclaim");
assert!(!should_reclaim(150, 100, false), "disabled: no reclaim");
assert!(!should_reclaim(150, 0, true), "max 0: no reclaim");
}
#[test]
fn reclaim_store_is_lossless_and_keeps_best() {
with_temp_data_dir(|| {
let mut items: Vec<Item> = (0..8).map(|rank| Item { rank }).collect();
let archived = reclaim_store(
MemoryStore::Patterns,
Some("p"),
&mut items,
8,
0.25,
true,
|a, b| a.rank.cmp(&b.rank),
)
.expect("reclaim succeeds");
assert_eq!(items.len(), 6);
assert_eq!(archived.len(), 2);
let mut all: Vec<u32> = items.iter().chain(&archived).map(|i| i.rank).collect();
all.sort_unstable();
assert_eq!(all, (0..8).collect::<Vec<_>>());
assert_eq!(
archived.iter().map(|i| i.rank).collect::<Vec<_>>(),
vec![6, 7]
);
});
}
#[test]
fn reclaim_store_noop_under_cap() {
with_temp_data_dir(|| {
let mut items: Vec<Item> = (0..3).map(|rank| Item { rank }).collect();
let archived = reclaim_store(
MemoryStore::History,
Some("p"),
&mut items,
10,
0.25,
true,
|a, b| a.rank.cmp(&b.rank),
)
.expect("reclaim succeeds");
assert!(archived.is_empty());
assert_eq!(items.len(), 3);
});
}
#[test]
fn reclaim_store_respects_disabled() {
with_temp_data_dir(|| {
let mut items: Vec<Item> = (0..20).map(|rank| Item { rank }).collect();
let archived = reclaim_store(
MemoryStore::Procedures,
Some("p"),
&mut items,
10,
0.25,
false,
|a, b| a.rank.cmp(&b.rank),
)
.expect("reclaim succeeds");
assert!(archived.is_empty());
assert_eq!(
items.len(),
20,
"disabled reclaim leaves the store untouched"
);
});
}
#[test]
fn reclaim_store_preserves_items_when_archive_persistence_fails() {
let _lock = crate::core::data_dir::test_env_lock();
let path = std::env::temp_dir().join(format!(
"lctx-capacity-file-{}-{}",
std::process::id(),
Utc::now().timestamp_nanos_opt().unwrap_or(0)
));
std::fs::write(&path, b"not a directory").expect("create blocking file");
crate::test_env::set_var("LEAN_CTX_DATA_DIR", path.to_str().unwrap());
let mut items: Vec<Item> = (0..8).map(|rank| Item { rank }).collect();
let original = items.clone();
let result = reclaim_store(
MemoryStore::Patterns,
Some("p"),
&mut items,
8,
0.25,
true,
|a, b| a.rank.cmp(&b.rank),
);
crate::test_env::remove_var("LEAN_CTX_DATA_DIR");
let _ = std::fs::remove_file(path);
assert!(result.is_err(), "archive persistence must fail");
assert_eq!(items, original, "failed reclaim must not mutate live items");
}
}