use std::sync::Arc;
use async_trait::async_trait;
use crate::error::Result;
use crate::memory::{ContextProvider, SessionContext};
use crate::types::{Content, Message, Role};
pub trait Tokenizer: Send + Sync {
fn count_tokens(&self, text: &str) -> usize;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ApproxTokenizer;
impl Tokenizer for ApproxTokenizer {
fn count_tokens(&self, text: &str) -> usize {
text.chars().count().div_ceil(4)
}
}
pub fn count_message_tokens(tokenizer: &dyn Tokenizer, message: &Message) -> usize {
message
.contents
.iter()
.map(|content| count_content_tokens(tokenizer, content))
.sum()
}
fn count_content_tokens(tokenizer: &dyn Tokenizer, content: &Content) -> usize {
fn value_text(value: &serde_json::Value) -> String {
match value {
serde_json::Value::String(s) => s.clone(),
other => other.to_string(),
}
}
match content {
Content::Text(t) => tokenizer.count_tokens(&t.text),
Content::TextReasoning(t) => tokenizer.count_tokens(&t.text),
Content::FunctionCall(fc) => {
let arguments = match &fc.arguments {
Some(crate::types::FunctionArguments::Raw(raw)) => raw.clone(),
Some(crate::types::FunctionArguments::Object(map)) => {
serde_json::to_string(map).unwrap_or_default()
}
None => String::new(),
};
tokenizer.count_tokens(&fc.name) + tokenizer.count_tokens(&arguments)
}
Content::FunctionResult(fr) => {
let mut total = fr
.result
.as_ref()
.map_or(0, |v| tokenizer.count_tokens(&value_text(v)));
if let Some(exception) = &fr.exception {
total += tokenizer.count_tokens(exception);
}
total
}
_ => 0,
}
}
pub trait CompactionStrategy: Send + Sync {
fn compact(&self, messages: &[Message], tokenizer: &dyn Tokenizer) -> Vec<Message>;
}
fn leading_system_count(messages: &[Message]) -> usize {
messages
.iter()
.take_while(|m| m.role == Role::system())
.count()
}
#[derive(Debug, Clone, Copy)]
pub struct Truncation {
pub max_messages: usize,
}
impl Truncation {
pub fn new(max_messages: usize) -> Self {
Self { max_messages }
}
}
impl CompactionStrategy for Truncation {
fn compact(&self, messages: &[Message], _tokenizer: &dyn Tokenizer) -> Vec<Message> {
if messages.len() <= self.max_messages {
return messages.to_vec();
}
let sys_count = leading_system_count(messages);
let mut out: Vec<Message> = messages[..sys_count].to_vec();
if sys_count >= self.max_messages {
out.truncate(self.max_messages);
return out;
}
let remaining_budget = self.max_messages - sys_count;
let rest = &messages[sys_count..];
let start = rest.len().saturating_sub(remaining_budget);
out.extend_from_slice(&rest[start..]);
out
}
}
#[derive(Debug, Clone, Copy)]
pub struct SlidingWindow {
pub window: usize,
}
impl SlidingWindow {
pub fn new(window: usize) -> Self {
Self { window }
}
}
impl CompactionStrategy for SlidingWindow {
fn compact(&self, messages: &[Message], _tokenizer: &dyn Tokenizer) -> Vec<Message> {
let sys_count = leading_system_count(messages);
let mut out: Vec<Message> = messages[..sys_count].to_vec();
let rest = &messages[sys_count..];
let start = rest.len().saturating_sub(self.window);
out.extend_from_slice(&rest[start..]);
out
}
}
#[derive(Debug, Clone, Copy)]
pub struct TokenBudget {
pub max_tokens: usize,
}
impl TokenBudget {
pub fn new(max_tokens: usize) -> Self {
Self { max_tokens }
}
}
impl CompactionStrategy for TokenBudget {
fn compact(&self, messages: &[Message], tokenizer: &dyn Tokenizer) -> Vec<Message> {
let sys_count = leading_system_count(messages);
let system_prefix = &messages[..sys_count];
let rest = &messages[sys_count..];
let mut used: usize = system_prefix
.iter()
.map(|m| count_message_tokens(tokenizer, m))
.sum();
let mut kept_rest: Vec<&Message> = Vec::new();
for message in rest.iter().rev() {
let cost = count_message_tokens(tokenizer, message);
if !kept_rest.is_empty() && used + cost > self.max_tokens {
break;
}
used += cost;
kept_rest.push(message);
}
kept_rest.reverse();
let mut out: Vec<Message> = system_prefix.to_vec();
out.extend(kept_rest.into_iter().cloned());
out
}
}
fn has_tool_result(message: &Message) -> bool {
message
.contents
.iter()
.any(|c| matches!(c, Content::FunctionResult(_)))
}
#[derive(Debug, Clone, Copy)]
pub struct SelectiveToolResult {
pub keep_last: usize,
}
impl SelectiveToolResult {
pub fn new(keep_last: usize) -> Self {
Self { keep_last }
}
}
pub const OMITTED_TOOL_RESULT: &str = "[tool result omitted by compaction]";
impl CompactionStrategy for SelectiveToolResult {
fn compact(&self, messages: &[Message], _tokenizer: &dyn Tokenizer) -> Vec<Message> {
let tool_result_count = messages.iter().filter(|m| has_tool_result(m)).count();
let mut strip_budget = tool_result_count.saturating_sub(self.keep_last);
let mut out = Vec::with_capacity(messages.len());
for message in messages {
if has_tool_result(message) && strip_budget > 0 {
strip_budget -= 1;
let mut compacted = message.clone();
for content in &mut compacted.contents {
if let Content::FunctionResult(fr) = content {
fr.result =
Some(serde_json::Value::String(OMITTED_TOOL_RESULT.to_string()));
if fr.exception.is_some() {
fr.exception = Some(OMITTED_TOOL_RESULT.to_string());
}
}
}
out.push(compacted);
} else {
out.push(message.clone());
}
}
out
}
}
fn drop_orphaned_tool_exchanges(
messages: &mut Vec<Message>,
origins: &mut Vec<Option<usize>>,
pending: &[Message],
) {
use std::collections::HashSet;
let paired: HashSet<Site> = {
let (pairs, _) = pair_tool_exchanges(messages.iter().chain(pending.iter()).enumerate());
pairs
.into_iter()
.flat_map(|(call_site, result_site)| [call_site, result_site])
.collect()
};
let has_orphan = messages.iter().enumerate().any(|(mi, message)| {
message.contents.iter().enumerate().any(|(ci, content)| {
matches!(
content,
Content::FunctionCall(_) | Content::FunctionResult(_)
) && !paired.contains(&(mi, ci))
})
});
if !has_orphan {
return;
}
let mut mi = 0;
let mut kept = 0;
messages.retain_mut(|message| {
let this = mi;
mi += 1;
let had_contents = !message.contents.is_empty();
let mut ci = 0;
message.contents.retain(|content| {
let keep = !matches!(
content,
Content::FunctionCall(_) | Content::FunctionResult(_)
) || paired.contains(&(this, ci));
ci += 1;
keep
});
let keep_message = !(had_contents && message.contents.is_empty());
if keep_message {
origins.swap(kept, this);
kept += 1;
}
keep_message
});
origins.truncate(kept);
}
fn ensure_non_system_message(
original: &[Message],
mut retained: Vec<Message>,
origins: &[Option<usize>],
) -> Vec<Message> {
if retained
.iter()
.any(|m| m.role != Role::system() && message_renders_on_every_provider(m))
{
return retained;
}
let plain = original
.iter()
.enumerate()
.rev()
.filter(|(_, m)| m.role != Role::system() && m.role != Role::tool())
.find_map(|(index, m)| {
let mut candidate = m.clone();
let role = candidate.role.clone();
candidate.contents.retain(|c| {
renders_on_every_provider(c)
&& !matches!(c, Content::FunctionCall(_) | Content::FunctionResult(_))
&& !(matches!(c, Content::Data(_)) && role == Role::assistant())
});
(!candidate.contents.is_empty()).then_some((index, candidate))
});
if let Some((index, plain)) = plain {
let position = chronological_position(origins, index);
retained.insert(position, plain);
return retained;
}
if let Some((origin_index, exchange)) = latest_complete_tool_exchange(original) {
let position = chronological_position(origins, origin_index);
for (offset, message) in exchange.into_iter().enumerate() {
retained.insert(position + offset, message);
}
}
retained
}
type Site = (usize, usize);
type ExchangePair = (Site, Site);
type UnansweredCall<'a> = (Site, &'a str);
fn pair_tool_exchanges<'a>(
messages: impl Iterator<Item = (usize, &'a Message)>,
) -> (Vec<ExchangePair>, Vec<UnansweredCall<'a>>) {
use std::collections::{HashMap, VecDeque};
let mut unanswered: HashMap<&str, VecDeque<Site>> = HashMap::new();
let mut pairs: Vec<ExchangePair> = Vec::new();
for (mi, message) in messages {
for (ci, content) in message.contents.iter().enumerate() {
match content {
Content::FunctionCall(fc) if !fc.call_id.is_empty() => {
unanswered
.entry(fc.call_id.as_str())
.or_default()
.push_back((mi, ci));
}
Content::FunctionResult(fr) if !fr.call_id.is_empty() => {
if let Some(call_site) = unanswered
.get_mut(fr.call_id.as_str())
.and_then(VecDeque::pop_back)
{
pairs.push((call_site, (mi, ci)));
}
}
_ => {}
}
}
}
let mut leftover: Vec<UnansweredCall<'_>> = unanswered
.into_iter()
.flat_map(|(id, sites)| sites.into_iter().map(move |site| (site, id)))
.collect();
leftover.sort_unstable_by_key(|(site, _)| *site);
(pairs, leftover)
}
fn origin_indices(retained: &[Message], original: &[Message]) -> Vec<Option<usize>> {
let mut assigned: Vec<Option<usize>> = vec![None; retained.len()];
let mut limit = original.len();
for index in (0..retained.len()).rev() {
match original[..limit]
.iter()
.rposition(|candidate| *candidate == retained[index])
{
Some(found) => {
assigned[index] = Some(found);
limit = found;
}
None => break,
}
}
assigned
}
fn chronological_position(origins: &[Option<usize>], origin_index: usize) -> usize {
origins
.iter()
.position(|slot| matches!(slot, Some(index) if *index > origin_index))
.unwrap_or(origins.len())
}
fn latest_complete_tool_exchange(messages: &[Message]) -> Option<(usize, Vec<Message>)> {
let (pairs, _) = pair_tool_exchanges(
messages
.iter()
.enumerate()
.filter(|(_, m)| m.role != Role::system()),
);
let ((call_mi, call_ci), (result_mi, result_ci)) = pairs.last().copied()?;
let messages = if call_mi == result_mi {
vec![reduced_call_with_signature(
&messages[call_mi],
call_ci,
&[result_ci],
)]
} else {
vec![
reduced_call_with_signature(&messages[call_mi], call_ci, &[]),
reduced_to(&messages[result_mi], &[result_ci]),
]
};
Some((call_mi, messages))
}
fn renders_on_every_provider(content: &Content) -> bool {
content.renders_on_every_provider()
}
fn message_renders_on_every_provider(message: &Message) -> bool {
message.contents.iter().any(|content| {
content.renders_on_every_provider()
&& !(matches!(content, Content::Data(_)) && message.role == Role::assistant())
})
}
fn reduced_call_with_signature(message: &Message, call_ci: usize, also_keep: &[usize]) -> Message {
let mut keep = vec![call_ci];
keep.extend_from_slice(also_keep);
keep.sort_unstable();
let mut out = reduced_to(message, &keep);
if let Some(Content::FunctionCall(fc)) = out
.contents
.iter_mut()
.find(|c| matches!(c, Content::FunctionCall(_)))
{
if fc.protected_data.is_none() {
fc.protected_data = preceding_reasoning_signature(message, call_ci);
}
}
out
}
fn reduced_to(message: &Message, keep: &[usize]) -> Message {
let mut out = message.clone();
let mut index = 0;
out.contents.retain(|_| {
let keep_this = keep.contains(&index);
index += 1;
keep_this
});
out
}
fn reinstate_calls_answered_by_pending(
original: &[Message],
retained: &mut Vec<Message>,
origins: &mut Vec<Option<usize>>,
pending: &[Message],
) {
let mut wanted = unanswered_result_counts(pending);
if wanted.is_empty() {
return;
}
let (_, retained_unanswered) = pair_tool_exchanges(retained.iter().enumerate());
let mut supplied: Vec<&crate::types::FunctionCallContent> = retained_unanswered
.iter()
.filter_map(|((mi, ci), _)| retained[*mi].contents[*ci].as_function_call())
.collect();
let (_, original_unanswered) = pair_tool_exchanges(original.iter().enumerate());
let mut chosen: Vec<Site> = Vec::new();
for (site, call_id) in original_unanswered.into_iter().rev() {
let Some(call) = original[site.0].contents[site.1].as_function_call() else {
continue;
};
if let Some(position) = supplied.iter().position(|supplied| *supplied == call) {
supplied.remove(position);
continue;
}
let Some(count) = wanted.get_mut(call_id) else {
continue;
};
if *count == 0 {
continue;
}
*count -= 1;
chosen.push(site);
}
chosen.sort_unstable();
let mut index = 0;
while index < chosen.len() {
let (mi, _) = chosen[index];
let mut group = Vec::new();
while index < chosen.len() && chosen[index].0 == mi {
group.push(chosen[index].1);
index += 1;
}
let (first, rest) = group.split_first().expect("a group has at least one call");
retained.push(reduced_call_with_signature(&original[mi], *first, rest));
origins.push(Some(mi));
}
}
fn preceding_reasoning_signature(message: &Message, index: usize) -> Option<String> {
for content in message.contents[..index].iter().rev() {
match content {
Content::TextReasoning(t) => {
if let Some(signature) = t.protected_data.as_deref().filter(|s| !s.is_empty()) {
return Some(signature.to_string());
}
}
Content::FunctionCall(_) => {}
_ => return None,
}
}
None
}
fn unanswered_result_counts(messages: &[Message]) -> std::collections::HashMap<&str, usize> {
use std::collections::HashMap;
let mut available_calls: HashMap<&str, usize> = HashMap::new();
let mut unanswered: HashMap<&str, usize> = HashMap::new();
for content in messages.iter().flat_map(|m| m.contents.iter()) {
match content {
Content::FunctionCall(fc) if !fc.call_id.is_empty() => {
*available_calls.entry(fc.call_id.as_str()).or_insert(0) += 1;
}
Content::FunctionResult(fr) if !fr.call_id.is_empty() => {
match available_calls.get_mut(fr.call_id.as_str()) {
Some(count) if *count > 0 => *count -= 1,
_ => *unanswered.entry(fr.call_id.as_str()).or_insert(0) += 1,
}
}
_ => {}
}
}
unanswered.retain(|_, count| *count > 0);
unanswered
}
fn finalize_compaction(
original: &[Message],
retained: Vec<Message>,
pending: &[Message],
) -> Vec<Message> {
let mut origins = origin_indices(&retained, original);
let mut retained = retained;
reinstate_calls_answered_by_pending(original, &mut retained, &mut origins, pending);
drop_orphaned_tool_exchanges(&mut retained, &mut origins, pending);
if pending
.iter()
.any(|m| m.role != Role::system() && message_renders_on_every_provider(m))
{
return retained;
}
ensure_non_system_message(original, retained, &origins)
}
pub fn compact(
messages: &[Message],
strategy: &dyn CompactionStrategy,
tokenizer: &dyn Tokenizer,
) -> Vec<Message> {
finalize_compaction(messages, strategy.compact(messages, tokenizer), &[])
}
pub struct CompactionProvider {
strategy: Arc<dyn CompactionStrategy>,
tokenizer: Box<dyn Tokenizer>,
}
impl CompactionProvider {
pub fn new(strategy: impl CompactionStrategy + 'static) -> Self {
Self::with_tokenizer(strategy, ApproxTokenizer)
}
pub fn with_tokenizer(
strategy: impl CompactionStrategy + 'static,
tokenizer: impl Tokenizer + 'static,
) -> Self {
Self {
strategy: Arc::new(strategy),
tokenizer: Box::new(tokenizer),
}
}
}
#[async_trait]
impl ContextProvider for CompactionProvider {
async fn before_run(&self, ctx: &mut SessionContext) -> Result<()> {
let retained = self.strategy.compact(&ctx.messages, &*self.tokenizer);
ctx.messages = finalize_compaction(&ctx.messages, retained, &ctx.input_messages);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::FunctionResultContent;
use serde_json::json;
fn text(role: Role, s: &str) -> Message {
Message::new(role, s)
}
fn tool_result_message(call_id: &str, result: &str) -> Message {
Message::with_contents(
Role::tool(),
vec![Content::FunctionResult(FunctionResultContent::new(
call_id,
Some(json!(result)),
))],
)
}
fn tool_call_message(call_id: &str, name: &str) -> Message {
Message::with_contents(
Role::assistant(),
vec![Content::FunctionCall(
crate::types::FunctionCallContent::new(call_id, name, None),
)],
)
}
fn call_ids_in(messages: &[Message], f: fn(&Content) -> bool) -> Vec<String> {
messages
.iter()
.flat_map(|m| m.contents.iter())
.filter(|c| f(c))
.filter_map(|c| match c {
Content::FunctionCall(fc) => Some(fc.call_id.clone()),
Content::FunctionResult(fr) => Some(fr.call_id.clone()),
_ => None,
})
.collect()
}
#[test]
fn selective_tool_result_compacts_a_failed_calls_exception_too() {
let failed = Message::with_contents(
Role::tool(),
vec![Content::FunctionResult(FunctionResultContent {
call_id: "c1".into(),
result: None,
exception: Some("Traceback: ...a very long stack trace...".into()),
})],
);
let messages = vec![
tool_call_message("c1", "t1"),
failed,
tool_call_message("c2", "t2"),
tool_result_message("c2", "fresh"),
];
let out = compact(&messages, &SelectiveToolResult::new(1), &ApproxTokenizer);
let compacted = &out[1].function_results()[0];
assert_eq!(compacted.exception.as_deref(), Some(OMITTED_TOOL_RESULT));
assert_eq!(out[3].function_results()[0].result, Some(json!("fresh")));
}
#[test]
fn selective_tool_result_does_not_invent_an_exception_on_a_successful_result() {
let messages = vec![
tool_call_message("c1", "t1"),
tool_result_message("c1", "stale"),
tool_call_message("c2", "t2"),
tool_result_message("c2", "fresh"),
];
let out = compact(&messages, &SelectiveToolResult::new(1), &ApproxTokenizer);
assert!(out[1].function_results()[0].exception.is_none());
}
#[test]
fn selective_tool_result_compacts_the_payload_without_orphaning_the_call() {
let messages = vec![
tool_call_message("call_1", "get_weather"),
tool_result_message("call_1", "a very long stale payload"),
tool_call_message("call_2", "get_time"),
tool_result_message("call_2", "noon"),
];
let out = compact(&messages, &SelectiveToolResult::new(1), &ApproxTokenizer);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["call_1", "call_2"]
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))),
vec!["call_1", "call_2"]
);
assert_eq!(
out[1].function_results()[0].result,
Some(json!(OMITTED_TOOL_RESULT))
);
assert_eq!(out[3].function_results()[0].result, Some(json!("noon")));
}
#[test]
fn a_budget_cut_between_call_and_result_orphans_neither() {
let mut call_msg = tool_call_message("call_1", "get_weather");
call_msg
.contents
.push(Content::text("let me look that up for you right now"));
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "hi"),
call_msg,
tool_result_message("call_1", "sunny"),
];
let out = compact(&messages, &TokenBudget::new(4), &ApproxTokenizer);
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))).is_empty());
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn a_reused_call_id_is_paired_by_occurrence_not_by_id() {
let messages = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
tool_call_message("c1", "get_weather"),
];
let out = compact(&messages, &SlidingWindow::new(10), &ApproxTokenizer);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"],
"the second, unanswered c1 call must be dropped"
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))),
vec!["c1"]
);
}
#[test]
fn two_full_exchanges_reusing_one_call_id_both_survive() {
let messages = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "rainy"),
];
let out = compact(&messages, &SlidingWindow::new(10), &ApproxTokenizer);
assert_eq!(out.len(), 4);
}
#[test]
fn a_result_preceding_its_call_is_not_treated_as_paired() {
let messages = vec![
tool_result_message("c1", "sunny"),
tool_call_message("c1", "get_weather"),
];
let out = compact(&messages, &SlidingWindow::new(10), &ApproxTokenizer);
assert!(out.is_empty(), "both halves are orphans, got {out:?}");
}
#[test]
fn an_intact_tool_exchange_is_left_alone() {
let messages = vec![
tool_call_message("call_1", "get_weather"),
tool_result_message("call_1", "sunny"),
];
let out = compact(&messages, &SlidingWindow::new(10), &ApproxTokenizer);
assert_eq!(out.len(), 2);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["call_1"]
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))),
vec!["call_1"]
);
}
#[test]
fn an_orphan_strip_keeps_the_rest_of_its_message() {
let mut call_msg = tool_call_message("call_1", "get_weather");
call_msg.contents.push(Content::text("checking now"));
let out = compact(&[call_msg], &SlidingWindow::new(10), &ApproxTokenizer);
assert_eq!(out.len(), 1);
assert_eq!(out[0].text(), "checking now");
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn compaction_never_returns_system_messages_only() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "hello"),
text(Role::assistant(), "hi"),
];
for out in [
compact(&messages, &Truncation::new(1), &ApproxTokenizer),
compact(&messages, &SlidingWindow::new(0), &ApproxTokenizer),
] {
assert!(
out.iter().any(|m| m.role != Role::system()),
"expected a non-system message to be retained, got {:?}",
out.iter().map(|m| m.role.as_str()).collect::<Vec<_>>()
);
assert_eq!(out.last().unwrap().text(), "hi");
}
}
#[test]
fn the_minimum_retained_message_is_never_a_half_exchange() {
let mut call_msg = tool_call_message("call_1", "get_weather");
call_msg.contents.push(Content::text("checking now"));
let messages = vec![
text(Role::system(), "sys"),
call_msg,
tool_result_message("call_1", "sunny"),
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert!(out.iter().any(|m| m.role != Role::system()));
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))).is_empty());
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn a_call_answered_by_the_runs_input_is_not_dropped_as_an_orphan() {
let history = vec![
text(Role::user(), "what's the weather?"),
tool_call_message("c1", "get_weather"),
];
let input = vec![tool_result_message("c1", "sunny")];
let retained = SlidingWindow::new(10).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"],
"the call must survive: its result arrives in this run's input"
);
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))).is_empty());
let retained = SlidingWindow::new(10).compact(&history, &ApproxTokenizer);
let blind = finalize_compaction(&history, retained, &[]);
assert!(call_ids_in(&blind, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn a_call_the_strategy_dropped_is_reinstated_when_the_input_answers_it() {
let history = vec![
text(Role::user(), "what's the weather?"),
tool_call_message("c1", "get_weather"),
];
let input = vec![tool_result_message("c1", "sunny")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
assert!(
call_ids_in(&retained, |c| matches!(c, Content::FunctionCall(_))).is_empty(),
"precondition: the strategy really did drop the call"
);
let out = finalize_compaction(&history, retained, &input);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"],
"the call must be put back so the incoming result answers something"
);
}
#[test]
fn a_reused_call_id_reinstates_the_unanswered_occurrence() {
let mut old_call = tool_call_message("c1", "get_weather");
old_call.contents = vec![Content::FunctionCall(
crate::types::FunctionCallContent::new(
"c1",
"get_weather",
Some(crate::types::FunctionArguments::Raw(
"{\"city\":\"old\"}".into(),
)),
),
)];
let mut new_call = tool_call_message("c1", "get_weather");
new_call.contents = vec![Content::FunctionCall(
crate::types::FunctionCallContent::new(
"c1",
"get_weather",
Some(crate::types::FunctionArguments::Raw(
"{\"city\":\"new\"}".into(),
)),
),
)];
let history = vec![old_call, tool_result_message("c1", "old answer"), new_call];
let input = vec![tool_result_message("c1", "new answer")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
let calls: Vec<&crate::types::FunctionCallContent> = out
.iter()
.flat_map(|m| m.contents.iter())
.filter_map(Content::as_function_call)
.collect();
assert_eq!(calls.len(), 1);
let args = calls[0].parse_arguments().unwrap();
assert_eq!(
args.get("city").and_then(|v| v.as_str()),
Some("new"),
"the unanswered (new) occurrence must be the one reinstated"
);
}
#[test]
fn a_self_contained_pending_exchange_reinstates_nothing() {
let history = vec![tool_call_message("c1", "get_weather")];
let input = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty(),
"pending answers itself; nothing should be reinstated, got {out:?}"
);
}
#[test]
fn a_pending_result_beyond_what_pending_answers_still_reinstates() {
let history = vec![tool_call_message("c1", "get_weather")];
let input = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "first"),
tool_result_message("c1", "second"),
];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"]
);
}
#[test]
fn a_pending_exchange_pairs_with_its_own_call_not_a_historical_one() {
let history = vec![tool_call_message("c1", "get_weather")];
let input = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
];
let retained = SlidingWindow::new(10).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty(),
"the unanswered historical call must be stripped, got {out:?}"
);
}
#[test]
fn a_reinstated_call_keeps_its_reasoning_signature() {
let signed = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
protected_data: Some("c2ln".into()),
..Default::default()
}),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c1",
"get_weather",
None,
)),
],
);
let history = vec![signed];
let input = vec![tool_result_message("c1", "sunny")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
let call = out
.iter()
.flat_map(|m| m.contents.iter())
.find_map(Content::as_function_call)
.expect("the call is reinstated");
assert_eq!(call.protected_data.as_deref(), Some("c2ln"));
}
#[test]
fn a_reinstated_call_gains_no_signature_from_unrelated_content() {
let mut msg = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
protected_data: Some("c2ln".into()),
..Default::default()
}),
Content::text("intervening"),
],
);
msg.contents.push(Content::FunctionCall(
crate::types::FunctionCallContent::new("c1", "get_weather", None),
));
let history = vec![msg];
let input = vec![tool_result_message("c1", "sunny")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
let call = out
.iter()
.flat_map(|m| m.contents.iter())
.find_map(Content::as_function_call)
.expect("the call is reinstated");
assert!(call.protected_data.is_none());
}
#[test]
fn the_fallback_restores_the_nearest_reused_call_not_the_oldest() {
let mk = |city: &str| {
Message::with_contents(
Role::assistant(),
vec![Content::FunctionCall(
crate::types::FunctionCallContent::new(
"c1",
"get_weather",
Some(crate::types::FunctionArguments::Raw(format!(
"{{\"city\":\"{city}\"}}"
))),
),
)],
)
};
let messages = vec![
text(Role::system(), "sys"),
mk("old"),
mk("new"),
tool_result_message("c1", "sunny"),
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
let call = out
.iter()
.flat_map(|m| m.contents.iter())
.find_map(Content::as_function_call)
.expect("the exchange is restored");
assert_eq!(
call.parse_arguments()
.unwrap()
.get("city")
.and_then(|v| v.as_str()),
Some("new")
);
}
#[test]
fn unknown_content_does_not_satisfy_the_retention_check() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "a real earlier turn"),
Message::with_contents(Role::assistant(), vec![Content::Unknown]),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter()
.any(|m| m.contents.iter().any(|c| matches!(c, Content::Text(_)))
&& m.role != Role::system()),
"expected a real turn to be restored, got {out:?}"
);
}
struct KeepIndices(Vec<usize>);
impl CompactionStrategy for KeepIndices {
fn compact(&self, messages: &[Message], _t: &dyn Tokenizer) -> Vec<Message> {
messages
.iter()
.enumerate()
.filter(|(i, _)| self.0.contains(i))
.map(|(_, m)| m.clone())
.collect()
}
}
#[test]
fn an_older_retained_occurrence_does_not_suppress_reinstating_the_newer() {
let mk = |city: &str| {
Message::with_contents(
Role::assistant(),
vec![Content::FunctionCall(
crate::types::FunctionCallContent::new(
"c1",
"get_weather",
Some(crate::types::FunctionArguments::Raw(format!(
"{{\"city\":\"{city}\"}}"
))),
),
)],
)
};
let history = vec![mk("old"), mk("new")];
let input = vec![tool_result_message("c1", "sunny")];
let retained = KeepIndices(vec![0]).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
let cities: Vec<String> = out
.iter()
.flat_map(|m| m.contents.iter())
.filter_map(Content::as_function_call)
.filter_map(|c| {
c.parse_arguments()
.ok()?
.get("city")?
.as_str()
.map(str::to_string)
})
.collect();
assert!(
cities.contains(&"new".to_string()),
"the newer occurrence the result answers must be present, got {cities:?}"
);
}
#[test]
fn non_image_media_alone_does_not_satisfy_the_retention_check() {
for media in [
Content::Data(crate::types::DataContent::from_bytes(b"aud", "audio/wav")),
Content::Uri(crate::types::UriContent {
uri: "https://example.com/doc.pdf".into(),
media_type: "application/pdf".into(),
}),
] {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "a real earlier turn"),
Message::with_contents(Role::assistant(), vec![media]),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected a real turn to be restored, got {out:?}"
);
}
}
#[test]
fn a_hosted_file_alone_does_not_satisfy_the_retention_check() {
let hosted = Message::with_contents(
Role::assistant(),
vec![Content::HostedFile(crate::types::HostedFileContent {
file_id: "file_1".into(),
})],
);
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "a real earlier turn"),
hosted,
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected a real turn to be restored, got {out:?}"
);
}
#[test]
fn a_completed_exchange_in_history_is_not_reinstated() {
let history = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
];
let input = vec![tool_result_message("c2", "noon")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn a_reinstated_call_carries_nothing_else_from_its_message() {
let mut call_msg = tool_call_message("c1", "get_weather");
call_msg
.contents
.push(Content::text("a large unrelated payload"));
let history = vec![call_msg];
let input = vec![tool_result_message("c1", "sunny")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert_eq!(out.len(), 1);
assert_eq!(out[0].contents.len(), 1, "only the call rides back in");
assert!(out[0].text().is_empty());
}
#[test]
fn nothing_is_reinstated_when_the_input_carries_no_results() {
let history = vec![tool_call_message("c1", "get_weather")];
let input = vec![text(Role::user(), "hello")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn an_already_retained_call_is_not_reinstated_twice() {
let history = vec![tool_call_message("c1", "get_weather")];
let input = vec![tool_result_message("c1", "sunny")];
let retained = SlidingWindow::new(10).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"]
);
}
#[test]
fn the_minimum_retention_fallback_never_picks_a_tool_role_message() {
let mut result_msg = tool_result_message("c1", "sunny");
result_msg.contents.push(Content::text("explanatory text"));
let messages = vec![
text(Role::system(), "sys"),
tool_call_message("c1", "get_weather"),
result_msg,
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"]
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))),
vec!["c1"]
);
for m in &out {
if m.role == Role::tool() {
assert!(
m.contents
.iter()
.any(|c| matches!(c, Content::FunctionResult(_))),
"a tool-role message must carry its result"
);
}
}
}
#[test]
fn a_call_with_no_answer_anywhere_is_still_dropped() {
let history = vec![tool_call_message("c1", "get_weather")];
let input = vec![text(Role::user(), "never mind")];
let retained = SlidingWindow::new(10).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))).is_empty());
}
#[test]
fn pending_input_satisfies_the_minimum_retention_rule() {
let history = vec![text(Role::system(), "sys"), text(Role::user(), "older")];
let input = vec![text(Role::user(), "current question")];
let retained = Truncation::new(1).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(out.iter().all(|m| m.role == Role::system()));
}
#[test]
fn a_conversation_of_only_a_tool_exchange_retains_it_whole() {
let messages = vec![
text(Role::system(), "sys"),
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
];
for out in [
compact(&messages, &Truncation::new(1), &ApproxTokenizer),
compact(&messages, &SlidingWindow::new(0), &ApproxTokenizer),
] {
assert!(
out.iter().any(|m| m.role != Role::system()),
"expected the tool exchange to be retained, got {:?}",
out.iter().map(|m| m.role.as_str()).collect::<Vec<_>>()
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"]
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))),
vec!["c1"]
);
}
}
#[test]
fn the_latest_complete_exchange_is_the_one_reinstated() {
let messages = vec![
text(Role::system(), "sys"),
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
tool_call_message("c2", "get_time"),
tool_result_message("c2", "noon"),
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c2"]
);
}
#[test]
fn an_incomplete_tool_exchange_is_not_reinstated() {
let messages = vec![
text(Role::system(), "sys"),
tool_call_message("c1", "get_weather"),
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert!(out.iter().all(|m| m.role == Role::system()));
}
#[test]
fn reasoning_alone_does_not_qualify_as_the_minimum_retained_message() {
let signed = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
..Default::default()
}),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c1",
"get_weather",
None,
)),
],
);
let messages = vec![
text(Role::system(), "sys"),
signed,
tool_result_message("c1", "sunny"),
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionCall(_))),
vec!["c1"],
"the complete exchange must be restored, got {out:?}"
);
assert_eq!(
call_ids_in(&out, |c| matches!(c, Content::FunctionResult(_))),
vec!["c1"]
);
}
#[test]
fn text_beside_reasoning_still_qualifies() {
let msg = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
..Default::default()
}),
Content::text("the answer"),
],
);
let messages = vec![text(Role::system(), "sys"), msg];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert_eq!(out.last().unwrap().text(), "the answer");
}
#[test]
fn the_complete_exchange_fallback_carries_the_reasoning_signature() {
let signed = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
protected_data: Some("c2ln".into()),
..Default::default()
}),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c1",
"get_weather",
None,
)),
],
);
let messages = vec![
text(Role::system(), "sys"),
signed,
tool_result_message("c1", "sunny"),
];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
let call = out
.iter()
.flat_map(|m| m.contents.iter())
.find_map(Content::as_function_call)
.expect("the exchange is restored");
assert_eq!(call.protected_data.as_deref(), Some("c2ln"));
}
#[test]
fn a_retained_message_that_renders_nothing_does_not_satisfy_retention() {
let reasoning_plus_orphan = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
..Default::default()
}),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"orphan",
"get_weather",
None,
)),
],
);
let messages = vec![
text(Role::system(), "sys"),
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
reasoning_plus_orphan,
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents
.iter()
.any(|c| !matches!(c, Content::TextReasoning(_)))),
"expected something that actually renders, got {out:?}"
);
}
#[test]
fn retained_reasoning_is_never_discarded_by_the_retention_check() {
let reasoning_only = Message::with_contents(
Role::assistant(),
vec![Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
raw_representation: Some(json!({"id": "rs_1"})),
..Default::default()
})],
);
let messages = vec![text(Role::system(), "sys"), reasoning_only];
let out = compact(&messages, &SlidingWindow::new(10), &ApproxTokenizer);
let kept = out
.iter()
.flat_map(|m| m.contents.iter())
.any(|c| matches!(c, Content::TextReasoning(_)));
assert!(kept, "reasoning must survive, got {out:?}");
}
#[test]
fn the_fallback_is_inserted_in_chronological_order() {
let messages = vec![
text(Role::user(), "old question"),
Message::with_contents(
Role::assistant(),
vec![Content::TextReasoning(crate::types::TextReasoningContent {
text: "newer thinking".into(),
..Default::default()
})],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert_eq!(out.len(), 2);
assert_eq!(out[0].text(), "old question", "got {out:?}");
assert!(matches!(out[1].contents[0], Content::TextReasoning(_)));
}
#[test]
fn a_fallback_newer_than_the_retained_message_is_appended() {
let messages = vec![
Message::with_contents(
Role::assistant(),
vec![Content::TextReasoning(crate::types::TextReasoningContent {
text: "older thinking".into(),
..Default::default()
})],
),
text(Role::user(), "newer question"),
];
let retained = KeepIndices(vec![0]).compact(&messages, &ApproxTokenizer);
let out = finalize_compaction(&messages, retained, &[]);
assert_eq!(out.len(), 2);
assert!(matches!(out[0].contents[0], Content::TextReasoning(_)));
assert_eq!(out[1].text(), "newer question");
}
#[test]
fn wire_empty_pending_input_does_not_satisfy_retention() {
let history = vec![text(Role::system(), "sys"), text(Role::user(), "real turn")];
let input = vec![Message::with_contents(
Role::user(),
vec![Content::HostedFile(crate::types::HostedFileContent {
file_id: "file_1".into(),
})],
)];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected the real turn to be restored, got {out:?}"
);
}
#[test]
fn tool_payloads_count_against_the_token_budget() {
let t = ApproxTokenizer;
let bulky = Message::with_contents(
Role::tool(),
vec![Content::FunctionResult(FunctionResultContent::new(
"c1",
Some(json!("x".repeat(400))),
))],
);
assert!(
count_message_tokens(&t, &bulky) >= 100,
"a large tool result must cost tokens, got {}",
count_message_tokens(&t, &bulky)
);
let call = Message::with_contents(
Role::assistant(),
vec![Content::FunctionCall(
crate::types::FunctionCallContent::new(
"c1",
"get_weather",
Some(crate::types::FunctionArguments::Raw("y".repeat(400))),
),
)],
);
assert!(count_message_tokens(&t, &call) >= 100);
let messages = vec![call.clone(), bulky.clone(), text(Role::user(), "tiny")];
let out = TokenBudget::new(10).compact(&messages, &t);
assert!(
out.len() < 3,
"the bulky exchange must not fit a 10-token budget, got {out:?}"
);
}
#[test]
fn duplicate_messages_do_not_alias_the_fallback_position() {
let reasoning = || {
Message::with_contents(
Role::assistant(),
vec![Content::TextReasoning(crate::types::TextReasoningContent {
text: "same thinking".into(),
..Default::default()
})],
)
};
let messages = vec![reasoning(), text(Role::user(), "old question"), reasoning()];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert_eq!(out.len(), 2);
assert_eq!(
out[0].text(),
"old question",
"the fallback must precede the retained later reasoning, got {out:?}"
);
}
#[test]
fn the_complete_exchange_fallback_is_also_inserted_chronologically() {
let messages = vec![
tool_call_message("c1", "get_weather"),
tool_result_message("c1", "sunny"),
Message::with_contents(
Role::assistant(),
vec![Content::TextReasoning(crate::types::TextReasoningContent {
text: "newer thinking".into(),
..Default::default()
})],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
let reasoning_at = out
.iter()
.position(|m| {
m.contents
.iter()
.any(|c| matches!(c, Content::TextReasoning(_)))
})
.expect("the retained reasoning is still present");
let call_at = out
.iter()
.position(|m| {
m.contents
.iter()
.any(|c| matches!(c, Content::FunctionCall(_)))
})
.expect("the exchange is restored");
assert!(
call_at < reasoning_at,
"the older exchange must precede the newer reasoning, got {out:?}"
);
let result_at = out
.iter()
.position(|m| {
m.contents
.iter()
.any(|c| matches!(c, Content::FunctionResult(_)))
})
.expect("the result is restored");
assert!(
call_at < result_at && result_at < reasoning_at,
"got {out:?}"
);
}
#[test]
fn parallel_reinstated_calls_stay_in_one_message() {
let parallel = Message::with_contents(
Role::assistant(),
vec![
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c1",
"get_weather",
None,
)),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c2", "get_time", None,
)),
],
);
let history = vec![parallel];
let input = vec![
tool_result_message("c1", "sunny"),
tool_result_message("c2", "noon"),
];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
let with_calls: Vec<&Message> = out
.iter()
.filter(|m| {
m.contents
.iter()
.any(|c| matches!(c, Content::FunctionCall(_)))
})
.collect();
assert_eq!(
with_calls.len(),
1,
"both calls must be reinstated in one message, got {out:?}"
);
assert_eq!(with_calls[0].contents.len(), 2);
}
#[test]
fn chronology_survives_the_orphan_repair_rewriting_a_message() {
let newer = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "newer thinking".into(),
..Default::default()
}),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"orphan",
"get_weather",
None,
)),
],
);
let messages = vec![text(Role::user(), "old question"), newer];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert_eq!(out.len(), 2);
assert_eq!(
out[0].text(),
"old question",
"the fallback must precede the retained (rewritten) message, got {out:?}"
);
}
#[test]
fn an_assistant_image_turn_does_not_satisfy_retention() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "a real earlier turn"),
Message::with_contents(
Role::assistant(),
vec![Content::Data(crate::types::DataContent::from_bytes(
b"img",
"image/png",
))],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected the real turn restored, got {out:?}"
);
}
#[test]
fn an_image_turn_satisfies_the_retention_check() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "an older turn"),
Message::with_contents(
Role::user(),
vec![Content::Data(crate::types::DataContent::from_bytes(
b"img",
"image/png",
))],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert_eq!(out.len(), 2, "no fallback should be added, got {out:?}");
assert!(out[1]
.contents
.iter()
.any(|c| matches!(c, Content::Data(_))));
}
#[test]
fn a_remote_image_uri_does_not_satisfy_retention() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "an older turn"),
Message::with_contents(
Role::user(),
vec![Content::Uri(crate::types::UriContent {
uri: "https://example.com/a.png".into(),
media_type: "image/png".into(),
})],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected the real turn restored, got {out:?}"
);
}
#[test]
fn an_image_format_outside_the_universal_set_does_not_satisfy_retention() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "an older turn"),
Message::with_contents(
Role::user(),
vec![Content::Data(crate::types::DataContent::from_bytes(
b"<svg/>",
"image/svg+xml",
))],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected the real turn restored, got {out:?}"
);
}
#[test]
fn a_reinstated_parallel_call_inherits_the_groups_reasoning_signature() {
let group = Message::with_contents(
Role::assistant(),
vec![
Content::TextReasoning(crate::types::TextReasoningContent {
text: "thinking".into(),
protected_data: Some("c2ln".into()),
..Default::default()
}),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c1",
"get_weather",
None,
)),
Content::FunctionCall(crate::types::FunctionCallContent::new(
"c2", "get_time", None,
)),
],
);
let history = vec![group];
let input = vec![tool_result_message("c2", "noon")];
let retained = SlidingWindow::new(0).compact(&history, &ApproxTokenizer);
let out = finalize_compaction(&history, retained, &input);
let call = out
.iter()
.flat_map(|m| m.contents.iter())
.find_map(Content::as_function_call)
.expect("c2 is reinstated");
assert_eq!(call.call_id, "c2");
assert_eq!(call.protected_data.as_deref(), Some("c2ln"));
}
#[test]
fn an_image_type_on_a_malformed_uri_does_not_satisfy_retention() {
for uri in ["not-a-data-uri", "data:image/png,AAAA", ""] {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "a real earlier turn"),
Message::with_contents(
Role::user(),
vec![Content::Data(crate::types::DataContent {
uri: uri.into(),
media_type: Some("image/png".into()),
})],
),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter().any(|m| m.role != Role::system()
&& m.contents.iter().any(|c| matches!(c, Content::Text(_)))),
"expected a real turn restored for {uri:?}, got {out:?}"
);
}
}
#[test]
fn blank_text_does_not_satisfy_the_retention_check() {
for blank in ["", " \n\t"] {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "a real earlier turn"),
text(Role::assistant(), blank),
];
let out = compact(&messages, &SlidingWindow::new(1), &ApproxTokenizer);
assert!(
out.iter()
.any(|m| m.role != Role::system() && !m.text().trim().is_empty()),
"expected the real turn restored for {blank:?}, got {out:?}"
);
}
}
#[test]
fn an_all_system_conversation_stays_all_system() {
let messages = vec![text(Role::system(), "a"), text(Role::system(), "b")];
let out = compact(&messages, &Truncation::new(1), &ApproxTokenizer);
assert!(out.iter().all(|m| m.role == Role::system()));
}
#[test]
fn approx_tokenizer_uses_four_chars_per_token_ceiling() {
let t = ApproxTokenizer;
assert_eq!(t.count_tokens(""), 0);
assert_eq!(t.count_tokens("abcd"), 1);
assert_eq!(t.count_tokens("abcde"), 2); assert_eq!(t.count_tokens("abcdefgh"), 2);
assert_eq!(t.count_tokens("abcdefghi"), 3); }
#[test]
fn approx_tokenizer_does_not_inflate_non_ascii_text() {
let t = ApproxTokenizer;
assert_eq!(t.count_tokens("日本語訳"), t.count_tokens("abcd"));
let msg = Message::with_contents(Role::user(), vec![Content::text("日本語訳")]);
assert_eq!(count_message_tokens(&t, &msg), 1);
}
#[test]
fn count_message_tokens_sums_text_content() {
let t = ApproxTokenizer;
let msg = Message::with_contents(
Role::user(),
vec![Content::text("abcd"), Content::text("abcdefgh")],
);
assert_eq!(count_message_tokens(&t, &msg), 3);
}
#[test]
fn truncation_keeps_most_recent_messages() {
let messages = vec![
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
text(Role::assistant(), "4"),
];
let strategy = Truncation::new(2);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 2);
assert_eq!(out[0].text(), "3");
assert_eq!(out[1].text(), "4");
}
#[test]
fn truncation_preserves_leading_system_messages() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
text(Role::assistant(), "4"),
];
let strategy = Truncation::new(2);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 2);
assert_eq!(out[0].role, Role::system());
assert_eq!(out[0].text(), "sys");
assert_eq!(out[1].text(), "4");
}
#[test]
fn truncation_preserves_multiple_leading_system_messages() {
let messages = vec![
text(Role::system(), "sys1"),
text(Role::system(), "sys2"),
text(Role::user(), "1"),
text(Role::assistant(), "2"),
];
let strategy = Truncation::new(3);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 3);
assert_eq!(out[0].text(), "sys1");
assert_eq!(out[1].text(), "sys2");
assert_eq!(out[2].text(), "2");
}
#[test]
fn truncation_noop_when_under_budget() {
let messages = vec![text(Role::user(), "1"), text(Role::assistant(), "2")];
let strategy = Truncation::new(10);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out, messages);
}
#[test]
fn sliding_window_keeps_system_plus_last_n_non_system() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
];
let strategy = SlidingWindow::new(2);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 3);
assert_eq!(out[0].text(), "sys");
assert_eq!(out[1].text(), "2");
assert_eq!(out[2].text(), "3");
}
#[test]
fn sliding_window_with_no_system_message() {
let messages = vec![
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
];
let strategy = SlidingWindow::new(1);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 1);
assert_eq!(out[0].text(), "3");
}
struct FixedTokenizer(usize);
impl Tokenizer for FixedTokenizer {
fn count_tokens(&self, _text: &str) -> usize {
self.0
}
}
#[test]
fn token_budget_keeps_only_what_fits_from_the_newest_backward() {
let messages = vec![
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
text(Role::assistant(), "4"),
];
let tokenizer = FixedTokenizer(10);
let strategy = TokenBudget::new(25);
let out = compact(&messages, &strategy, &tokenizer);
assert_eq!(out.len(), 2);
assert_eq!(out[0].text(), "3");
assert_eq!(out[1].text(), "4");
}
#[test]
fn token_budget_preserves_leading_system_message_and_counts_it() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
];
let tokenizer = FixedTokenizer(10);
let strategy = TokenBudget::new(20);
let out = compact(&messages, &strategy, &tokenizer);
assert_eq!(out.len(), 2);
assert_eq!(out[0].role, Role::system());
assert_eq!(out[1].text(), "3");
}
#[test]
fn token_budget_keeps_at_least_the_newest_message_even_if_it_alone_exceeds_budget() {
let messages = vec![text(Role::user(), "1"), text(Role::assistant(), "2")];
let tokenizer = FixedTokenizer(100);
let strategy = TokenBudget::new(1);
let out = compact(&messages, &strategy, &tokenizer);
assert_eq!(out.len(), 1);
assert_eq!(out[0].text(), "2");
}
#[test]
fn token_budget_keeps_everything_when_it_all_fits() {
let messages = vec![text(Role::user(), "1"), text(Role::assistant(), "2")];
let tokenizer = FixedTokenizer(1);
let strategy = TokenBudget::new(1000);
let out = compact(&messages, &strategy, &tokenizer);
assert_eq!(out, messages);
}
#[test]
fn selective_tool_result_strips_stale_results_and_keeps_recent_ones() {
let messages = vec![
tool_call_message("c1", "t1"),
tool_result_message("c1", "result 1"),
tool_call_message("c2", "t2"),
tool_result_message("c2", "result 2"),
tool_call_message("c3", "t3"),
tool_result_message("c3", "result 3"),
];
let strategy = SelectiveToolResult::new(1);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 6);
assert_eq!(
out[1].function_results()[0].result,
Some(json!(OMITTED_TOOL_RESULT))
);
assert_eq!(
out[3].function_results()[0].result,
Some(json!(OMITTED_TOOL_RESULT))
);
assert_eq!(out[5].function_results()[0].result, Some(json!("result 3")));
}
#[test]
fn selective_tool_result_keeps_text_alongside_a_stripped_tool_result() {
let mixed = Message::with_contents(
Role::tool(),
vec![
Content::text("some accompanying text"),
Content::FunctionResult(FunctionResultContent::new("c1", Some(json!("r1")))),
],
);
let messages = vec![
tool_call_message("c1", "t1"),
mixed,
tool_call_message("c2", "t2"),
tool_result_message("c2", "result 2"),
tool_call_message("c3", "t3"),
tool_result_message("c3", "result 3"),
];
let strategy = SelectiveToolResult::new(2);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out.len(), 6);
assert_eq!(out[1].text(), "some accompanying text");
assert_eq!(
out[1].function_results()[0].result,
Some(json!(OMITTED_TOOL_RESULT))
);
assert_eq!(out[3].function_results()[0].result, Some(json!("result 2")));
assert_eq!(out[5].function_results()[0].result, Some(json!("result 3")));
}
#[test]
fn selective_tool_result_noop_when_keep_last_covers_all() {
let messages = vec![
tool_call_message("c1", "t1"),
tool_result_message("c1", "result 1"),
tool_call_message("c2", "t2"),
tool_result_message("c2", "result 2"),
];
let strategy = SelectiveToolResult::new(5);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out, messages);
}
#[test]
fn selective_tool_result_ignores_messages_without_tool_results() {
let messages = vec![
text(Role::system(), "sys"),
text(Role::user(), "hi"),
text(Role::assistant(), "hello"),
];
let strategy = SelectiveToolResult::new(0);
let out = compact(&messages, &strategy, &ApproxTokenizer);
assert_eq!(out, messages);
}
#[tokio::test]
async fn compaction_provider_before_run_replaces_ctx_messages_with_compacted_subset() {
let provider = CompactionProvider::new(Truncation::new(2));
let mut ctx = SessionContext::new(vec![]);
ctx.messages = vec![
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
text(Role::assistant(), "4"),
];
provider.before_run(&mut ctx).await.unwrap();
assert_eq!(ctx.messages.len(), 2);
assert_eq!(ctx.messages[0].text(), "3");
assert_eq!(ctx.messages[1].text(), "4");
}
#[tokio::test]
async fn compaction_provider_with_tokenizer_uses_the_supplied_tokenizer() {
struct FixedTokenizer(usize);
impl Tokenizer for FixedTokenizer {
fn count_tokens(&self, _text: &str) -> usize {
self.0
}
}
let provider = CompactionProvider::with_tokenizer(TokenBudget::new(25), FixedTokenizer(10));
let mut ctx = SessionContext::new(vec![]);
ctx.messages = vec![
text(Role::user(), "1"),
text(Role::assistant(), "2"),
text(Role::user(), "3"),
text(Role::assistant(), "4"),
];
provider.before_run(&mut ctx).await.unwrap();
assert_eq!(ctx.messages.len(), 2);
assert_eq!(ctx.messages[0].text(), "3");
assert_eq!(ctx.messages[1].text(), "4");
}
#[tokio::test]
async fn compaction_provider_after_run_is_a_noop() {
let provider = CompactionProvider::new(Truncation::new(1));
provider
.after_run(&[Message::new(Role::user(), "hi")], &[], None)
.await
.unwrap();
}
}