use crate::providers::traits::ChatMessage;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use zeroclaw_macros::Configurable;
fn default_max_tokens() -> usize {
8192
}
fn default_keep_recent() -> usize {
4
}
fn default_collapse() -> bool {
true
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Configurable)]
#[prefix = "agent.history-pruning"]
pub struct HistoryPrunerConfig {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_max_tokens")]
pub max_tokens: usize,
#[serde(default = "default_keep_recent")]
pub keep_recent: usize,
#[serde(default = "default_collapse")]
pub collapse_tool_results: bool,
}
impl Default for HistoryPrunerConfig {
fn default() -> Self {
Self {
enabled: false,
max_tokens: 8192,
keep_recent: 4,
collapse_tool_results: true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PruneStats {
pub messages_before: usize,
pub messages_after: usize,
pub collapsed_pairs: usize,
pub dropped_messages: usize,
}
fn estimate_tokens(messages: &[ChatMessage]) -> usize {
messages.iter().map(|m| m.content.len() / 4).sum()
}
fn protected_indices(messages: &[ChatMessage], keep_recent: usize) -> Vec<bool> {
let len = messages.len();
let mut protected = vec![false; len];
for (i, msg) in messages.iter().enumerate() {
if msg.role == "system" {
protected[i] = true;
}
}
let recent_start = len.saturating_sub(keep_recent);
for p in protected.iter_mut().skip(recent_start) {
*p = true;
}
protected
}
pub fn prune_history(messages: &mut Vec<ChatMessage>, config: &HistoryPrunerConfig) -> PruneStats {
let messages_before = messages.len();
if !config.enabled || messages.is_empty() {
return PruneStats {
messages_before,
messages_after: messages_before,
collapsed_pairs: 0,
dropped_messages: 0,
};
}
let mut collapsed_pairs: usize = 0;
if config.collapse_tool_results {
let mut i = 0;
while i + 1 < messages.len() {
let protected = protected_indices(messages, config.keep_recent);
if messages[i].role == "assistant"
&& messages[i + 1].role == "tool"
&& !protected[i]
&& !protected[i + 1]
{
let tool_content = &messages[i + 1].content;
let truncated: String = tool_content.chars().take(100).collect();
let summary = format!("[Tool result: {truncated}...]");
messages[i] = ChatMessage {
role: "assistant".to_string(),
content: summary,
};
messages.remove(i + 1);
collapsed_pairs += 1;
} else {
i += 1;
}
}
}
let mut dropped_messages: usize = 0;
while estimate_tokens(messages) > config.max_tokens {
let protected = protected_indices(messages, config.keep_recent);
if let Some(idx) = protected
.iter()
.enumerate()
.find(|&(_, &p)| !p)
.map(|(i, _)| i)
{
messages.remove(idx);
dropped_messages += 1;
} else {
break;
}
}
PruneStats {
messages_before,
messages_after: messages.len(),
collapsed_pairs,
dropped_messages,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn msg(role: &str, content: &str) -> ChatMessage {
ChatMessage {
role: role.to_string(),
content: content.to_string(),
}
}
#[test]
fn prune_disabled_is_noop() {
let mut messages = vec![
msg("system", "You are helpful."),
msg("user", "Hello"),
msg("assistant", "Hi there!"),
];
let config = HistoryPrunerConfig {
enabled: false,
..Default::default()
};
let stats = prune_history(&mut messages, &config);
assert_eq!(messages.len(), 3);
assert_eq!(messages[0].content, "You are helpful.");
assert_eq!(stats.messages_before, 3);
assert_eq!(stats.messages_after, 3);
assert_eq!(stats.collapsed_pairs, 0);
}
#[test]
fn prune_under_budget_no_change() {
let mut messages = vec![
msg("system", "You are helpful."),
msg("user", "Hello"),
msg("assistant", "Hi!"),
];
let config = HistoryPrunerConfig {
enabled: true,
max_tokens: 8192,
keep_recent: 2,
collapse_tool_results: false,
};
let stats = prune_history(&mut messages, &config);
assert_eq!(messages.len(), 3);
assert_eq!(stats.collapsed_pairs, 0);
assert_eq!(stats.dropped_messages, 0);
}
#[test]
fn prune_collapses_tool_pairs() {
let tool_result = "a".repeat(160);
let mut messages = vec![
msg("system", "sys"),
msg("assistant", "calling tool X"),
msg("tool", &tool_result),
msg("user", "thanks"),
msg("assistant", "done"),
];
let config = HistoryPrunerConfig {
enabled: true,
max_tokens: 100_000,
keep_recent: 2,
collapse_tool_results: true,
};
let stats = prune_history(&mut messages, &config);
assert_eq!(stats.collapsed_pairs, 1);
assert_eq!(messages.len(), 4);
assert_eq!(messages[1].role, "assistant");
assert!(messages[1].content.starts_with("[Tool result: "));
}
#[test]
fn prune_preserves_system_and_recent() {
let big = "x".repeat(40_000);
let mut messages = vec![
msg("system", "system prompt"),
msg("user", &big),
msg("assistant", "old reply"),
msg("user", "recent1"),
msg("assistant", "recent2"),
];
let config = HistoryPrunerConfig {
enabled: true,
max_tokens: 100,
keep_recent: 2,
collapse_tool_results: false,
};
let stats = prune_history(&mut messages, &config);
assert!(messages.iter().any(|m| m.role == "system"));
assert!(messages.iter().any(|m| m.content == "recent1"));
assert!(messages.iter().any(|m| m.content == "recent2"));
assert!(stats.dropped_messages > 0);
}
#[test]
fn prune_drops_oldest_when_over_budget() {
let filler = "y".repeat(400);
let mut messages = vec![
msg("system", "sys"),
msg("user", &filler),
msg("assistant", &filler),
msg("user", "recent-user"),
msg("assistant", "recent-assistant"),
];
let config = HistoryPrunerConfig {
enabled: true,
max_tokens: 150,
keep_recent: 2,
collapse_tool_results: false,
};
let stats = prune_history(&mut messages, &config);
assert!(stats.dropped_messages >= 1);
assert_eq!(messages[0].role, "system");
assert!(messages.iter().any(|m| m.content == "recent-user"));
assert!(messages.iter().any(|m| m.content == "recent-assistant"));
}
#[test]
fn prune_empty_messages() {
let mut messages: Vec<ChatMessage> = vec![];
let config = HistoryPrunerConfig {
enabled: true,
..Default::default()
};
let stats = prune_history(&mut messages, &config);
assert_eq!(stats.messages_before, 0);
assert_eq!(stats.messages_after, 0);
}
}