use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MessageRole {
System,
User,
Assistant,
Tool,
}
#[derive(Debug, Clone)]
pub struct MessageMeta {
pub index: usize,
pub role: MessageRole,
pub droppable: bool,
pub has_tool_call: bool,
pub is_tool_result: bool,
pub tool_id: Option<String>,
pub token_count: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EvictionRange {
pub start: usize,
pub end: usize,
}
impl EvictionRange {
pub fn new(start: usize, end: usize) -> Self {
Self { start, end }
}
pub fn len(&self) -> usize {
self.end.saturating_sub(self.start)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CompactionStrategy {
Evict(f64),
Retain(usize),
Min(Box<CompactionStrategy>, Box<CompactionStrategy>),
Max(Box<CompactionStrategy>, Box<CompactionStrategy>),
}
impl Default for CompactionStrategy {
fn default() -> Self {
Self::Min(Box::new(Self::Evict(0.6)), Box::new(Self::Retain(10)))
}
}
impl CompactionStrategy {
pub fn calculate_eviction_range(
&self,
messages: &[MessageMeta],
retention_window: usize,
) -> Option<EvictionRange> {
if messages.len() <= retention_window {
return None; }
let raw_end = self.calculate_raw_end(messages.len(), retention_window);
let start = Self::find_safe_start(messages);
if start >= raw_end {
return None; }
let end = Self::adjust_end_for_tool_safety(messages, raw_end, retention_window);
if start >= end {
return None;
}
Some(EvictionRange::new(start, end))
}
fn calculate_raw_end(&self, total: usize, retention_window: usize) -> usize {
match self {
Self::Evict(fraction) => {
let evict_count = (total as f64 * fraction).floor() as usize;
total.saturating_sub(retention_window).min(evict_count)
}
Self::Retain(keep) => total.saturating_sub(*keep.max(&retention_window)),
Self::Min(a, b) => {
let end_a = a.calculate_raw_end(total, retention_window);
let end_b = b.calculate_raw_end(total, retention_window);
end_a.min(end_b)
}
Self::Max(a, b) => {
let end_a = a.calculate_raw_end(total, retention_window);
let end_b = b.calculate_raw_end(total, retention_window);
end_a.max(end_b)
}
}
}
fn find_safe_start(messages: &[MessageMeta]) -> usize {
messages
.iter()
.position(|m| m.role == MessageRole::Assistant)
.unwrap_or(0)
}
fn adjust_end_for_tool_safety(
messages: &[MessageMeta],
mut end: usize,
retention_window: usize,
) -> usize {
let min_end = messages.len().saturating_sub(retention_window);
if end > min_end {
end = min_end;
}
if end == 0 || end >= messages.len() {
return end;
}
let last_evicted = &messages[end - 1];
if last_evicted.has_tool_call {
if let Some(tool_id) = &last_evicted.tool_id {
for (i, msg) in messages.iter().enumerate().skip(end).take(5) {
if msg.is_tool_result && msg.tool_id.as_ref() == Some(tool_id) {
end = i + 1;
break;
}
}
}
}
let msg_at_end = messages.get(end);
if let Some(msg) = msg_at_end
&& msg.is_tool_result
{
while end > 0 {
let prev = &messages[end - 1];
if prev.is_tool_result || prev.has_tool_call {
end -= 1;
} else {
break;
}
}
}
while end > 0 && end < messages.len() {
if messages[end].is_tool_result {
end -= 1;
} else {
break;
}
}
end
}
pub fn filter_droppable(messages: &[MessageMeta], range: &EvictionRange) -> Vec<usize> {
(range.start..range.end)
.filter(|&i| !messages[i].droppable)
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_messages(roles: &[(MessageRole, bool, bool)]) -> Vec<MessageMeta> {
roles
.iter()
.enumerate()
.map(|(i, (role, has_tool_call, is_tool_result))| MessageMeta {
index: i,
role: *role,
droppable: false,
has_tool_call: *has_tool_call,
is_tool_result: *is_tool_result,
tool_id: if *has_tool_call || *is_tool_result {
Some(format!("tool_{}", i))
} else {
None
},
token_count: 100,
})
.collect()
}
#[test]
fn test_eviction_range_empty() {
let strategy = CompactionStrategy::Retain(10);
let messages = make_messages(&[
(MessageRole::System, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, false, false),
]);
let range = strategy.calculate_eviction_range(&messages, 5);
assert!(range.is_none());
}
#[test]
fn test_eviction_starts_at_assistant() {
let strategy = CompactionStrategy::Evict(0.5);
let messages = make_messages(&[
(MessageRole::System, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, false, false),
]);
let range = strategy.calculate_eviction_range(&messages, 2);
assert!(range.is_some());
let range = range.unwrap();
assert_eq!(range.start, 2);
}
#[test]
fn test_tool_call_result_adjacency() {
let mut messages = make_messages(&[
(MessageRole::System, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, true, false), (MessageRole::Tool, false, true), (MessageRole::Assistant, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, false, false),
]);
messages[2].tool_id = Some("call_1".to_string());
messages[3].tool_id = Some("call_1".to_string());
let strategy = CompactionStrategy::Retain(2);
let range = strategy.calculate_eviction_range(&messages, 2);
if let Some(range) = range {
if range.end > 2 && range.end <= 3 {
panic!("Eviction split tool call from result!");
}
}
}
#[test]
fn test_filter_droppable() {
let mut messages = make_messages(&[
(MessageRole::System, false, false),
(MessageRole::User, false, false),
(MessageRole::Assistant, false, false),
(MessageRole::User, false, false), (MessageRole::Assistant, false, false),
]);
messages[3].droppable = true;
let range = EvictionRange::new(0, 5);
let non_droppable = CompactionStrategy::filter_droppable(&messages, &range);
assert_eq!(non_droppable.len(), 4);
assert!(!non_droppable.contains(&3));
}
#[test]
fn test_min_strategy() {
let strategy = CompactionStrategy::Min(
Box::new(CompactionStrategy::Evict(0.8)),
Box::new(CompactionStrategy::Retain(5)),
);
let messages = make_messages(&vec![(MessageRole::Assistant, false, false); 10]);
let range = strategy.calculate_eviction_range(&messages, 3);
assert!(range.is_some());
}
}