1use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub struct CompactionConfig {
5 pub preserve_recent_messages: usize,
6 pub max_estimated_tokens: usize,
7}
8
9impl Default for CompactionConfig {
10 fn default() -> Self {
11 Self {
12 preserve_recent_messages: 4,
13 max_estimated_tokens: 10_000,
14 }
15 }
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct CompactionResult {
20 pub summary: String,
21 pub formatted_summary: String,
22 pub compacted_session: Session,
23 pub removed_message_count: usize,
24}
25
26#[must_use]
27pub fn estimate_session_tokens(session: &Session) -> usize {
28 session.messages.iter().map(estimate_message_tokens).sum()
29}
30
31#[must_use]
32pub fn should_compact(session: &Session, config: CompactionConfig) -> bool {
33 session.messages.len() > config.preserve_recent_messages
34 && estimate_session_tokens(session) >= config.max_estimated_tokens
35}
36
37#[must_use]
38pub fn format_compact_summary(summary: &str) -> String {
39 let without_analysis = strip_tag_block(summary, "analysis");
40 let formatted = if let Some(content) = extract_tag_block(&without_analysis, "summary") {
41 without_analysis.replace(
42 &format!("<summary>{content}</summary>"),
43 &format!("Summary:\n{}", content.trim()),
44 )
45 } else {
46 without_analysis
47 };
48
49 collapse_blank_lines(&formatted).trim().to_string()
50}
51
52#[must_use]
53pub fn get_compact_continuation_message(
54 summary: &str,
55 suppress_follow_up_questions: bool,
56 recent_messages_preserved: bool,
57) -> String {
58 let mut base = format!(
59 "This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.\n\n{}",
60 format_compact_summary(summary)
61 );
62
63 if recent_messages_preserved {
64 base.push_str("\n\nRecent messages are preserved verbatim.");
65 }
66
67 if suppress_follow_up_questions {
68 base.push_str("\nContinue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, and do not preface with continuation text.");
69 }
70
71 base
72}
73
74#[must_use]
75pub fn compact_session(session: &Session, config: CompactionConfig) -> CompactionResult {
76 if !should_compact(session, config) {
77 return CompactionResult {
78 summary: String::new(),
79 formatted_summary: String::new(),
80 compacted_session: session.clone(),
81 removed_message_count: 0,
82 };
83 }
84
85 let keep_from = session
86 .messages
87 .len()
88 .saturating_sub(config.preserve_recent_messages);
89 let removed = &session.messages[..keep_from];
90 let preserved = session.messages[keep_from..].to_vec();
91 let summary = summarize_messages(removed);
92 let formatted_summary = format_compact_summary(&summary);
93 let continuation = get_compact_continuation_message(&summary, true, !preserved.is_empty());
94
95 let mut compacted_messages = vec![ConversationMessage {
96 role: MessageRole::System,
97 blocks: vec![ContentBlock::Text { text: continuation }],
98 usage: None,
99 }];
100 compacted_messages.extend(preserved);
101
102 CompactionResult {
103 summary,
104 formatted_summary,
105 compacted_session: Session {
106 version: session.version,
107 messages: compacted_messages,
108 },
109 removed_message_count: removed.len(),
110 }
111}
112
113fn summarize_messages(messages: &[ConversationMessage]) -> String {
114 let user_messages = messages
115 .iter()
116 .filter(|message| message.role == MessageRole::User)
117 .count();
118 let assistant_messages = messages
119 .iter()
120 .filter(|message| message.role == MessageRole::Assistant)
121 .count();
122 let tool_messages = messages
123 .iter()
124 .filter(|message| message.role == MessageRole::Tool)
125 .count();
126
127 let mut tool_names = messages
128 .iter()
129 .flat_map(|message| message.blocks.iter())
130 .filter_map(|block| match block {
131 ContentBlock::ToolUse { name, .. } => Some(name.as_str()),
132 ContentBlock::ToolResult { tool_name, .. } => Some(tool_name.as_str()),
133 ContentBlock::Text { .. } | ContentBlock::Image { .. } => None,
134 })
135 .collect::<Vec<_>>();
136 tool_names.sort_unstable();
137 tool_names.dedup();
138
139 let mut lines = vec![
140 "<summary>".to_string(),
141 "Conversation summary:".to_string(),
142 format!(
143 "- Scope: {} earlier messages compacted (user={}, assistant={}, tool={}).",
144 messages.len(),
145 user_messages,
146 assistant_messages,
147 tool_messages
148 ),
149 ];
150
151 if !tool_names.is_empty() {
152 lines.push(format!("- Tools mentioned: {}.", tool_names.join(", ")));
153 }
154
155 let recent_user_requests = collect_recent_role_summaries(messages, MessageRole::User, 3);
156 if !recent_user_requests.is_empty() {
157 lines.push("- Recent user requests:".to_string());
158 lines.extend(
159 recent_user_requests
160 .into_iter()
161 .map(|request| format!(" - {request}")),
162 );
163 }
164
165 let pending_work = infer_pending_work(messages);
166 if !pending_work.is_empty() {
167 lines.push("- Pending work:".to_string());
168 lines.extend(pending_work.into_iter().map(|item| format!(" - {item}")));
169 }
170
171 let key_files = collect_key_files(messages);
172 if !key_files.is_empty() {
173 lines.push(format!("- Key files referenced: {}.", key_files.join(", ")));
174 }
175
176 if let Some(current_work) = infer_current_work(messages) {
177 lines.push(format!("- Current work: {current_work}"));
178 }
179
180 lines.push("- Key timeline:".to_string());
181 for message in messages {
182 let role = match message.role {
183 MessageRole::System => "system",
184 MessageRole::User => "user",
185 MessageRole::Assistant => "assistant",
186 MessageRole::Tool => "tool",
187 };
188 let content = message
189 .blocks
190 .iter()
191 .map(summarize_block)
192 .collect::<Vec<_>>()
193 .join(" | ");
194 lines.push(format!(" - {role}: {content}"));
195 }
196 lines.push("</summary>".to_string());
197 lines.join("\n")
198}
199
200fn summarize_block(block: &ContentBlock) -> String {
201 let raw = match block {
202 ContentBlock::Text { text } => text.clone(),
203 ContentBlock::Image { media_type, .. } => format!("[image: {media_type}]"),
204 ContentBlock::ToolUse { name, input, .. } => format!("tool_use {name}({input})"),
205 ContentBlock::ToolResult {
206 tool_name,
207 output,
208 is_error,
209 ..
210 } => format!(
211 "tool_result {tool_name}: {}{output}",
212 if *is_error { "error " } else { "" }
213 ),
214 };
215 truncate_summary(&raw, 160)
216}
217
218fn collect_recent_role_summaries(
219 messages: &[ConversationMessage],
220 role: MessageRole,
221 limit: usize,
222) -> Vec<String> {
223 messages
224 .iter()
225 .filter(|message| message.role == role)
226 .rev()
227 .filter_map(|message| first_text_block(message))
228 .take(limit)
229 .map(|text| truncate_summary(text, 160))
230 .collect::<Vec<_>>()
231 .into_iter()
232 .rev()
233 .collect()
234}
235
236fn infer_pending_work(messages: &[ConversationMessage]) -> Vec<String> {
237 messages
238 .iter()
239 .rev()
240 .filter_map(first_text_block)
241 .filter(|text| {
242 let lowered = text.to_ascii_lowercase();
243 lowered.contains("todo")
244 || lowered.contains("next")
245 || lowered.contains("pending")
246 || lowered.contains("follow up")
247 || lowered.contains("remaining")
248 })
249 .take(3)
250 .map(|text| truncate_summary(text, 160))
251 .collect::<Vec<_>>()
252 .into_iter()
253 .rev()
254 .collect()
255}
256
257fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
258 let mut files = messages
259 .iter()
260 .flat_map(|message| message.blocks.iter())
261 .filter_map(|block| match block {
262 ContentBlock::Text { text } => Some(text.as_str()),
263 ContentBlock::ToolUse { input, .. } => Some(input.as_str()),
264 ContentBlock::ToolResult { output, .. } => Some(output.as_str()),
265 ContentBlock::Image { .. } => None,
266 })
267 .flat_map(extract_file_candidates)
268 .collect::<Vec<_>>();
269 files.sort();
270 files.dedup();
271 files.into_iter().take(8).collect()
272}
273
274fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
275 messages
276 .iter()
277 .rev()
278 .filter_map(first_text_block)
279 .find(|text| !text.trim().is_empty())
280 .map(|text| truncate_summary(text, 200))
281}
282
283fn first_text_block(message: &ConversationMessage) -> Option<&str> {
284 message.blocks.iter().find_map(|block| match block {
285 ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
286 ContentBlock::ToolUse { .. }
287 | ContentBlock::ToolResult { .. }
288 | ContentBlock::Text { .. }
289 | ContentBlock::Image { .. } => None,
290 })
291}
292
293fn has_interesting_extension(candidate: &str) -> bool {
294 std::path::Path::new(candidate)
295 .extension()
296 .and_then(|extension| extension.to_str())
297 .is_some_and(|extension| {
298 ["rs", "ts", "tsx", "js", "json", "md"]
299 .iter()
300 .any(|expected| extension.eq_ignore_ascii_case(expected))
301 })
302}
303
304fn extract_file_candidates(content: &str) -> Vec<String> {
305 content
306 .split_whitespace()
307 .filter_map(|token| {
308 let candidate = token.trim_matches(|char: char| {
309 matches!(char, ',' | '.' | ':' | ';' | ')' | '(' | '"' | '\'' | '`')
310 });
311 if candidate.contains('/') && has_interesting_extension(candidate) {
312 Some(candidate.to_string())
313 } else {
314 None
315 }
316 })
317 .collect()
318}
319
320fn truncate_summary(content: &str, max_chars: usize) -> String {
321 if content.chars().count() <= max_chars {
322 return content.to_string();
323 }
324 let mut truncated = content.chars().take(max_chars).collect::<String>();
325 truncated.push('…');
326 truncated
327}
328
329fn estimate_message_tokens(message: &ConversationMessage) -> usize {
330 message
331 .blocks
332 .iter()
333 .map(|block| match block {
334 ContentBlock::Text { text } => text.len() / 4 + 1,
335 ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
336 ContentBlock::ToolResult { tool_name, output, .. } => (tool_name.len() + output.len()) / 4 + 1,
337 ContentBlock::Image { data, .. } => data.len() / 4 + 1,
338 })
339 .sum()
340}
341
342fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
343 let start = format!("<{tag}>");
344 let end = format!("</{tag}>");
345 let start_index = content.find(&start)? + start.len();
346 let end_index = content[start_index..].find(&end)? + start_index;
347 Some(content[start_index..end_index].to_string())
348}
349
350fn strip_tag_block(content: &str, tag: &str) -> String {
351 let start = format!("<{tag}>");
352 let end = format!("</{tag}>");
353 if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
354 let end_index = end_index_rel + end.len();
355 let mut stripped = String::new();
356 stripped.push_str(&content[..start_index]);
357 stripped.push_str(&content[end_index..]);
358 stripped
359 } else {
360 content.to_string()
361 }
362}
363
364fn collapse_blank_lines(content: &str) -> String {
365 let mut result = String::new();
366 let mut last_blank = false;
367 for line in content.lines() {
368 let is_blank = line.trim().is_empty();
369 if is_blank && last_blank {
370 continue;
371 }
372 result.push_str(line);
373 result.push('\n');
374 last_blank = is_blank;
375 }
376 result
377}
378
379#[cfg(test)]
380mod tests {
381 use super::{
382 collect_key_files, compact_session, estimate_session_tokens, format_compact_summary,
383 infer_pending_work, should_compact, CompactionConfig,
384 };
385 use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
386
387 #[test]
388 fn formats_compact_summary_like_upstream() {
389 let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
390 assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
391 }
392
393 #[test]
394 fn leaves_small_sessions_unchanged() {
395 let session = Session {
396 version: 1,
397 messages: vec![ConversationMessage::user_text("hello")],
398 };
399
400 let result = compact_session(&session, CompactionConfig::default());
401 assert_eq!(result.removed_message_count, 0);
402 assert_eq!(result.compacted_session, session);
403 assert!(result.summary.is_empty());
404 assert!(result.formatted_summary.is_empty());
405 }
406
407 #[test]
408 fn compacts_older_messages_into_a_system_summary() {
409 let session = Session {
410 version: 1,
411 messages: vec![
412 ConversationMessage::user_text("one ".repeat(200)),
413 ConversationMessage::assistant(vec![ContentBlock::Text {
414 text: "two ".repeat(200),
415 }]),
416 ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
417 ConversationMessage {
418 role: MessageRole::Assistant,
419 blocks: vec![ContentBlock::Text {
420 text: "recent".to_string(),
421 }],
422 usage: None,
423 },
424 ],
425 };
426
427 let result = compact_session(
428 &session,
429 CompactionConfig {
430 preserve_recent_messages: 2,
431 max_estimated_tokens: 1,
432 },
433 );
434
435 assert_eq!(result.removed_message_count, 2);
436 assert_eq!(
437 result.compacted_session.messages[0].role,
438 MessageRole::System
439 );
440 assert!(matches!(
441 &result.compacted_session.messages[0].blocks[0],
442 ContentBlock::Text { text } if text.contains("Summary:")
443 ));
444 assert!(result.formatted_summary.contains("Scope:"));
445 assert!(result.formatted_summary.contains("Key timeline:"));
446 assert!(should_compact(
447 &session,
448 CompactionConfig {
449 preserve_recent_messages: 2,
450 max_estimated_tokens: 1,
451 }
452 ));
453 assert!(
454 estimate_session_tokens(&result.compacted_session) < estimate_session_tokens(&session)
455 );
456 }
457
458 #[test]
459 fn truncates_long_blocks_in_summary() {
460 let summary = super::summarize_block(&ContentBlock::Text {
461 text: "x".repeat(400),
462 });
463 assert!(summary.ends_with('…'));
464 assert!(summary.chars().count() <= 161);
465 }
466
467 #[test]
468 fn extracts_key_files_from_message_content() {
469 let files = collect_key_files(&[ConversationMessage::user_text(
470 "Update rust/crates/runtime/src/compact.rs and rust/crates/rusty-ternlang-cli/src/main.rs next.",
471 )]);
472 assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
473 assert!(files.contains(&"rust/crates/rusty-ternlang-cli/src/main.rs".to_string()));
474 }
475
476 #[test]
477 fn infers_pending_work_from_recent_messages() {
478 let pending = infer_pending_work(&[
479 ConversationMessage::user_text("done"),
480 ConversationMessage::assistant(vec![ContentBlock::Text {
481 text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
482 }]),
483 ]);
484 assert_eq!(pending.len(), 1);
485 assert!(pending[0].contains("Next: update tests"));
486 }
487}