use compact_str::CompactString;
use serde::{Deserialize, Serialize};
use crate::context::pressure::PressureAction;
pub type HandleId = u32;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HandleKind {
ToolResult,
MemoryPage,
KnowledgeEntry,
SubAgentJoin,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Residency {
Resident,
External {
payload_ref: String,
digest: String,
original_size: u64,
},
PagedOut { payload_ref: String, digest: String },
Collapsed,
}
impl Residency {
pub fn label(&self) -> &'static str {
match self {
Self::Resident => "resident",
Self::External { .. } => "external",
Self::PagedOut { .. } => "paged_out",
Self::Collapsed => "collapsed",
}
}
pub fn occupies_context(&self) -> bool {
matches!(self, Self::Resident)
}
pub fn payload_ref(&self) -> Option<&str> {
match self {
Self::External { payload_ref, .. } | Self::PagedOut { payload_ref, .. } => {
Some(payload_ref.as_str())
}
Self::Resident | Self::Collapsed => None,
}
}
pub fn digest(&self) -> Option<&str> {
match self {
Self::External { digest, .. } | Self::PagedOut { digest, .. } => Some(digest.as_str()),
Self::Resident | Self::Collapsed => None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Handle {
pub id: HandleId,
pub kind: HandleKind,
pub residency: Residency,
pub tokens: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub source: Option<CompactString>,
}
impl Handle {
pub fn resident(id: HandleId, kind: HandleKind, tokens: u32) -> Self {
Self {
id,
kind,
residency: Residency::Resident,
tokens,
source: None,
}
}
pub fn resident_for(
id: HandleId,
kind: HandleKind,
tokens: u32,
source: impl Into<CompactString>,
) -> Self {
Self {
id,
kind,
residency: Residency::Resident,
tokens,
source: Some(source.into()),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct HandleTable {
handles: Vec<Handle>,
}
impl HandleTable {
pub fn new() -> Self {
Self::default()
}
pub fn insert(&mut self, handle: Handle) {
if let Some(existing) = self.handles.iter_mut().find(|h| h.id == handle.id) {
*existing = handle;
} else {
self.handles.push(handle);
}
}
pub fn get(&self, id: HandleId) -> Option<&Handle> {
self.handles.iter().find(|h| h.id == id)
}
pub fn get_mut(&mut self, id: HandleId) -> Option<&mut Handle> {
self.handles.iter_mut().find(|h| h.id == id)
}
pub fn all(&self) -> &[Handle] {
&self.handles
}
pub fn all_mut(&mut self) -> &mut [Handle] {
&mut self.handles
}
pub fn retain(&mut self, keep: impl FnMut(&Handle) -> bool) {
self.handles.retain(keep);
}
pub fn residency_for_source(&self, source: &str) -> Option<&Residency> {
self.handles
.iter()
.find(|h| h.source.as_deref() == Some(source))
.map(|h| &h.residency)
}
pub fn tool_result_handles_mut(&mut self) -> impl Iterator<Item = &mut Handle> {
self.handles
.iter_mut()
.filter(|h| matches!(h.kind, HandleKind::ToolResult))
}
pub fn resident_tokens(&self) -> u32 {
self.handles
.iter()
.filter(|h| h.residency.occupies_context())
.map(|h| h.tokens)
.sum()
}
pub fn non_resident_tokens(&self) -> u32 {
self.handles
.iter()
.filter(|h| !h.residency.occupies_context())
.map(|h| h.tokens)
.sum()
}
}
#[derive(Debug, Clone)]
pub enum EvictionOp {
Snip { per_msg_ratio: f64 },
TimeDecayMicro,
Collapse { target_tokens: u32 },
AutoCompact { preserve_turns: usize },
}
impl EvictionOp {
pub fn label(&self) -> &'static str {
match self {
Self::Snip { .. } => "snip",
Self::TimeDecayMicro => "time_decay_micro",
Self::Collapse { .. } => "collapse",
Self::AutoCompact { .. } => "auto_compact",
}
}
pub fn invalidates_prefix_at(&self) -> Option<usize> {
match self {
Self::Snip { .. } => Some(0), Self::TimeDecayMicro => None,
Self::Collapse { .. } => Some(0),
Self::AutoCompact { .. } => Some(0),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct EvictionPlan {
pub ops: Vec<EvictionOp>,
}
impl EvictionPlan {
pub fn empty() -> Self {
Self::default()
}
pub fn is_empty(&self) -> bool {
self.ops.is_empty()
}
pub fn has_time_decay(&self) -> bool {
self.ops
.iter()
.any(|op| matches!(op, EvictionOp::TimeDecayMicro))
}
pub fn from_legacy_action(
action: PressureAction,
target_tokens: u32,
preserve_turns: usize,
) -> Self {
let ops = match action {
PressureAction::None => vec![],
PressureAction::SnipCompact => vec![EvictionOp::Snip {
per_msg_ratio: 0.10,
}],
PressureAction::MicroCompact => vec![EvictionOp::TimeDecayMicro],
PressureAction::ContextCollapse => vec![EvictionOp::Collapse { target_tokens }],
PressureAction::AutoCompact => vec![EvictionOp::AutoCompact { preserve_turns }],
};
Self { ops }
}
}
pub fn plan_eviction(
recommended: PressureAction,
idle_decay: bool,
target_tokens: u32,
preserve_turns: usize,
) -> EvictionPlan {
let mut ops = Vec::new();
if idle_decay {
ops.push(EvictionOp::TimeDecayMicro);
}
if recommended != PressureAction::None {
ops.extend(
EvictionPlan::from_legacy_action(recommended, target_tokens, preserve_turns).ops,
);
}
EvictionPlan { ops }
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resident_tokens_counts_only_resident() {
let mut table = HandleTable::new();
table.insert(Handle::resident(1, HandleKind::ToolResult, 100));
table.insert(Handle {
id: 2,
kind: HandleKind::ToolResult,
residency: Residency::External {
payload_ref: "payload:x".into(),
digest: "sha256:".to_string() + &"a".repeat(64),
original_size: 5_000,
},
tokens: 5000,
source: None,
});
table.insert(Handle {
id: 3,
kind: HandleKind::MemoryPage,
residency: Residency::Collapsed,
tokens: 200,
source: None,
});
assert_eq!(table.resident_tokens(), 100);
}
#[test]
fn handle_table_insert_is_idempotent_by_id() {
let mut table = HandleTable::new();
table.insert(Handle::resident(1, HandleKind::ToolResult, 100));
table.insert(Handle::resident(1, HandleKind::ToolResult, 250));
assert_eq!(table.all().len(), 1);
assert_eq!(table.get(1).unwrap().tokens, 250);
}
#[test]
fn residency_occupies_context_only_when_resident() {
assert!(Residency::Resident.occupies_context());
assert!(!Residency::Collapsed.occupies_context());
assert!(!paged_out().occupies_context());
assert!(!external().occupies_context());
}
fn external() -> Residency {
Residency::External {
payload_ref: "payload:01J".into(),
digest: "sha256:".to_string() + &"a".repeat(64),
original_size: 90_000,
}
}
fn paged_out() -> Residency {
Residency::PagedOut {
payload_ref: "payload:02K".into(),
digest: "sha256:".to_string() + &"b".repeat(64),
}
}
#[test]
fn only_externally_backed_residencies_expose_a_locator_and_a_digest() {
for residency in [external(), paged_out()] {
assert!(residency.payload_ref().is_some(), "{residency:?}");
assert!(residency.digest().is_some(), "{residency:?}");
}
for residency in [Residency::Resident, Residency::Collapsed] {
assert_eq!(residency.payload_ref(), None, "{residency:?}");
assert_eq!(residency.digest(), None, "{residency:?}");
}
}
#[test]
fn external_and_paged_out_are_distinguishable_states() {
assert_eq!(external().label(), "external");
assert_eq!(paged_out().label(), "paged_out");
assert_ne!(external(), paged_out());
}
#[test]
fn plan_eviction_empty_when_no_pressure_and_no_idle() {
assert!(plan_eviction(PressureAction::None, false, 50_000, 2).is_empty());
}
#[test]
fn plan_eviction_emits_specific_op_for_recommended_action() {
let plan = plan_eviction(PressureAction::AutoCompact, false, 50_000, 3);
assert!(matches!(
&plan.ops[..],
[EvictionOp::AutoCompact { preserve_turns: 3 }]
));
}
#[test]
fn plan_eviction_collapse_carries_caller_target_tokens() {
let plan = plan_eviction(PressureAction::ContextCollapse, false, 12_345, 2);
assert!(matches!(
&plan.ops[..],
[EvictionOp::Collapse {
target_tokens: 12_345
}]
));
}
#[test]
fn plan_eviction_orders_time_decay_before_pressure() {
let plan = plan_eviction(PressureAction::ContextCollapse, true, 50_000, 2);
assert_eq!(plan.ops.len(), 2);
assert!(matches!(plan.ops[0], EvictionOp::TimeDecayMicro));
assert!(matches!(plan.ops[1], EvictionOp::Collapse { .. }));
}
#[test]
fn plan_eviction_time_decay_only() {
let plan = plan_eviction(PressureAction::None, true, 50_000, 2);
assert_eq!(plan.ops.len(), 1);
assert!(matches!(plan.ops[0], EvictionOp::TimeDecayMicro));
}
#[test]
fn plan_eviction_micro_compact_emits_time_decay_without_idle() {
let plan = plan_eviction(PressureAction::MicroCompact, false, 50_000, 2);
assert!(
plan.has_time_decay(),
"MicroCompact yields a time-decay op even when not idle"
);
for recommended in [
PressureAction::None,
PressureAction::MicroCompact,
PressureAction::AutoCompact,
PressureAction::ContextCollapse,
] {
for idle in [false, true] {
let p = plan_eviction(recommended, idle, 50_000, 2);
assert!(
!idle || p.has_time_decay(),
"idle_decay must imply a time-decay op"
);
}
}
}
#[test]
fn eviction_op_labels() {
assert_eq!(EvictionOp::Snip { per_msg_ratio: 0.1 }.label(), "snip");
assert_eq!(EvictionOp::TimeDecayMicro.label(), "time_decay_micro");
assert_eq!(
EvictionOp::Collapse {
target_tokens: 5000
}
.label(),
"collapse"
);
assert_eq!(
EvictionOp::AutoCompact { preserve_turns: 2 }.label(),
"auto_compact"
);
}
}