use std::collections::BTreeSet;
use runtime_kernel_api::{ContextLimits, KernelFailure};
use runtime_ports::{ModelContent, ModelMessage, ModelRole, ModelToolDefinition};
use crate::{
checkpoint::compile_checkpoint,
context_measurement::{ContextMeasurement, measure_context},
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum CompactionReason {
HighWatermark,
Emergency,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct CompactionDecision {
pub(crate) reason: CompactionReason,
pub(crate) before: ContextMeasurement,
pub(crate) target_tokens: u64,
pub(crate) reclaimable_tokens: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct CompactionOutcome {
pub(crate) before: ContextMeasurement,
pub(crate) after: ContextMeasurement,
pub(crate) removed_messages: usize,
pub(crate) used_fallback: bool,
}
#[derive(Default)]
pub(crate) struct CompactionState {
generation: usize,
last_turn: Option<usize>,
}
pub(crate) struct CompactionController {
limits: ContextLimits,
state: CompactionState,
}
impl CompactionController {
pub(crate) fn new(limits: ContextLimits) -> Self {
Self {
limits,
state: CompactionState::default(),
}
}
pub(crate) fn measurement(
&self,
messages: &[ModelMessage],
tools: &[ModelToolDefinition],
) -> ContextMeasurement {
measure_context(
messages,
tools,
self.limits.protocol_overhead_tokens,
self.limits.reserved_output_tokens,
)
}
pub(crate) fn decide(
&self,
messages: &[ModelMessage],
tools: &[ModelToolDefinition],
current_user_prompt: &str,
turn: usize,
) -> Result<Option<CompactionDecision>, KernelFailure> {
let before = self.measurement(messages, tools);
let resident = resident_messages(messages, current_user_prompt);
let pinned = self.measurement(&resident, tools).total_tokens;
let reclaimable_tokens = before.total_tokens.saturating_sub(pinned);
let hard = self.limits.max_input_tokens;
let trigger = percent(hard, self.limits.compaction.trigger_percent);
let target = percent(hard, self.limits.compaction.target_percent);
let emergency = percent(hard, self.limits.compaction.emergency_percent);
let minimum_reclaim = percent(hard, self.limits.compaction.min_reclaim_percent);
if before.total_tokens >= hard && reclaimable_tokens < minimum_reclaim {
return Err(KernelFailure::new(
"CONTEXT_REQUIRED_TOO_LARGE",
format!(
"model context requires {} tokens but only {} are reclaimable within the {} token limit",
before.total_tokens, reclaimable_tokens, hard
),
));
}
if !self.limits.auto_compact {
if before.total_tokens > hard {
return Err(KernelFailure::new(
"CONTEXT_LIMIT_EXCEEDED",
"model context exceeds max_input_tokens and auto compaction is disabled",
));
}
return Ok(None);
}
if reclaimable_tokens < minimum_reclaim {
return Ok(None);
}
let reason = if before.total_tokens >= emergency {
Some(CompactionReason::Emergency)
} else if before.total_tokens >= trigger
&& !self.state.last_turn.is_some_and(|last| {
turn.saturating_sub(last) < self.limits.compaction.min_turns_between
})
{
Some(CompactionReason::HighWatermark)
} else {
None
};
Ok(reason.map(|reason| CompactionDecision {
reason,
before,
target_tokens: target,
reclaimable_tokens,
}))
}
pub(crate) fn compact(
&mut self,
messages: &mut Vec<ModelMessage>,
tools: &[ModelToolDefinition],
current_user_prompt: &str,
checkpoint_candidate: Option<&str>,
turn: usize,
) -> Result<CompactionOutcome, KernelFailure> {
if has_unclosed_tool_transaction(messages) {
return Err(KernelFailure::new(
"CONTEXT_COMPACTION_DEFERRED",
"model transcript contains an unclosed ToolUse/ToolResult transaction",
));
}
let before = self.measurement(messages, tools);
let compiled = compile_checkpoint(
checkpoint_candidate,
messages,
current_user_prompt,
self.limits.compaction.checkpoint_max_chars,
)?;
let before_len = messages.len();
let mut resident = resident_messages(messages, current_user_prompt);
resident.push(compiled.message);
*messages = resident;
let after = self.measurement(messages, tools);
if after.total_tokens >= before.total_tokens
|| after.total_tokens > self.limits.max_input_tokens
{
return Err(KernelFailure::new(
"CONTEXT_COMPACTION_INSUFFICIENT",
format!(
"checkpoint compaction changed context from {} to {} tokens against a {} token limit",
before.total_tokens, after.total_tokens, self.limits.max_input_tokens
),
));
}
self.state.generation = self.state.generation.saturating_add(1);
self.state.last_turn = Some(turn);
Ok(CompactionOutcome {
before,
after,
removed_messages: before_len.saturating_sub(messages.len()),
used_fallback: compiled.used_fallback,
})
}
}
fn resident_messages(messages: &[ModelMessage], current_user_prompt: &str) -> Vec<ModelMessage> {
let mut resident = Vec::with_capacity(2);
if let Some(system) = messages
.first()
.filter(|message| message.role == ModelRole::System)
{
resident.push(system.clone());
}
let original_task = messages
.iter()
.find(|message| {
message.role == ModelRole::User
&& message.content.iter().any(|content| {
matches!(content, ModelContent::Text { text } if text == current_user_prompt)
})
})
.cloned()
.unwrap_or_else(|| ModelMessage::text(ModelRole::User, current_user_prompt));
resident.push(original_task);
resident
}
pub(crate) fn has_unclosed_tool_transaction(messages: &[ModelMessage]) -> bool {
let mut open = BTreeSet::<String>::new();
for message in messages {
for content in &message.content {
match content {
ModelContent::ToolUse { id, .. } => {
open.insert(id.clone());
}
ModelContent::ToolResult { tool_use_id, .. } => {
open.remove(tool_use_id);
}
ModelContent::Text { .. } => {}
}
}
}
!open.is_empty()
}
fn percent(value: u64, percentage: u8) -> u64 {
value
.saturating_mul(u64::from(percentage))
.saturating_div(100)
}
#[cfg(test)]
mod tests {
use runtime_kernel_api::{ContextCompactionLimits, ContextLimits};
use serde_json::json;
use super::*;
fn limits() -> ContextLimits {
ContextLimits {
max_messages: 64,
max_input_tokens: 1_000,
reserved_output_tokens: 100,
protocol_overhead_tokens: 50,
auto_compact: true,
compaction: ContextCompactionLimits {
trigger_percent: 50,
target_percent: 30,
emergency_percent: 90,
min_reclaim_percent: 5,
min_turns_between: 3,
checkpoint_max_chars: 1_024,
},
}
}
#[test]
fn unclosed_tool_transaction_is_detected() {
let messages = vec![ModelMessage {
role: ModelRole::Assistant,
content: vec![ModelContent::ToolUse {
id: "call".into(),
name: "run".into(),
arguments: json!({}),
}],
}];
assert!(has_unclosed_tool_transaction(&messages));
}
#[test]
fn high_watermark_compacts_to_resident_set_and_cooldown_is_deterministic() {
let mut controller = CompactionController::new(limits());
let mut messages = vec![
ModelMessage::text(ModelRole::System, "system"),
ModelMessage::text(ModelRole::User, "task"),
ModelMessage::text(ModelRole::Assistant, "x".repeat(2_000)),
];
assert!(
controller
.decide(&messages, &[], "task", 1)
.unwrap()
.is_some()
);
let outcome = controller
.compact(&mut messages, &[], "task", None, 1)
.unwrap();
assert!(outcome.used_fallback);
assert_eq!(messages.len(), 3);
assert!(
controller
.decide(&messages, &[], "task", 2)
.unwrap()
.is_none()
);
}
}