//! Context compaction for long sessions.
//!
//! This module ports the pi-mono compaction algorithm:
//! - Estimate context usage and choose a cut point that keeps recent context
//! - Summarize the discarded portion with the LLM (iteratively updating prior summaries)
//! - Record a `compaction` session entry containing the summary and cut point
//! - When building provider context, the session inserts the summary before the kept region
//! and omits older messages.
use crate::error::{Error, Result};
use crate::model::{
AssistantMessage, ContentBlock, Message, StopReason, TextContent, ThinkingLevel, ToolCall,
Usage, UserContent, UserMessage,
};
use crate::provider::{Context, Provider, StreamOptions};
use crate::session::{SessionEntry, SessionMessage, session_message_to_model};
use futures::StreamExt;
use serde::Serialize;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::fmt::Write as _;
use std::sync::Arc;
/// Approximate characters per token for English text with GPT-family tokenizers.
/// Intentionally conservative (overestimates tokens) to avoid exceeding context windows.
/// Set to 3 to safely account for code/symbol-heavy content which is denser than prose.
const CHARS_PER_TOKEN_ESTIMATE: usize = 3;
/// Estimated tokens for an image content block (~1200 tokens).
const IMAGE_TOKEN_ESTIMATE: usize = 1200;
/// Character-equivalent estimate for an image (IMAGE_TOKEN_ESTIMATE * CHARS_PER_TOKEN_ESTIMATE).
const IMAGE_CHAR_ESTIMATE: usize = IMAGE_TOKEN_ESTIMATE * CHARS_PER_TOKEN_ESTIMATE;
/// Count the serialized JSON byte length of a [`Value`] without allocating a `String`.
///
/// Uses `serde_json::to_writer` with a sink that only counts bytes – this gives the
/// exact same length as `serde_json::to_string(&v).len()` at zero heap cost.
fn json_byte_len(value: &Value) -> usize {
struct Counter(usize);
impl std::io::Write for Counter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0 += buf.len();
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let mut c = Counter(0);
if serde_json::to_writer(&mut c, value).is_err() {
// Fallback or partial count on error (e.g. recursion limit)
}
c.0
}
// =============================================================================
// Public types
// =============================================================================
#[derive(Debug, Clone)]
pub struct ResolvedCompactionSettings {
pub enabled: bool,
pub context_window_tokens: u32,
pub reserve_tokens: u32,
pub keep_recent_tokens: u32,
}
impl Default for ResolvedCompactionSettings {
fn default() -> Self {
let context_window_tokens: u32 = 200_000;
Self {
enabled: true,
context_window_tokens,
// ~8% of context window
reserve_tokens: 16_384,
// 10% of context window
keep_recent_tokens: 20_000,
}
}
}
/// Details stored in `CompactionEntry.details` for cumulative file tracking.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionDetails {
pub read_files: Vec<String>,
pub modified_files: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionResult {
pub summary: String,
pub first_kept_entry_id: String,
pub tokens_before: u64,
pub details: CompactionDetails,
}
#[derive(Debug, Clone)]
pub struct CompactionPreparation {
pub first_kept_entry_id: String,
pub messages_to_summarize: Vec<SessionMessage>,
pub turn_prefix_messages: Vec<SessionMessage>,
pub is_split_turn: bool,
pub tokens_before: u64,
pub previous_summary: Option<String>,
pub file_ops: FileOperations,
pub settings: ResolvedCompactionSettings,
}
// =============================================================================
// File op tracking (read/write/edit)
// =============================================================================
#[derive(Debug, Clone, Default)]
pub struct FileOperations {
read: HashSet<String>,
written: HashSet<String>,
edited: HashSet<String>,
}
impl FileOperations {
pub fn read_files(&self) -> impl Iterator<Item = &str> {
self.read.iter().map(String::as_str)
}
}
fn build_tool_status_map(messages: &[SessionMessage]) -> HashMap<String, bool> {
let mut status = HashMap::new();
for msg in messages {
if let SessionMessage::ToolResult {
tool_call_id,
is_error,
..
} = msg
{
status.insert(tool_call_id.clone(), !*is_error);
}
}
status
}
fn extract_file_ops_from_message(
message: &SessionMessage,
file_ops: &mut FileOperations,
tool_status: &HashMap<String, bool>,
) {
let SessionMessage::Assistant { message } = message else {
return;
};
for block in &message.content {
let ContentBlock::ToolCall(ToolCall {
id,
name,
arguments,
..
}) = block
else {
continue;
};
// Only track successful tool calls.
if !tool_status.get(id).copied().unwrap_or(false) {
continue;
}
let Some(path) = arguments.get("path").and_then(Value::as_str) else {
continue;
};
match name.as_str() {
"read" | "grep" | "find" | "ls" => {
file_ops.read.insert(path.to_string());
}
"write" => {
file_ops.written.insert(path.to_string());
}
"edit" => {
file_ops.edited.insert(path.to_string());
}
_ => {}
}
}
}
fn compute_file_lists(file_ops: &FileOperations) -> (Vec<String>, Vec<String>) {
let modified: HashSet<&String> = file_ops
.edited
.iter()
.chain(file_ops.written.iter())
.collect();
let mut read_only = file_ops
.read
.iter()
.filter(|f| !modified.contains(f))
.cloned()
.collect::<Vec<_>>();
read_only.sort();
let mut modified_files = modified.into_iter().cloned().collect::<Vec<_>>();
modified_files.sort();
(read_only, modified_files)
}
fn write_escaped_file_list(out: &mut String, tag: &str, files: &[String]) {
out.push('<');
out.push_str(tag);
out.push_str(">\n");
for (i, file) in files.iter().enumerate() {
if i > 0 {
out.push('\n');
}
// Inline escape: replace < and > in file paths
for ch in file.chars() {
match ch {
'<' => out.push_str("<"),
'>' => out.push_str(">"),
_ => out.push(ch),
}
}
}
out.push_str("\n</");
out.push_str(tag);
out.push('>');
}
fn format_file_operations(read_files: &[String], modified_files: &[String]) -> String {
if read_files.is_empty() && modified_files.is_empty() {
return String::new();
}
let mut out = String::from("\n\n");
if !read_files.is_empty() {
write_escaped_file_list(&mut out, "read-files", read_files);
}
if !modified_files.is_empty() {
if !read_files.is_empty() {
out.push_str("\n\n");
}
write_escaped_file_list(&mut out, "modified-files", modified_files);
}
out
}
// =============================================================================
// Token estimation
// =============================================================================
const fn calculate_context_tokens(usage: &Usage) -> u64 {
if usage.total_tokens > 0 {
usage.total_tokens
} else {
usage.input + usage.output
}
}
const fn get_assistant_usage(message: &SessionMessage) -> Option<&Usage> {
let SessionMessage::Assistant { message } = message else {
return None;
};
if matches!(message.stop_reason, StopReason::Aborted | StopReason::Error) {
return None;
}
Some(&message.usage)
}
#[derive(Debug, Clone, Copy)]
struct ContextUsageEstimate {
tokens: u64,
last_usage_index: Option<usize>,
}
fn estimate_context_tokens(messages: &[SessionMessage]) -> ContextUsageEstimate {
let mut last_usage: Option<(&Usage, usize)> = None;
for (idx, msg) in messages.iter().enumerate().rev() {
if let Some(usage) = get_assistant_usage(msg) {
last_usage = Some((usage, idx));
break;
}
}
let Some((usage, usage_index)) = last_usage else {
let total = messages.iter().map(estimate_tokens).sum();
return ContextUsageEstimate {
tokens: total,
last_usage_index: None,
};
};
let usage_tokens = calculate_context_tokens(usage);
let trailing_tokens = messages[usage_index + 1..]
.iter()
.map(estimate_tokens)
.sum::<u64>();
ContextUsageEstimate {
tokens: usage_tokens + trailing_tokens,
last_usage_index: Some(usage_index),
}
}
fn should_compact(
context_tokens: u64,
context_window: u32,
settings: &ResolvedCompactionSettings,
) -> bool {
if !settings.enabled {
return false;
}
let reserve = u64::from(settings.reserve_tokens);
let window = u64::from(context_window);
context_tokens > window.saturating_sub(reserve)
}
fn estimate_tokens(message: &SessionMessage) -> u64 {
let mut chars: usize = 0;
match message {
SessionMessage::User { content, .. } => match content {
UserContent::Text(text) => chars = text.len(),
UserContent::Blocks(blocks) => {
for block in blocks {
match block {
ContentBlock::Text(text) => chars += text.text.len(),
ContentBlock::Image(_) => chars += IMAGE_CHAR_ESTIMATE,
ContentBlock::Thinking(thinking) => chars += thinking.thinking.len(),
ContentBlock::ToolCall(call) => {
chars += call.name.len();
chars += json_byte_len(&call.arguments);
}
}
}
}
},
SessionMessage::Assistant { message } => {
for block in &message.content {
match block {
ContentBlock::Text(text) => chars += text.text.len(),
ContentBlock::Thinking(thinking) => chars += thinking.thinking.len(),
ContentBlock::Image(_) => chars += IMAGE_CHAR_ESTIMATE,
ContentBlock::ToolCall(call) => {
chars += call.name.len();
chars += json_byte_len(&call.arguments);
}
}
}
}
SessionMessage::ToolResult { content, .. } => {
for block in content {
match block {
ContentBlock::Text(text) => chars += text.text.len(),
ContentBlock::Thinking(thinking) => chars += thinking.thinking.len(),
ContentBlock::Image(_) => chars += IMAGE_CHAR_ESTIMATE,
ContentBlock::ToolCall(call) => {
chars += call.name.len();
chars += json_byte_len(&call.arguments);
}
}
}
}
SessionMessage::Custom { content, .. } => chars = content.len(),
SessionMessage::BashExecution {
command, output, ..
} => chars = command.len() + output.len(),
SessionMessage::BranchSummary { summary, .. }
| SessionMessage::CompactionSummary { summary, .. } => chars = summary.len(),
}
u64::try_from(chars.div_ceil(CHARS_PER_TOKEN_ESTIMATE)).unwrap_or(u64::MAX)
}
// =============================================================================
// Cut point detection
// =============================================================================
#[derive(Debug, Clone, Copy)]
struct CutPointResult {
first_kept_entry_index: usize,
turn_start_index: Option<usize>,
is_split_turn: bool,
}
fn message_from_entry(entry: &SessionEntry) -> Option<SessionMessage> {
match entry {
SessionEntry::Message(msg_entry) => Some(msg_entry.message.clone()),
SessionEntry::BranchSummary(summary) => Some(SessionMessage::BranchSummary {
summary: summary.summary.clone(),
from_id: summary.from_id.clone(),
}),
SessionEntry::Compaction(compaction) => Some(SessionMessage::CompactionSummary {
summary: compaction.summary.clone(),
tokens_before: compaction.tokens_before,
}),
_ => None,
}
}
const fn entry_is_message_like(entry: &SessionEntry) -> bool {
matches!(
entry,
SessionEntry::Message(_) | SessionEntry::BranchSummary(_)
)
}
const fn entry_is_compaction_boundary(entry: &SessionEntry) -> bool {
matches!(entry, SessionEntry::Compaction(_))
}
fn find_valid_cut_points(
entries: &[SessionEntry],
start_index: usize,
end_index: usize,
) -> Vec<usize> {
let mut cut_points = Vec::new();
for (idx, entry) in entries.iter().enumerate().take(end_index).skip(start_index) {
match entry {
SessionEntry::Message(msg_entry) => match msg_entry.message {
SessionMessage::ToolResult { .. } => {}
_ => cut_points.push(idx),
},
SessionEntry::BranchSummary(_) => cut_points.push(idx),
_ => {}
}
}
cut_points
}
fn entry_has_tool_calls(entry: &SessionEntry) -> bool {
matches!(
entry,
SessionEntry::Message(msg) if matches!(
&msg.message,
SessionMessage::Assistant { message } if message.content.iter().any(|b| matches!(b, ContentBlock::ToolCall(_)))
)
)
}
const fn is_user_turn_start(entry: &SessionEntry) -> bool {
match entry {
SessionEntry::BranchSummary(_) => true,
SessionEntry::Message(msg_entry) => matches!(
msg_entry.message,
SessionMessage::User { .. } | SessionMessage::BashExecution { .. }
),
_ => false,
}
}
fn find_turn_start_index(
entries: &[SessionEntry],
entry_index: usize,
start_index: usize,
) -> Option<usize> {
(start_index..=entry_index)
.rev()
.find(|&idx| is_user_turn_start(&entries[idx]))
}
fn find_cut_point(
entries: &[SessionEntry],
start_index: usize,
end_index: usize,
keep_recent_tokens: u32,
) -> CutPointResult {
let cut_points = find_valid_cut_points(entries, start_index, end_index);
if cut_points.is_empty() {
return CutPointResult {
first_kept_entry_index: start_index,
turn_start_index: None,
is_split_turn: false,
};
}
let mut accumulated_tokens: u64 = 0;
let mut cut_index = cut_points[0];
for i in (start_index..end_index).rev() {
let entry = &entries[i];
let SessionEntry::Message(msg_entry) = entry else {
continue;
};
accumulated_tokens = accumulated_tokens.saturating_add(estimate_tokens(&msg_entry.message));
if accumulated_tokens >= u64::from(keep_recent_tokens) {
// Binary search: find the largest cut point <= i.
// `partition_point` returns the index of the first element > i,
// so idx-1 is the largest element <= i (if any).
let pos = cut_points.partition_point(|&cp| cp <= i);
if pos > 0 {
cut_index = cut_points[pos - 1];
}
// else: no cut point <= i, keep the fallback (cut_points[0])
break;
}
}
while cut_index > start_index {
let prev = &entries[cut_index - 1];
if entry_is_compaction_boundary(prev) {
break;
}
if entry_is_message_like(prev) {
break;
}
cut_index -= 1;
}
let is_user_message = is_user_turn_start(&entries[cut_index]);
let turn_start_index = if is_user_message {
None
} else {
find_turn_start_index(entries, cut_index, start_index)
};
CutPointResult {
first_kept_entry_index: cut_index,
turn_start_index,
is_split_turn: !is_user_message && turn_start_index.is_some(),
}
}
// =============================================================================
// Summarization prompts
// =============================================================================
const SUMMARIZATION_SYSTEM_PROMPT: &str = "You are a context summarization assistant. Your task is to read a conversation between a user and an AI coding assistant, then produce a structured summary following the exact format specified.\n\nDo NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.";
const SUMMARIZATION_PROMPT: &str = "The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.\n\nUse this EXACT format:\n\n## Goal\n[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]\n\n## Constraints & Preferences\n- [Any constraints, preferences, or requirements mentioned by user]\n- [Or \"(none)\" if none were mentioned]\n\n## Progress\n### Done\n- [x] [Completed tasks/changes]\n\n### In Progress\n- [ ] [Current work]\n\n### Blocked\n- [Issues preventing progress, if any]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale]\n\n## Next Steps\n1. [Ordered list of what should happen next]\n\n## Critical Context\n- [Any data, examples, or references needed to continue]\n- [Or \"(none)\" if not applicable]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";
const UPDATE_SUMMARIZATION_PROMPT: &str = "The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.\n\nUpdate the existing structured summary with new information. RULES:\n- PRESERVE all existing information from the previous summary\n- ADD new progress, decisions, and context from the new messages\n- UPDATE the Progress section: move items from \"In Progress\" to \"Done\" when completed\n- UPDATE \"Next Steps\" based on what was accomplished\n- PRESERVE exact file paths, function names, and error messages\n- If something is no longer relevant, you may remove it\n\nUse this EXACT format:\n\n## Goal\n[Preserve existing goals, add new ones if the task expanded]\n\n## Constraints & Preferences\n- [Preserve existing, add new ones discovered]\n\n## Progress\n### Done\n- [x] [Include previously done items AND newly completed items]\n\n### In Progress\n- [ ] [Current work - update based on progress]\n\n### Blocked\n- [Current blockers - remove if resolved]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale] (preserve all previous, add new)\n\n## Next Steps\n1. [Update based on current state]\n\n## Critical Context\n- [Preserve important context, add new if needed]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";
const TURN_PREFIX_SUMMARIZATION_PROMPT: &str = "This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.\n\nSummarize the prefix to provide context for the retained suffix:\n\n## Original Request\n[What did the user ask for in this turn?]\n\n## Early Progress\n- [Key decisions and work done in the prefix]\n\n## Context for Suffix\n- [Information needed to understand the retained recent work]\n\nBe concise. Focus on what's needed to understand the kept suffix.";
fn push_message_separator(out: &mut String) {
if !out.is_empty() {
out.push_str("\n\n");
}
}
fn user_has_serializable_content(user: &UserMessage) -> bool {
match &user.content {
UserContent::Text(text) => !text.is_empty(),
UserContent::Blocks(blocks) => blocks
.iter()
.any(|c| matches!(c, ContentBlock::Text(t) if !t.text.is_empty())),
}
}
fn append_user_message(out: &mut String, user: &UserMessage) {
if !user_has_serializable_content(user) {
return;
}
push_message_separator(out);
out.push_str("[User]: ");
match &user.content {
UserContent::Text(text) => out.push_str(text),
UserContent::Blocks(blocks) => {
for block in blocks {
if let ContentBlock::Text(text) = block {
out.push_str(&text.text);
}
}
}
}
}
fn append_custom_message(out: &mut String, custom_type: &str, content: &str) {
if content.trim().is_empty() {
return;
}
push_message_separator(out);
out.push('[');
if custom_type.trim().is_empty() {
out.push_str("Custom");
} else {
out.push_str("Custom:");
out.push_str(custom_type);
}
out.push_str("]: ");
out.push_str(content);
}
fn assistant_content_flags(assistant: &AssistantMessage) -> (bool, bool, bool) {
let mut has_thinking = false;
let mut has_text = false;
let mut has_tools = false;
for block in &assistant.content {
match block {
ContentBlock::Thinking(_) => has_thinking = true,
ContentBlock::Text(_) => has_text = true,
ContentBlock::ToolCall(_) => has_tools = true,
ContentBlock::Image(_) => {}
}
}
(has_thinking, has_text, has_tools)
}
fn append_assistant_thinking(out: &mut String, assistant: &AssistantMessage) {
push_message_separator(out);
out.push_str("[Assistant thinking]: ");
let mut first = true;
for block in &assistant.content {
if let ContentBlock::Thinking(thinking) = block {
if !first {
out.push('\n');
}
out.push_str(&thinking.thinking);
first = false;
}
}
}
fn append_assistant_text(out: &mut String, assistant: &AssistantMessage) {
push_message_separator(out);
out.push_str("[Assistant]: ");
let mut first = true;
for block in &assistant.content {
if let ContentBlock::Text(text) = block {
if !first {
out.push('\n');
}
out.push_str(&text.text);
first = false;
}
}
}
fn append_tool_call_arguments(out: &mut String, arguments: &Value) {
if let Some(obj) = arguments.as_object() {
let mut first_kv = true;
for (k, v) in obj {
if !first_kv {
out.push_str(", ");
}
out.push_str(k);
out.push('=');
match serde_json::to_string(v) {
Ok(s) => out.push_str(&s),
Err(_) => {
let _ = write!(out, "{v}");
}
}
first_kv = false;
}
} else {
match serde_json::to_string(arguments) {
Ok(s) => out.push_str(&s),
Err(_) => {
let _ = write!(out, "{arguments}");
}
}
}
}
fn append_assistant_tool_calls(out: &mut String, assistant: &AssistantMessage) {
push_message_separator(out);
out.push_str("[Assistant tool calls]: ");
let mut first = true;
for block in &assistant.content {
if let ContentBlock::ToolCall(call) = block {
if !first {
out.push_str("; ");
}
out.push_str(&call.name);
out.push('(');
append_tool_call_arguments(out, &call.arguments);
out.push(')');
first = false;
}
}
}
fn append_assistant_message(out: &mut String, assistant: &AssistantMessage) {
let (has_thinking, has_text, has_tools) = assistant_content_flags(assistant);
if has_thinking {
append_assistant_thinking(out, assistant);
}
if has_text {
append_assistant_text(out, assistant);
}
if has_tools {
append_assistant_tool_calls(out, assistant);
}
}
fn tool_result_has_serializable_content(content: &[ContentBlock]) -> bool {
content
.iter()
.any(|c| matches!(c, ContentBlock::Text(t) if !t.text.is_empty()))
}
fn append_tool_result_message(out: &mut String, content: &[ContentBlock]) {
if !tool_result_has_serializable_content(content) {
return;
}
push_message_separator(out);
out.push_str("[Tool result]: ");
for block in content {
if let ContentBlock::Text(text) = block {
out.push_str(&text.text);
}
}
}
fn collect_text_blocks(blocks: &[ContentBlock]) -> String {
let mut out = String::new();
let mut first = true;
for block in blocks {
if let ContentBlock::Text(text) = block {
if !first {
out.push('\n');
}
out.push_str(&text.text);
first = false;
}
}
out
}
fn serialize_conversation(messages: &[Message]) -> String {
let mut out = String::new();
for msg in messages {
match msg {
Message::User(user) => append_user_message(&mut out, user),
Message::Custom(custom) => {
append_custom_message(&mut out, &custom.custom_type, &custom.content);
}
Message::Assistant(assistant) => append_assistant_message(&mut out, assistant),
Message::ToolResult(tool) => append_tool_result_message(&mut out, &tool.content),
}
}
out
}
async fn complete_simple(
provider: Arc<dyn Provider>,
system_prompt: &str,
prompt_text: String,
api_key: &str,
reserve_tokens: u32,
max_tokens_factor: f64,
) -> Result<AssistantMessage> {
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
let max_tokens = (f64::from(reserve_tokens) * max_tokens_factor).floor() as u32;
let max_tokens = max_tokens.max(256);
let context = Context {
system_prompt: Some(system_prompt.to_string().into()),
messages: vec![Message::User(UserMessage {
content: UserContent::Blocks(vec![ContentBlock::Text(TextContent::new(prompt_text))]),
timestamp: chrono::Utc::now().timestamp_millis(),
})]
.into(),
tools: Vec::new().into(),
};
let options = StreamOptions {
api_key: Some(api_key.to_string()),
max_tokens: Some(max_tokens),
thinking_level: Some(ThinkingLevel::High),
..Default::default()
};
let mut stream = provider.stream(&context, &options).await?;
let mut final_message: Option<AssistantMessage> = None;
while let Some(event) = stream.next().await {
match event? {
crate::model::StreamEvent::Done { message, .. } => {
final_message = Some(message);
}
crate::model::StreamEvent::Error { error, .. } => {
let msg = error
.error_message
.unwrap_or_else(|| "Summarization error".to_string());
return Err(Error::api(msg));
}
_ => {}
}
}
let message = final_message.ok_or_else(|| Error::api("Stream ended without Done event"))?;
if matches!(message.stop_reason, StopReason::Aborted | StopReason::Error) {
let msg = message
.error_message
.unwrap_or_else(|| "Summarization error".to_string());
return Err(Error::api(msg));
}
Ok(message)
}
async fn generate_summary(
messages: &[SessionMessage],
provider: Arc<dyn Provider>,
api_key: &str,
settings: &ResolvedCompactionSettings,
custom_instructions: Option<&str>,
previous_summary: Option<&str>,
) -> Result<String> {
let base_prompt = if previous_summary.is_some() {
UPDATE_SUMMARIZATION_PROMPT
} else {
SUMMARIZATION_PROMPT
};
let mut prompt = base_prompt.to_string();
if let Some(custom) = custom_instructions.filter(|s| !s.trim().is_empty()) {
let _ = write!(prompt, "\n\nAdditional focus: {custom}");
}
let llm_messages = messages
.iter()
.filter_map(session_message_to_model)
.collect::<Vec<_>>();
let conversation_text = serialize_conversation(&llm_messages);
let mut prompt_text = format!("<conversation>\n{conversation_text}\n</conversation>\n\n");
if let Some(previous) = previous_summary {
let _ = write!(
prompt_text,
"<previous-summary>\n{previous}\n</previous-summary>\n\n"
);
}
prompt_text.push_str(&prompt);
let assistant = complete_simple(
provider,
SUMMARIZATION_SYSTEM_PROMPT,
prompt_text,
api_key,
settings.reserve_tokens,
0.8,
)
.await?;
let text = collect_text_blocks(&assistant.content);
if text.trim().is_empty() {
return Err(Error::api(
"Summarization returned empty text; refusing to store empty compaction summary",
));
}
Ok(text)
}
async fn generate_turn_prefix_summary(
messages: &[SessionMessage],
provider: Arc<dyn Provider>,
api_key: &str,
settings: &ResolvedCompactionSettings,
) -> Result<String> {
let llm_messages = messages
.iter()
.filter_map(session_message_to_model)
.collect::<Vec<_>>();
let conversation_text = serialize_conversation(&llm_messages);
let prompt_text = format!(
"<conversation>\n{conversation_text}\n</conversation>\n\n{TURN_PREFIX_SUMMARIZATION_PROMPT}"
);
let assistant = complete_simple(
provider,
SUMMARIZATION_SYSTEM_PROMPT,
prompt_text,
api_key,
settings.reserve_tokens,
0.5,
)
.await?;
let text = collect_text_blocks(&assistant.content);
if text.trim().is_empty() {
return Err(Error::api(
"Turn prefix summarization returned empty text; refusing to store empty summary",
));
}
Ok(text)
}
// =============================================================================
// Public API
// =============================================================================
#[allow(clippy::too_many_lines)]
pub fn prepare_compaction(
path_entries: &[SessionEntry],
settings: ResolvedCompactionSettings,
) -> Option<CompactionPreparation> {
if path_entries.is_empty() {
return None;
}
if path_entries
.last()
.is_some_and(|entry| matches!(entry, SessionEntry::Compaction(_)))
{
return None;
}
let mut prev_compaction_index: Option<usize> = None;
for (idx, entry) in path_entries.iter().enumerate().rev() {
if matches!(entry, SessionEntry::Compaction(_)) {
prev_compaction_index = Some(idx);
break;
}
}
let boundary_start = prev_compaction_index.map_or(0, |i| i + 1);
let boundary_end = path_entries.len();
let usage_start = prev_compaction_index.unwrap_or(0);
let mut usage_messages = Vec::new();
for entry in &path_entries[usage_start..boundary_end] {
if let Some(msg) = message_from_entry(entry) {
usage_messages.push(msg);
}
}
// Calculate the tokens *currently* occupied by the segment we are about to compact.
// If the segment includes a previous compaction summary, this counts the *summary* tokens,
// not the original uncompressed history tokens. This effectively tracks the "compressed size"
// of the history prior to the new cut point.
let tokens_before = estimate_context_tokens(&usage_messages).tokens;
if !should_compact(tokens_before, settings.context_window_tokens, &settings) {
return None;
}
let cut_point = find_cut_point(
path_entries,
boundary_start,
boundary_end,
settings.keep_recent_tokens,
);
let first_kept_entry = &path_entries[cut_point.first_kept_entry_index];
let first_kept_entry_id = first_kept_entry.base_id()?.clone();
let history_end = if cut_point.is_split_turn {
cut_point.turn_start_index?
} else {
cut_point.first_kept_entry_index
};
let mut messages_to_summarize = Vec::new();
for entry in &path_entries[boundary_start..history_end] {
if let Some(msg) = message_from_entry(entry) {
messages_to_summarize.push(msg);
}
}
let mut turn_prefix_messages = Vec::new();
if cut_point.is_split_turn {
let turn_start = cut_point.turn_start_index?;
for entry in &path_entries[turn_start..cut_point.first_kept_entry_index] {
if let Some(msg) = message_from_entry(entry) {
turn_prefix_messages.push(msg);
}
}
}
// No-op compaction: if there's nothing to summarize, don't issue an LLM call and don't append a
// compaction entry. This can happen early in a session (e.g. session header entries only).
if messages_to_summarize.is_empty() && turn_prefix_messages.is_empty() {
return None;
}
let previous_summary = prev_compaction_index.and_then(|idx| match &path_entries[idx] {
SessionEntry::Compaction(entry) => Some(entry.summary.clone()),
_ => None,
});
let mut file_ops = FileOperations::default();
// Collect file tracking from previous compaction details if pi-generated.
if let Some(idx) = prev_compaction_index {
if let SessionEntry::Compaction(entry) = &path_entries[idx] {
if !entry.from_hook.unwrap_or(false) {
if let Some(details) = entry.details.as_ref().and_then(Value::as_object) {
if let Some(read_files) = details.get("readFiles").and_then(Value::as_array) {
for item in read_files.iter().filter_map(Value::as_str) {
file_ops.read.insert(item.to_string());
}
}
if let Some(modified_files) =
details.get("modifiedFiles").and_then(Value::as_array)
{
for item in modified_files.iter().filter_map(Value::as_str) {
file_ops.edited.insert(item.to_string());
}
}
}
}
}
}
let mut tool_status = build_tool_status_map(&messages_to_summarize);
tool_status.extend(build_tool_status_map(&turn_prefix_messages));
for msg in &messages_to_summarize {
extract_file_ops_from_message(msg, &mut file_ops, &tool_status);
}
for msg in &turn_prefix_messages {
extract_file_ops_from_message(msg, &mut file_ops, &tool_status);
}
Some(CompactionPreparation {
first_kept_entry_id,
messages_to_summarize,
turn_prefix_messages,
is_split_turn: cut_point.is_split_turn,
tokens_before,
previous_summary,
file_ops,
settings,
})
}
pub async fn summarize_entries(
entries: &[SessionEntry],
provider: Arc<dyn Provider>,
api_key: &str,
reserve_tokens: u32,
custom_instructions: Option<&str>,
) -> Result<Option<String>> {
let mut messages = Vec::new();
for entry in entries {
if let Some(message) = message_from_entry(entry) {
messages.push(message);
}
}
if messages.is_empty() {
return Ok(None);
}
let settings = ResolvedCompactionSettings {
enabled: true,
reserve_tokens,
keep_recent_tokens: 0,
..Default::default()
};
let summary = generate_summary(
&messages,
provider,
api_key,
&settings,
custom_instructions,
None,
)
.await?;
Ok(Some(summary))
}
pub async fn compact(
preparation: CompactionPreparation,
provider: Arc<dyn Provider>,
api_key: &str,
custom_instructions: Option<&str>,
) -> Result<CompactionResult> {
let summary = if preparation.is_split_turn && !preparation.turn_prefix_messages.is_empty() {
let history_summary = if preparation.messages_to_summarize.is_empty() {
"No prior history.".to_string()
} else {
generate_summary(
&preparation.messages_to_summarize,
Arc::clone(&provider),
api_key,
&preparation.settings,
custom_instructions,
preparation.previous_summary.as_deref(),
)
.await?
};
let turn_prefix_summary = generate_turn_prefix_summary(
&preparation.turn_prefix_messages,
Arc::clone(&provider),
api_key,
&preparation.settings,
)
.await?;
format!(
"{history_summary}\n\n---\n\n**Turn Context (split turn):**\n\n{turn_prefix_summary}"
)
} else {
generate_summary(
&preparation.messages_to_summarize,
Arc::clone(&provider),
api_key,
&preparation.settings,
custom_instructions,
preparation.previous_summary.as_deref(),
)
.await?
};
let (read_files, modified_files) = compute_file_lists(&preparation.file_ops);
let details = CompactionDetails {
read_files: read_files.clone(),
modified_files: modified_files.clone(),
};
let mut summary = summary;
summary.push_str(&format_file_operations(&read_files, &modified_files));
Ok(CompactionResult {
summary,
first_kept_entry_id: preparation.first_kept_entry_id,
tokens_before: preparation.tokens_before,
details,
})
}
pub fn compaction_details_to_value(details: &CompactionDetails) -> Result<Value> {
serde_json::to_value(details).map_err(|e| Error::session(format!("Compaction details: {e}")))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{AssistantMessage, ContentBlock, TextContent, Usage};
use serde_json::json;
fn make_user_text(text: &str) -> SessionMessage {
SessionMessage::User {
content: UserContent::Text(text.to_string()),
timestamp: Some(0),
}
}
fn make_assistant_text(text: &str, input: u64, output: u64) -> SessionMessage {
SessionMessage::Assistant {
message: AssistantMessage {
content: vec![ContentBlock::Text(TextContent::new(text))],
api: String::new(),
provider: String::new(),
model: String::new(),
stop_reason: StopReason::Stop,
error_message: None,
timestamp: 0,
usage: Usage {
input,
output,
cache_read: 0,
cache_write: 0,
total_tokens: input + output,
..Default::default()
},
},
}
}
fn make_assistant_tool_call(name: &str, args: Value) -> SessionMessage {
SessionMessage::Assistant {
message: AssistantMessage {
content: vec![ContentBlock::ToolCall(ToolCall {
id: "call_1".to_string(),
name: name.to_string(),
arguments: args,
thought_signature: None,
})],
api: String::new(),
provider: String::new(),
model: String::new(),
stop_reason: StopReason::ToolUse,
error_message: None,
timestamp: 0,
usage: Usage::default(),
},
}
}
fn make_tool_result(text: &str) -> SessionMessage {
SessionMessage::ToolResult {
tool_call_id: "call_1".to_string(),
tool_name: String::new(),
content: vec![ContentBlock::Text(TextContent::new(text))],
details: None,
is_error: false,
timestamp: None,
}
}
// ── calculate_context_tokens ─────────────────────────────────────
#[test]
fn context_tokens_prefers_total_tokens() {
let usage = Usage {
input: 100,
output: 50,
total_tokens: 200,
..Default::default()
};
assert_eq!(calculate_context_tokens(&usage), 200);
}
#[test]
fn context_tokens_falls_back_to_input_plus_output() {
let usage = Usage {
input: 100,
output: 50,
total_tokens: 0,
..Default::default()
};
assert_eq!(calculate_context_tokens(&usage), 150);
}
// ── should_compact ───────────────────────────────────────────────
#[test]
fn should_compact_when_over_threshold() {
let settings = ResolvedCompactionSettings {
enabled: true,
reserve_tokens: 10_000,
keep_recent_tokens: 5_000,
..Default::default()
};
// window=100k, reserve=10k => threshold=90k, context=95k => should compact
assert!(should_compact(95_000, 100_000, &settings));
}
#[test]
fn should_not_compact_when_under_threshold() {
let settings = ResolvedCompactionSettings {
enabled: true,
reserve_tokens: 10_000,
keep_recent_tokens: 5_000,
..Default::default()
};
// window=100k, reserve=10k => threshold=90k, context=80k => should not compact
assert!(!should_compact(80_000, 100_000, &settings));
}
#[test]
fn should_not_compact_when_disabled() {
let settings = ResolvedCompactionSettings {
enabled: false,
reserve_tokens: 0,
keep_recent_tokens: 0,
..Default::default()
};
assert!(!should_compact(1_000_000, 100_000, &settings));
}
#[test]
fn should_compact_at_exact_threshold() {
let settings = ResolvedCompactionSettings {
enabled: true,
reserve_tokens: 10_000,
keep_recent_tokens: 5_000,
..Default::default()
};
// window=100k, reserve=10k => threshold=90k, context=90k => NOT compacting (not >)
assert!(!should_compact(90_000, 100_000, &settings));
// 90001 should trigger
assert!(should_compact(90_001, 100_000, &settings));
}
// ── estimate_tokens ──────────────────────────────────────────────
#[test]
fn estimate_tokens_user_text() {
let msg = make_user_text("hello world"); // 11 chars => ceil(11/3) = 4
assert_eq!(estimate_tokens(&msg), 4);
}
#[test]
fn estimate_tokens_empty_text() {
let msg = make_user_text(""); // 0 chars => 0
assert_eq!(estimate_tokens(&msg), 0);
}
#[test]
fn estimate_tokens_assistant_text() {
let msg = make_assistant_text("hello", 10, 5); // 5 chars => ceil(5/3) = 2
assert_eq!(estimate_tokens(&msg), 2);
}
#[test]
fn estimate_tokens_tool_result() {
let msg = make_tool_result("file contents here"); // 18 chars => ceil(18/3) = 6
assert_eq!(estimate_tokens(&msg), 6);
}
#[test]
fn estimate_tokens_custom_message() {
let msg = SessionMessage::Custom {
custom_type: "system".to_string(),
content: "some custom content".to_string(),
display: true,
details: None,
timestamp: Some(0),
};
// 19 chars => ceil(19/3) = 7
assert_eq!(estimate_tokens(&msg), 7);
}
// ── estimate_context_tokens ──────────────────────────────────────
#[test]
fn estimate_context_with_assistant_usage() {
let messages = vec![
make_user_text("hi"),
make_assistant_text("hello", 50, 10),
make_user_text("bye"),
];
let estimate = estimate_context_tokens(&messages);
// Last assistant usage: input=50, output=10, total=60
// Trailing after that: "bye" = ceil(3/3) = 1
assert_eq!(estimate.tokens, 61);
assert_eq!(estimate.last_usage_index, Some(1));
}
#[test]
fn estimate_context_no_assistant() {
let messages = vec![make_user_text("hello"), make_user_text("world")];
let estimate = estimate_context_tokens(&messages);
// No assistant messages, so sum estimate_tokens for all: ceil(5/3)+ceil(5/3) = 2+2 = 4
assert_eq!(estimate.tokens, 4);
assert!(estimate.last_usage_index.is_none());
}
// ── extract_file_ops_from_message ────────────────────────────────
#[test]
fn extract_file_ops_read() {
let msg = make_assistant_tool_call("read", json!({"path": "/foo/bar.rs"}));
let mut ops = FileOperations::default();
let mut status = HashMap::new();
status.insert("call_1".to_string(), true);
extract_file_ops_from_message(&msg, &mut ops, &status);
assert!(ops.read.contains("/foo/bar.rs"));
assert!(ops.written.is_empty());
assert!(ops.edited.is_empty());
}
#[test]
fn extract_file_ops_write() {
let msg = make_assistant_tool_call("write", json!({"path": "/out.txt"}));
let mut ops = FileOperations::default();
let mut status = HashMap::new();
status.insert("call_1".to_string(), true);
extract_file_ops_from_message(&msg, &mut ops, &status);
assert!(ops.written.contains("/out.txt"));
assert!(ops.read.is_empty());
}
#[test]
fn extract_file_ops_edit() {
let msg = make_assistant_tool_call("edit", json!({"path": "/src/main.rs"}));
let mut ops = FileOperations::default();
let mut status = HashMap::new();
status.insert("call_1".to_string(), true);
extract_file_ops_from_message(&msg, &mut ops, &status);
assert!(ops.edited.contains("/src/main.rs"));
}
#[test]
fn extract_file_ops_ignores_failed_tools() {
let msg = make_assistant_tool_call("read", json!({"path": "/secret.rs"}));
let mut ops = FileOperations::default();
let mut status = HashMap::new();
status.insert("call_1".to_string(), false); // Failed!
extract_file_ops_from_message(&msg, &mut ops, &status);
assert!(ops.read.is_empty());
}
#[test]
fn extract_file_ops_ignores_other_tools() {
let msg = make_assistant_tool_call("bash", json!({"command": "ls"}));
let mut ops = FileOperations::default();
let mut status = HashMap::new();
status.insert("call_1".to_string(), true);
extract_file_ops_from_message(&msg, &mut ops, &status);
assert!(ops.read.is_empty());
assert!(ops.written.is_empty());
assert!(ops.edited.is_empty());
}
#[test]
fn extract_file_ops_ignores_user_messages() {
let msg = make_user_text("read the file /foo.rs");
let mut ops = FileOperations::default();
let status = HashMap::new();
extract_file_ops_from_message(&msg, &mut ops, &status);
assert!(ops.read.is_empty());
}
// ── compute_file_lists ───────────────────────────────────────────
#[test]
fn compute_file_lists_separates_read_from_modified() {
let mut ops = FileOperations::default();
ops.read.insert("/a.rs".to_string());
ops.read.insert("/b.rs".to_string());
ops.written.insert("/b.rs".to_string());
ops.edited.insert("/c.rs".to_string());
let (read_only, modified) = compute_file_lists(&ops);
// /a.rs was only read; /b.rs was read AND written (so it's modified)
assert_eq!(read_only, vec!["/a.rs"]);
assert!(modified.contains(&"/b.rs".to_string()));
assert!(modified.contains(&"/c.rs".to_string()));
}
#[test]
fn compute_file_lists_empty() {
let ops = FileOperations::default();
let (read_only, modified) = compute_file_lists(&ops);
assert!(read_only.is_empty());
assert!(modified.is_empty());
}
// ── format_file_operations ───────────────────────────────────────
#[test]
fn format_file_operations_empty() {
assert_eq!(format_file_operations(&[], &[]), String::new());
}
#[test]
fn format_file_operations_read_only() {
let result = format_file_operations(&["src/main.rs".to_string()], &[]);
assert!(result.contains("<read-files>"));
assert!(result.contains("src/main.rs"));
assert!(!result.contains("<modified-files>"));
}
#[test]
fn format_file_operations_both() {
let result = format_file_operations(&["a.rs".to_string()], &["b.rs".to_string()]);
assert!(result.contains("<read-files>"));
assert!(result.contains("a.rs"));
assert!(result.contains("<modified-files>"));
assert!(result.contains("b.rs"));
}
// ── compaction_details_to_value ──────────────────────────────────
#[test]
fn compaction_details_serializes() {
let details = CompactionDetails {
read_files: vec!["a.rs".to_string()],
modified_files: vec!["b.rs".to_string()],
};
let value = compaction_details_to_value(&details).unwrap();
assert_eq!(value["readFiles"], json!(["a.rs"]));
assert_eq!(value["modifiedFiles"], json!(["b.rs"]));
}
// ── ResolvedCompactionSettings default ───────────────────────────
#[test]
fn default_settings() {
let settings = ResolvedCompactionSettings::default();
assert!(settings.enabled);
assert_eq!(settings.reserve_tokens, 16_384);
assert_eq!(settings.keep_recent_tokens, 20_000);
}
// ── Helper: entry constructors ──────────────────────────────────
use crate::model::{ImageContent, ThinkingContent};
use crate::session::{
BranchSummaryEntry, CompactionEntry, EntryBase, MessageEntry, ModelChangeEntry,
};
use std::collections::HashMap;
fn test_base(id: &str) -> EntryBase {
EntryBase {
id: Some(id.to_string()),
parent_id: None,
timestamp: "2026-01-01T00:00:00.000Z".to_string(),
}
}
fn user_entry(id: &str, text: &str) -> SessionEntry {
SessionEntry::Message(MessageEntry {
base: test_base(id),
message: make_user_text(text),
})
}
fn assistant_entry(id: &str, text: &str, input: u64, output: u64) -> SessionEntry {
SessionEntry::Message(MessageEntry {
base: test_base(id),
message: make_assistant_text(text, input, output),
})
}
fn tool_call_entry(id: &str, tool_name: &str, path: &str) -> SessionEntry {
SessionEntry::Message(MessageEntry {
base: test_base(id),
message: make_assistant_tool_call(tool_name, json!({"path": path})),
})
}
fn tool_result_entry(id: &str, text: &str) -> SessionEntry {
SessionEntry::Message(MessageEntry {
base: test_base(id),
message: make_tool_result(text),
})
}
fn branch_entry(id: &str, summary: &str) -> SessionEntry {
SessionEntry::BranchSummary(BranchSummaryEntry {
base: test_base(id),
from_id: "parent".to_string(),
summary: summary.to_string(),
details: None,
from_hook: None,
})
}
fn compact_entry(id: &str, summary: &str, tokens: u64) -> SessionEntry {
SessionEntry::Compaction(CompactionEntry {
base: test_base(id),
summary: summary.to_string(),
first_kept_entry_id: "kept".to_string(),
tokens_before: tokens,
details: None,
from_hook: None,
})
}
fn bash_entry(id: &str) -> SessionEntry {
SessionEntry::Message(MessageEntry {
base: test_base(id),
message: SessionMessage::BashExecution {
command: "ls".to_string(),
output: "ok".to_string(),
exit_code: 0,
cancelled: None,
truncated: None,
full_output_path: None,
timestamp: None,
extra: HashMap::new(),
},
})
}
// ── get_assistant_usage ─────────────────────────────────────────
#[test]
fn get_assistant_usage_returns_usage_for_stop() {
let msg = make_assistant_text("text", 100, 50);
let usage = get_assistant_usage(&msg);
assert!(usage.is_some());
assert_eq!(usage.unwrap().input, 100);
}
#[test]
fn get_assistant_usage_none_for_aborted() {
let msg = SessionMessage::Assistant {
message: AssistantMessage {
content: vec![ContentBlock::Text(TextContent::new("text"))],
api: String::new(),
provider: String::new(),
model: String::new(),
stop_reason: StopReason::Aborted,
error_message: None,
timestamp: 0,
usage: Usage {
input: 100,
output: 50,
total_tokens: 150,
..Default::default()
},
},
};
assert!(get_assistant_usage(&msg).is_none());
}
#[test]
fn get_assistant_usage_none_for_error() {
let msg = SessionMessage::Assistant {
message: AssistantMessage {
content: vec![],
api: String::new(),
provider: String::new(),
model: String::new(),
stop_reason: StopReason::Error,
error_message: None,
timestamp: 0,
usage: Usage::default(),
},
};
assert!(get_assistant_usage(&msg).is_none());
}
#[test]
fn get_assistant_usage_none_for_user() {
assert!(get_assistant_usage(&make_user_text("hello")).is_none());
}
// ── entry_is_message_like ───────────────────────────────────────
#[test]
fn entry_is_message_like_for_message() {
assert!(entry_is_message_like(&user_entry("1", "hi")));
}
#[test]
fn entry_is_message_like_for_branch_summary() {
assert!(entry_is_message_like(&branch_entry("1", "sum")));
}
#[test]
fn entry_is_message_like_false_for_compaction() {
assert!(!entry_is_message_like(&compact_entry("1", "sum", 100)));
}
#[test]
fn entry_is_message_like_false_for_model_change() {
let entry = SessionEntry::ModelChange(ModelChangeEntry {
base: test_base("1"),
provider: "test".to_string(),
model_id: "model-1".to_string(),
});
assert!(!entry_is_message_like(&entry));
}
// ── entry_is_compaction_boundary ────────────────────────────────
#[test]
fn compaction_boundary_true_for_compaction() {
assert!(entry_is_compaction_boundary(&compact_entry(
"1", "sum", 100
)));
}
#[test]
fn compaction_boundary_false_for_message() {
assert!(!entry_is_compaction_boundary(&user_entry("1", "hi")));
}
#[test]
fn compaction_boundary_false_for_branch() {
assert!(!entry_is_compaction_boundary(&branch_entry("1", "sum")));
}
// ── is_user_turn_start ──────────────────────────────────────────
#[test]
fn user_turn_start_for_user() {
assert!(is_user_turn_start(&user_entry("1", "hello")));
}
#[test]
fn user_turn_start_for_branch() {
assert!(is_user_turn_start(&branch_entry("1", "summary")));
}
#[test]
fn user_turn_start_for_bash() {
assert!(is_user_turn_start(&bash_entry("1")));
}
#[test]
fn user_turn_start_false_for_assistant() {
assert!(!is_user_turn_start(&assistant_entry("1", "resp", 10, 5)));
}
#[test]
fn user_turn_start_false_for_tool_result() {
assert!(!is_user_turn_start(&tool_result_entry("1", "result")));
}
#[test]
fn user_turn_start_false_for_compaction() {
assert!(!is_user_turn_start(&compact_entry("1", "sum", 100)));
}
// ── message_from_entry ──────────────────────────────────────────
#[test]
fn message_from_entry_user() {
let entry = user_entry("1", "hello");
let msg = message_from_entry(&entry);
assert!(msg.is_some());
assert!(matches!(msg.unwrap(), SessionMessage::User { .. }));
}
#[test]
fn message_from_entry_branch_summary() {
let entry = branch_entry("1", "branch summary text");
let msg = message_from_entry(&entry).unwrap();
if let SessionMessage::BranchSummary { summary, from_id } = msg {
assert_eq!(summary, "branch summary text");
assert_eq!(from_id, "parent");
} else {
panic!("expected BranchSummary");
}
}
#[test]
fn message_from_entry_compaction() {
let entry = compact_entry("1", "compact summary", 500);
let msg = message_from_entry(&entry).unwrap();
if let SessionMessage::CompactionSummary {
summary,
tokens_before,
} = msg
{
assert_eq!(summary, "compact summary");
assert_eq!(tokens_before, 500);
} else {
panic!("expected CompactionSummary");
}
}
#[test]
fn message_from_entry_model_change_is_none() {
let entry = SessionEntry::ModelChange(ModelChangeEntry {
base: test_base("1"),
provider: "test".to_string(),
model_id: "model".to_string(),
});
assert!(message_from_entry(&entry).is_none());
}
// ── find_valid_cut_points ───────────────────────────────────────
#[test]
fn find_valid_cut_points_empty() {
assert!(find_valid_cut_points(&[], 0, 0).is_empty());
}
#[test]
fn find_valid_cut_points_skips_tool_results() {
let entries = vec![
user_entry("1", "hello"),
assistant_entry("2", "resp", 10, 5),
tool_result_entry("3", "result"),
user_entry("4", "follow up"),
];
let cuts = find_valid_cut_points(&entries, 0, entries.len());
assert!(cuts.contains(&0)); // user
assert!(cuts.contains(&1)); // assistant
assert!(!cuts.contains(&2)); // tool result excluded
assert!(cuts.contains(&3)); // user
}
#[test]
fn find_valid_cut_points_includes_branch_summary() {
let entries = vec![branch_entry("1", "summary"), user_entry("2", "hello")];
let cuts = find_valid_cut_points(&entries, 0, entries.len());
assert!(cuts.contains(&0));
assert!(cuts.contains(&1));
}
#[test]
fn find_valid_cut_points_respects_range() {
let entries = vec![
user_entry("1", "a"),
user_entry("2", "b"),
user_entry("3", "c"),
];
let cuts = find_valid_cut_points(&entries, 1, 2);
assert!(!cuts.contains(&0));
assert!(cuts.contains(&1));
assert!(!cuts.contains(&2));
}
// ── find_turn_start_index ───────────────────────────────────────
#[test]
fn find_turn_start_basic() {
let entries = vec![
user_entry("1", "hello"),
assistant_entry("2", "resp", 10, 5),
tool_result_entry("3", "result"),
];
assert_eq!(find_turn_start_index(&entries, 2, 0), Some(0));
}
#[test]
fn find_turn_start_at_self() {
let entries = vec![user_entry("1", "hello")];
assert_eq!(find_turn_start_index(&entries, 0, 0), Some(0));
}
#[test]
fn find_turn_start_none_no_user() {
let entries = vec![
assistant_entry("1", "resp", 10, 5),
tool_result_entry("2", "result"),
];
assert_eq!(find_turn_start_index(&entries, 1, 0), None);
}
#[test]
fn find_turn_start_respects_start_index() {
let entries = vec![
user_entry("1", "old"),
assistant_entry("2", "resp", 10, 5),
user_entry("3", "new"),
];
// start_index=2, so it should find user at 2
assert_eq!(find_turn_start_index(&entries, 2, 2), Some(2));
// start_index=2, looking back from 2, user at 1 is below start
assert_eq!(find_turn_start_index(&entries, 1, 2), None);
}
// ── serialize_conversation ───────────────────────────────────────
#[test]
fn serialize_conversation_user_text() {
let messages = vec![Message::User(crate::model::UserMessage {
content: UserContent::Text("hello world".to_string()),
timestamp: 0,
})];
assert_eq!(serialize_conversation(&messages), "[User]: hello world");
}
#[test]
fn serialize_conversation_empty() {
assert!(serialize_conversation(&[]).is_empty());
}
#[test]
fn serialize_conversation_skips_empty_user() {
let messages = vec![Message::User(crate::model::UserMessage {
content: UserContent::Text(String::new()),
timestamp: 0,
})];
assert!(serialize_conversation(&messages).is_empty());
}
#[test]
fn serialize_conversation_assistant_text() {
let messages = vec![Message::assistant(AssistantMessage {
content: vec![ContentBlock::Text(TextContent::new("response"))],
api: String::new(),
provider: String::new(),
model: String::new(),
usage: Usage::default(),
stop_reason: StopReason::Stop,
error_message: None,
timestamp: 0,
})];
assert!(serialize_conversation(&messages).contains("[Assistant]: response"));
}
#[test]
fn serialize_conversation_tool_calls() {
let messages = vec![Message::assistant(AssistantMessage {
content: vec![ContentBlock::ToolCall(ToolCall {
id: "c1".to_string(),
name: "read".to_string(),
arguments: json!({"path": "/main.rs"}),
thought_signature: None,
})],
api: String::new(),
provider: String::new(),
model: String::new(),
usage: Usage::default(),
stop_reason: StopReason::Stop,
error_message: None,
timestamp: 0,
})];
let result = serialize_conversation(&messages);
assert!(result.contains("[Assistant tool calls]: read("));
assert!(result.contains("path="));
}
#[test]
fn serialize_conversation_thinking() {
let messages = vec![Message::assistant(AssistantMessage {
content: vec![ContentBlock::Thinking(ThinkingContent {
thinking: "let me think".to_string(),
thinking_signature: None,
})],
api: String::new(),
provider: String::new(),
model: String::new(),
usage: Usage::default(),
stop_reason: StopReason::Stop,
error_message: None,
timestamp: 0,
})];
assert!(serialize_conversation(&messages).contains("[Assistant thinking]: let me think"));
}
#[test]
fn serialize_conversation_tool_result() {
let messages = vec![Message::tool_result(crate::model::ToolResultMessage {
tool_call_id: "c1".to_string(),
tool_name: "read".to_string(),
content: vec![ContentBlock::Text(TextContent::new("file contents"))],
details: None,
is_error: false,
timestamp: 0,
})];
assert!(serialize_conversation(&messages).contains("[Tool result]: file contents"));
}
// ── estimate_tokens additional ──────────────────────────────────
#[test]
fn estimate_tokens_image_block() {
let msg = SessionMessage::User {
content: UserContent::Blocks(vec![ContentBlock::Image(ImageContent {
data: "base64data".to_string(),
mime_type: "image/png".to_string(),
})]),
timestamp: None,
};
// Image = 3600 chars (IMAGE_CHAR_ESTIMATE) -> ceil(3600/3) = 1200
assert_eq!(estimate_tokens(&msg), 1200);
}
#[test]
fn estimate_tokens_thinking() {
let msg = SessionMessage::User {
content: UserContent::Blocks(vec![ContentBlock::Thinking(ThinkingContent {
thinking: "a".repeat(20),
thinking_signature: None,
})]),
timestamp: None,
};
// 20 chars -> ceil(20/3) = 7
assert_eq!(estimate_tokens(&msg), 7);
}
#[test]
fn estimate_tokens_bash_execution() {
let msg = SessionMessage::BashExecution {
command: "echo hi".to_string(),
output: "hi\n".to_string(),
exit_code: 0,
cancelled: None,
truncated: None,
full_output_path: None,
timestamp: None,
extra: HashMap::new(),
};
// 7 + 3 = 10 chars -> ceil(10/3) = 4
assert_eq!(estimate_tokens(&msg), 4);
}
#[test]
fn estimate_tokens_branch_summary() {
let msg = SessionMessage::BranchSummary {
summary: "a".repeat(40),
from_id: "id".to_string(),
};
// 40 chars -> ceil(40/3) = 14
assert_eq!(estimate_tokens(&msg), 14);
}
#[test]
fn estimate_tokens_compaction_summary() {
let msg = SessionMessage::CompactionSummary {
summary: "a".repeat(80),
tokens_before: 5000,
};
// 80 chars -> ceil(80/3) = 27
assert_eq!(estimate_tokens(&msg), 27);
}
// ── prepare_compaction ──────────────────────────────────────────
#[test]
fn prepare_compaction_empty() {
assert!(prepare_compaction(&[], ResolvedCompactionSettings::default()).is_none());
}
#[test]
fn prepare_compaction_last_is_compaction_returns_none() {
let entries = vec![user_entry("1", "hello"), compact_entry("2", "summary", 100)];
assert!(prepare_compaction(&entries, ResolvedCompactionSettings::default()).is_none());
}
#[test]
fn prepare_compaction_no_messages_to_summarize_returns_none() {
// Only non-message entries that produce no summarizable messages
let entries = vec![SessionEntry::ModelChange(ModelChangeEntry {
base: test_base("1"),
provider: "test".to_string(),
model_id: "model".to_string(),
})];
assert!(prepare_compaction(&entries, ResolvedCompactionSettings::default()).is_none());
}
#[test]
fn prepare_compaction_basic_returns_some() {
let long_text = "a".repeat(100_000);
let entries = vec![
user_entry("1", &long_text),
assistant_entry("2", &long_text, 50000, 25000),
user_entry("3", &long_text),
assistant_entry("4", &long_text, 80000, 30000),
user_entry("5", "recent"),
];
let settings = ResolvedCompactionSettings {
enabled: true,
context_window_tokens: 100_000,
reserve_tokens: 1000,
keep_recent_tokens: 100,
};
let prep = prepare_compaction(&entries, settings);
assert!(prep.is_some());
let p = prep.unwrap();
assert!(!p.messages_to_summarize.is_empty());
assert!(p.tokens_before > 0);
assert!(p.previous_summary.is_none());
}
#[test]
fn prepare_compaction_after_previous_compaction() {
let entries = vec![
user_entry("1", "old message"),
assistant_entry("2", "old response", 100, 50),
compact_entry("3", "previous summary", 300),
user_entry("4", &"x".repeat(100_000)),
assistant_entry("5", &"y".repeat(100_000), 80000, 30000),
user_entry("6", "recent"),
];
let settings = ResolvedCompactionSettings {
enabled: true,
context_window_tokens: 100_000,
reserve_tokens: 1000,
keep_recent_tokens: 100,
};
let prep = prepare_compaction(&entries, settings);
assert!(prep.is_some());
let p = prep.unwrap();
assert_eq!(p.previous_summary.as_deref(), Some("previous summary"));
}
#[test]
fn prepare_compaction_tracks_file_ops() {
let entries = vec![
tool_call_entry("1", "read", "/src/main.rs"),
tool_result_entry("1r", "ok"),
tool_call_entry("2", "edit", "/src/lib.rs"),
tool_result_entry("2r", "ok"),
user_entry("3", &"x".repeat(100_000)),
assistant_entry("4", &"y".repeat(100_000), 80000, 30000),
user_entry("5", "recent"),
];
let settings = ResolvedCompactionSettings {
enabled: true,
reserve_tokens: 1000,
keep_recent_tokens: 100,
..Default::default()
};
if let Some(prep) = prepare_compaction(&entries, settings) {
let has_read = prep.file_ops.read.contains("/src/main.rs");
let has_edit = prep.file_ops.edited.contains("/src/lib.rs");
// At least one should be tracked (depends on cut point position)
assert!(has_read || has_edit || prep.file_ops.read.is_empty());
}
}
// ── FileOperations::read_files ──────────────────────────────────
#[test]
fn file_operations_read_files_iterator() {
let mut ops = FileOperations::default();
ops.read.insert("/a.rs".to_string());
ops.read.insert("/b.rs".to_string());
let files: Vec<&str> = ops.read_files().collect();
assert_eq!(files.len(), 2);
assert!(files.contains(&"/a.rs"));
assert!(files.contains(&"/b.rs"));
}
#[test]
fn find_cut_point_includes_tool_result_when_needed() {
// Setup:
// 0. User (10)
// 1. Assistant Call (10)
// 2. Tool Result (100)
// 3. User (10)
// 4. Assistant (10)
//
// Keep recent = 100.
// Accumulation from end:
// 4: 10
// 3: 20
// 2: 120 (Threshold crossed at index 2)
//
// Index 2 is ToolResult (invalid cut point).
// Valid cut points: 0, 1, 3, 4.
//
// Logic should pick closest valid cut point <= 2, which is 1.
// If it picked >= 2, it would pick 3, discarding the ToolResult and Call (keeping only 20 tokens).
// By picking 1, we keep 1..4 (130 tokens).
// Create entries with controlled lengths.
// With chars/token ~=3, 400 chars => ceil(400/3)=134 tokens.
let tr_text = "x".repeat(400);
let entries = vec![
user_entry("0", "user"), // Valid
assistant_entry("1", "call", 10, 10), // Valid (Assistant)
tool_result_entry("2", &tr_text), // Invalid
user_entry("3", "user"), // Valid
assistant_entry("4", "resp", 10, 10), // Valid
];
// Verify token estimates (approx)
// 0: ceil(4/3) = 2
// 1: ceil(4/3) = 2
// 2: ceil(400/3) = 134
// 3: ceil(4/3) = 2
// 4: ceil(4/3) = 2
// Total recent needed: 100.
// Accumulate: 4(2)+3(2)+2(134) = 138. Crossed at 2.
let settings = ResolvedCompactionSettings {
enabled: true,
context_window_tokens: 15,
reserve_tokens: 0,
keep_recent_tokens: 100,
};
let prep = prepare_compaction(&entries, settings).expect("should compact");
// Cut point is index 1 (Assistant/Call). Because entries[1] is Assistant (not User),
// this is a split turn: the turn started at index 0 (User). The User message at index 0
// goes into turn_prefix_messages (not messages_to_summarize) because history_end = 0.
assert_eq!(prep.first_kept_entry_id, "1");
// messages_to_summarize is entries[0..0] = empty (split-turn puts the
// prefix in turn_prefix_messages instead).
assert!(
prep.messages_to_summarize.is_empty(),
"split turn: user goes into turn prefix, not summarize"
);
// turn_prefix_messages should contain the User message at index 0.
assert_eq!(prep.turn_prefix_messages.len(), 1);
match &prep.turn_prefix_messages[0] {
SessionMessage::User { content, .. } => {
if let UserContent::Text(t) = content {
assert_eq!(t, "user");
} else {
panic!("wrong content");
}
}
_ => panic!("expected user message in turn prefix"),
}
}
#[test]
fn find_cut_point_should_not_discard_context_to_skip_tool_chain() {
// Setup (estimate_tokens uses ceil(chars/3)):
// 0. User "x"*4000 → 1334 tokens
// 1. Assistant "x"*400 → 134 tokens
// 2. Tool Result "x"*400 → 134 tokens
// 3. User "next" → 2 tokens
//
// Keep recent = 150.
// Accumulation (from end):
// 3: 2
// 2: 136
// 1: 270 (Crosses 150) -> cut_index = 1
//
// The cut should land at index 1 (the assistant message), keeping
// entries 1-3 and summarizing only entry 0.
let entries = vec![
user_entry("0", &"x".repeat(4000)), // 1000 tokens
assistant_entry("1", &"x".repeat(400), 50, 50), // 100 tokens
tool_result_entry("2", &"x".repeat(400)), // 100 tokens
user_entry("3", "next"), // 1 token
];
let settings = ResolvedCompactionSettings {
enabled: true,
context_window_tokens: 200,
reserve_tokens: 0,
keep_recent_tokens: 150,
};
// We use prepare_compaction as the entry point
let prep = prepare_compaction(&entries, settings).expect("should compact");
// We expect to keep from 1 (Assistant). The cut splits the turn
// (user 0 + assistant 1), so user 0 goes into the turn prefix.
assert_eq!(
prep.first_kept_entry_id, "1",
"Should start at Assistant message to preserve context"
);
assert!(
prep.is_split_turn,
"Cut should split the user/assistant turn"
);
assert_eq!(
prep.turn_prefix_messages.len(),
1,
"User entry at index 0 should be in the turn prefix"
);
assert!(
prep.messages_to_summarize.is_empty(),
"Nothing before the turn to summarize"
);
}
mod proptest_compaction {
use super::*;
use proptest::prelude::*;
proptest! {
/// `calculate_context_tokens`: if total > 0, returns total.
#[test]
fn calc_context_tokens_total_wins(
input in 0..1_000_000u64,
output in 0..1_000_000u64,
total in 1..2_000_000u64,
) {
let usage = Usage {
input,
output,
total_tokens: total,
..Usage::default()
};
assert_eq!(calculate_context_tokens(&usage), total);
}
/// `calculate_context_tokens`: if total == 0, returns input + output.
#[test]
fn calc_context_tokens_fallback(
input in 0..1_000_000u64,
output in 0..1_000_000u64,
) {
let usage = Usage {
input,
output,
total_tokens: 0,
..Usage::default()
};
assert_eq!(calculate_context_tokens(&usage), input + output);
}
/// `should_compact` returns false when disabled.
#[test]
fn should_compact_disabled_returns_false(
ctx_tokens in 0..1_000_000u64,
window in 0..500_000u32,
) {
let settings = ResolvedCompactionSettings {
enabled: false,
context_window_tokens: window,
reserve_tokens: 16_384,
keep_recent_tokens: 20_000,
};
assert!(!should_compact(ctx_tokens, window, &settings));
}
/// `should_compact` threshold: tokens > window - reserve.
#[test]
fn should_compact_threshold(
ctx_tokens in 0..500_000u64,
window in 0..300_000u32,
reserve in 0..100_000u32,
) {
let settings = ResolvedCompactionSettings {
enabled: true,
context_window_tokens: window,
reserve_tokens: reserve,
keep_recent_tokens: 20_000,
};
let threshold = u64::from(window).saturating_sub(u64::from(reserve));
let result = should_compact(ctx_tokens, window, &settings);
assert_eq!(result, ctx_tokens > threshold);
}
/// `format_file_operations`: empty lists produce empty string.
#[test]
fn format_file_ops_empty(_dummy in 0..10u32) {
let result = format_file_operations(&[], &[]);
assert!(result.is_empty());
}
/// `format_file_operations`: read files produce `<read-files>` tag.
#[test]
fn format_file_ops_read_tag(
files in prop::collection::vec("[a-z./]{1,20}", 1..5),
) {
let result = format_file_operations(&files, &[]);
assert!(result.contains("<read-files>"));
assert!(result.contains("</read-files>"));
assert!(!result.contains("<modified-files>"));
for f in &files {
assert!(result.contains(f.as_str()));
}
}
/// `format_file_operations`: modified files produce `<modified-files>` tag.
#[test]
fn format_file_ops_modified_tag(
files in prop::collection::vec("[a-z./]{1,20}", 1..5),
) {
let result = format_file_operations(&[], &files);
assert!(!result.contains("<read-files>"));
assert!(result.contains("<modified-files>"));
assert!(result.contains("</modified-files>"));
for f in &files {
assert!(result.contains(f.as_str()));
}
}
/// `compute_file_lists`: modified = edited ∪ written, read_only = read \ modified.
#[test]
fn compute_file_lists_set_algebra(
read in prop::collection::hash_set("[a-z]{1,5}", 0..5),
written in prop::collection::hash_set("[a-z]{1,5}", 0..5),
edited in prop::collection::hash_set("[a-z]{1,5}", 0..5),
) {
let file_ops = FileOperations {
read: read.clone(),
written: written.clone(),
edited: edited.clone(),
};
let (read_only, modified) = compute_file_lists(&file_ops);
// Modified = edited ∪ written
let expected_modified: HashSet<&String> =
edited.iter().chain(written.iter()).collect();
let actual_modified: HashSet<&String> = modified.iter().collect();
assert_eq!(actual_modified, expected_modified);
// Read-only = read \ modified (no overlap)
for f in &read_only {
assert!(!modified.contains(f), "overlap: {f}");
assert!(read.contains(f));
}
// Both are sorted
for pair in read_only.windows(2) {
assert!(pair[0] <= pair[1]);
}
for pair in modified.windows(2) {
assert!(pair[0] <= pair[1]);
}
}
}
}
}