1use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
2
3const COMPACT_CONTINUATION_PREAMBLE: &str =
4 "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";
5const COMPACT_RECENT_MESSAGES_NOTE: &str = "Recent messages are preserved verbatim.";
6const COMPACT_DIRECT_RESUME_INSTRUCTION: &str = "Continue 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.";
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub struct CompactionConfig {
10 pub preserve_recent_messages: usize,
11 pub max_estimated_tokens: usize,
12}
13
14impl Default for CompactionConfig {
15 fn default() -> Self {
16 Self {
17 preserve_recent_messages: 4,
18 max_estimated_tokens: 10_000,
19 }
20 }
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct CompactionResult {
25 pub summary: String,
26 pub formatted_summary: String,
27 pub compacted_session: Session,
28 pub removed_message_count: usize,
29}
30
31#[must_use]
32pub fn estimate_session_tokens(session: &Session) -> usize {
33 session.messages.iter().map(estimate_message_tokens).sum()
34}
35
36#[must_use]
37pub fn should_compact(session: &Session, config: CompactionConfig) -> bool {
38 let start = compacted_summary_prefix_len(session);
39 let compactable = &session.messages[start..];
40
41 compactable.len() > config.preserve_recent_messages
42 && compactable
43 .iter()
44 .map(estimate_message_tokens)
45 .sum::<usize>()
46 >= config.max_estimated_tokens
47}
48
49#[must_use]
50pub fn format_compact_summary(summary: &str) -> String {
51 let without_analysis = strip_tag_block(summary, "analysis");
52 let formatted = if let Some(content) = extract_tag_block(&without_analysis, "summary") {
53 without_analysis.replace(
54 &format!("<summary>{content}</summary>"),
55 &format!("Summary:\n{}", content.trim()),
56 )
57 } else {
58 without_analysis
59 };
60
61 collapse_blank_lines(&formatted).trim().to_string()
62}
63
64#[must_use]
65pub fn get_compact_continuation_message(
66 summary: &str,
67 suppress_follow_up_questions: bool,
68 recent_messages_preserved: bool,
69) -> String {
70 let mut base = format!(
71 "{COMPACT_CONTINUATION_PREAMBLE}{}",
72 format_compact_summary(summary)
73 );
74
75 if recent_messages_preserved {
76 base.push_str("\n\n");
77 base.push_str(COMPACT_RECENT_MESSAGES_NOTE);
78 }
79
80 if suppress_follow_up_questions {
81 base.push('\n');
82 base.push_str(COMPACT_DIRECT_RESUME_INSTRUCTION);
83 }
84
85 base
86}
87
88#[must_use]
89pub fn compact_session(session: &Session, config: CompactionConfig) -> CompactionResult {
90 if !should_compact(session, config) {
91 return CompactionResult {
92 summary: String::new(),
93 formatted_summary: String::new(),
94 compacted_session: session.clone(),
95 removed_message_count: 0,
96 };
97 }
98
99 let existing_summary = session
100 .messages
101 .first()
102 .and_then(extract_existing_compacted_summary);
103 let compacted_prefix_len = usize::from(existing_summary.is_some());
104 let keep_from = session
105 .messages
106 .len()
107 .saturating_sub(config.preserve_recent_messages);
108 let removed = &session.messages[compacted_prefix_len..keep_from];
109 let preserved = session.messages[keep_from..].to_vec();
110 let summary =
111 merge_compact_summaries(existing_summary.as_deref(), &summarize_messages(removed));
112 let formatted_summary = format_compact_summary(&summary);
113 let continuation = get_compact_continuation_message(&summary, true, !preserved.is_empty());
114
115 let mut compacted_messages = vec![ConversationMessage {
116 role: MessageRole::System,
117 blocks: vec![ContentBlock::Text { text: continuation }],
118 usage: None,
119 }];
120 compacted_messages.extend(preserved);
121
122 CompactionResult {
123 summary,
124 formatted_summary,
125 compacted_session: Session {
126 version: session.version,
127 messages: compacted_messages,
128 },
129 removed_message_count: removed.len(),
130 }
131}
132
133fn compacted_summary_prefix_len(session: &Session) -> usize {
134 usize::from(
135 session
136 .messages
137 .first()
138 .and_then(extract_existing_compacted_summary)
139 .is_some(),
140 )
141}
142
143fn summarize_messages(messages: &[ConversationMessage]) -> String {
144 let user_messages = messages
145 .iter()
146 .filter(|message| message.role == MessageRole::User)
147 .count();
148 let assistant_messages = messages
149 .iter()
150 .filter(|message| message.role == MessageRole::Assistant)
151 .count();
152 let tool_messages = messages
153 .iter()
154 .filter(|message| message.role == MessageRole::Tool)
155 .count();
156
157 let mut tool_names = messages
158 .iter()
159 .flat_map(|message| message.blocks.iter())
160 .filter_map(|block| match block {
161 ContentBlock::ToolUse { name, .. } => Some(name.as_str()),
162 ContentBlock::ToolResult { tool_name, .. } => Some(tool_name.as_str()),
163 ContentBlock::Text { .. } | ContentBlock::Image { .. } => None,
164 })
165 .collect::<Vec<_>>();
166 tool_names.sort_unstable();
167 tool_names.dedup();
168
169 let mut lines = vec![
170 "<summary>".to_string(),
171 "Conversation summary:".to_string(),
172 format!(
173 "- Scope: {} earlier messages compacted (user={}, assistant={}, tool={}).",
174 messages.len(),
175 user_messages,
176 assistant_messages,
177 tool_messages
178 ),
179 ];
180
181 if !tool_names.is_empty() {
182 lines.push(format!("- Tools mentioned: {}.", tool_names.join(", ")));
183 }
184
185 let recent_user_requests = collect_recent_role_summaries(messages, MessageRole::User, 3);
186 if !recent_user_requests.is_empty() {
187 lines.push("- Recent user requests:".to_string());
188 lines.extend(
189 recent_user_requests
190 .into_iter()
191 .map(|request| format!(" - {request}")),
192 );
193 }
194
195 let pending_work = infer_pending_work(messages);
196 if !pending_work.is_empty() {
197 lines.push("- Pending work:".to_string());
198 lines.extend(pending_work.into_iter().map(|item| format!(" - {item}")));
199 }
200
201 let key_files = collect_key_files(messages);
202 if !key_files.is_empty() {
203 lines.push(format!("- Key files referenced: {}.", key_files.join(", ")));
204 }
205
206 if let Some(current_work) = infer_current_work(messages) {
207 lines.push(format!("- Current work: {current_work}"));
208 }
209
210 lines.push("- Key timeline:".to_string());
211 for message in messages {
212 let role = match message.role {
213 MessageRole::System => "system",
214 MessageRole::User => "user",
215 MessageRole::Assistant => "assistant",
216 MessageRole::Tool => "tool",
217 };
218 let content = message
219 .blocks
220 .iter()
221 .map(summarize_block)
222 .collect::<Vec<_>>()
223 .join(" | ");
224 lines.push(format!(" - {role}: {content}"));
225 }
226 lines.push("</summary>".to_string());
227 lines.join("\n")
228}
229
230fn merge_compact_summaries(existing_summary: Option<&str>, new_summary: &str) -> String {
231 let Some(existing_summary) = existing_summary else {
232 return new_summary.to_string();
233 };
234
235 let previous_highlights = extract_summary_highlights(existing_summary);
236 let new_formatted_summary = format_compact_summary(new_summary);
237 let new_highlights = extract_summary_highlights(&new_formatted_summary);
238 let new_timeline = extract_summary_timeline(&new_formatted_summary);
239
240 let mut lines = vec!["<summary>".to_string(), "Conversation summary:".to_string()];
241
242 if !previous_highlights.is_empty() {
243 lines.push("- Previously compacted context:".to_string());
244 lines.extend(
245 previous_highlights
246 .into_iter()
247 .map(|line| format!(" {line}")),
248 );
249 }
250
251 if !new_highlights.is_empty() {
252 lines.push("- Newly compacted context:".to_string());
253 lines.extend(new_highlights.into_iter().map(|line| format!(" {line}")));
254 }
255
256 if !new_timeline.is_empty() {
257 lines.push("- Key timeline:".to_string());
258 lines.extend(new_timeline.into_iter().map(|line| format!(" {line}")));
259 }
260
261 lines.push("</summary>".to_string());
262 lines.join("\n")
263}
264
265fn summarize_block(block: &ContentBlock) -> String {
266 let raw = match block {
267 ContentBlock::Text { text } => text.clone(),
268 ContentBlock::Image { media_type, .. } => format!("[image: {media_type}]"),
269 ContentBlock::ToolUse { name, input, .. } => format!("tool_use {name}({input})"),
270 ContentBlock::ToolResult {
271 tool_name,
272 output,
273 is_error,
274 ..
275 } => format!(
276 "tool_result {tool_name}: {}{output}",
277 if *is_error { "error " } else { "" }
278 ),
279 };
280 truncate_summary(&raw, 160)
281}
282
283fn collect_recent_role_summaries(
284 messages: &[ConversationMessage],
285 role: MessageRole,
286 limit: usize,
287) -> Vec<String> {
288 messages
289 .iter()
290 .filter(|message| message.role == role)
291 .rev()
292 .filter_map(|message| first_text_block(message))
293 .take(limit)
294 .map(|text| truncate_summary(text, 160))
295 .collect::<Vec<_>>()
296 .into_iter()
297 .rev()
298 .collect()
299}
300
301fn infer_pending_work(messages: &[ConversationMessage]) -> Vec<String> {
302 messages
303 .iter()
304 .rev()
305 .filter_map(first_text_block)
306 .filter(|text| {
307 let lowered = text.to_ascii_lowercase();
308 lowered.contains("todo")
309 || lowered.contains("next")
310 || lowered.contains("pending")
311 || lowered.contains("follow up")
312 || lowered.contains("remaining")
313 })
314 .take(3)
315 .map(|text| truncate_summary(text, 160))
316 .collect::<Vec<_>>()
317 .into_iter()
318 .rev()
319 .collect()
320}
321
322fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
323 let mut files = messages
324 .iter()
325 .flat_map(|message| message.blocks.iter())
326 .filter_map(|block| match block {
327 ContentBlock::Text { text } => Some(text.as_str()),
328 ContentBlock::ToolUse { input, .. } => Some(input.as_str()),
329 ContentBlock::ToolResult { output, .. } => Some(output.as_str()),
330 ContentBlock::Image { .. } => None,
331 })
332 .flat_map(extract_file_candidates)
333 .collect::<Vec<_>>();
334 files.sort();
335 files.dedup();
336 files.into_iter().take(8).collect()
337}
338
339fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
340 messages
341 .iter()
342 .rev()
343 .filter_map(first_text_block)
344 .find(|text| !text.trim().is_empty())
345 .map(|text| truncate_summary(text, 200))
346}
347
348fn first_text_block(message: &ConversationMessage) -> Option<&str> {
349 message.blocks.iter().find_map(|block| match block {
350 ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
351 ContentBlock::ToolUse { .. }
352 | ContentBlock::ToolResult { .. }
353 | ContentBlock::Image { .. }
354 | ContentBlock::Text { .. } => None,
355 })
356}
357
358fn has_interesting_extension(candidate: &str) -> bool {
359 std::path::Path::new(candidate)
360 .extension()
361 .and_then(|extension| extension.to_str())
362 .is_some_and(|extension| {
363 ["rs", "ts", "tsx", "js", "json", "md"]
364 .iter()
365 .any(|expected| extension.eq_ignore_ascii_case(expected))
366 })
367}
368
369fn extract_file_candidates(content: &str) -> Vec<String> {
370 content
371 .split_whitespace()
372 .filter_map(|token| {
373 let candidate = token.trim_matches(|char: char| {
374 matches!(char, ',' | '.' | ':' | ';' | ')' | '(' | '"' | '\'' | '`')
375 });
376 if (candidate.contains('/') || candidate.contains('\\'))
377 && has_interesting_extension(candidate)
378 {
379 Some(candidate.to_string())
380 } else {
381 None
382 }
383 })
384 .collect()
385}
386
387fn truncate_summary(content: &str, max_chars: usize) -> String {
388 if content.chars().count() <= max_chars {
389 return content.to_string();
390 }
391 let mut truncated = content.chars().take(max_chars).collect::<String>();
392 truncated.push('…');
393 truncated
394}
395
396fn estimate_message_tokens(message: &ConversationMessage) -> usize {
397 message
398 .blocks
399 .iter()
400 .map(|block| match block {
401 ContentBlock::Text { text } => text.len() / 4 + 1,
402 ContentBlock::Image { data, .. } => data.len() / 4 + 1,
403 ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
404 ContentBlock::ToolResult {
405 tool_name, output, ..
406 } => (tool_name.len() + output.len()) / 4 + 1,
407 })
408 .sum()
409}
410
411fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
412 let start = format!("<{tag}>");
413 let end = format!("</{tag}>");
414 let start_index = content.find(&start)? + start.len();
415 let end_index = content[start_index..].find(&end)? + start_index;
416 Some(content[start_index..end_index].to_string())
417}
418
419fn strip_tag_block(content: &str, tag: &str) -> String {
420 let start = format!("<{tag}>");
421 let end = format!("</{tag}>");
422 if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
423 let end_index = end_index_rel + end.len();
424 let mut stripped = String::new();
425 stripped.push_str(&content[..start_index]);
426 stripped.push_str(&content[end_index..]);
427 stripped
428 } else {
429 content.to_string()
430 }
431}
432
433fn collapse_blank_lines(content: &str) -> String {
434 let mut result = String::new();
435 let mut last_blank = false;
436 for line in content.lines() {
437 let is_blank = line.trim().is_empty();
438 if is_blank && last_blank {
439 continue;
440 }
441 result.push_str(line);
442 result.push('\n');
443 last_blank = is_blank;
444 }
445 result
446}
447
448fn extract_existing_compacted_summary(message: &ConversationMessage) -> Option<String> {
449 if message.role != MessageRole::System {
450 return None;
451 }
452
453 let text = first_text_block(message)?;
454 let summary = text.strip_prefix(COMPACT_CONTINUATION_PREAMBLE)?;
455 let summary = summary
456 .split_once(&format!("\n\n{COMPACT_RECENT_MESSAGES_NOTE}"))
457 .map_or(summary, |(value, _)| value);
458 let summary = summary
459 .split_once(&format!("\n{COMPACT_DIRECT_RESUME_INSTRUCTION}"))
460 .map_or(summary, |(value, _)| value);
461 Some(summary.trim().to_string())
462}
463
464fn extract_summary_highlights(summary: &str) -> Vec<String> {
465 let mut lines = Vec::new();
466 let mut in_timeline = false;
467
468 for line in format_compact_summary(summary).lines() {
469 let trimmed = line.trim_end();
470 if trimmed.is_empty() || trimmed == "Summary:" || trimmed == "Conversation summary:" {
471 continue;
472 }
473 if trimmed == "- Key timeline:" {
474 in_timeline = true;
475 continue;
476 }
477 if in_timeline {
478 continue;
479 }
480 lines.push(trimmed.to_string());
481 }
482
483 lines
484}
485
486fn extract_summary_timeline(summary: &str) -> Vec<String> {
487 let mut lines = Vec::new();
488 let mut in_timeline = false;
489
490 for line in format_compact_summary(summary).lines() {
491 let trimmed = line.trim_end();
492 if trimmed == "- Key timeline:" {
493 in_timeline = true;
494 continue;
495 }
496 if !in_timeline {
497 continue;
498 }
499 if trimmed.is_empty() {
500 break;
501 }
502 lines.push(trimmed.to_string());
503 }
504
505 lines
506}
507
508#[cfg(test)]
509mod tests {
510 use super::{
511 collect_key_files, compact_session, estimate_session_tokens, format_compact_summary,
512 get_compact_continuation_message, infer_pending_work, should_compact, CompactionConfig,
513 };
514 use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
515
516 #[test]
517 fn formats_compact_summary_like_upstream() {
518 let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
519 assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
520 }
521
522 #[test]
523 fn leaves_small_sessions_unchanged() {
524 let session = Session {
525 version: 1,
526 messages: vec![ConversationMessage::user_text("hello")],
527 };
528
529 let result = compact_session(&session, CompactionConfig::default());
530 assert_eq!(result.removed_message_count, 0);
531 assert_eq!(result.compacted_session, session);
532 assert!(result.summary.is_empty());
533 assert!(result.formatted_summary.is_empty());
534 }
535
536 #[test]
537 fn compacts_older_messages_into_a_system_summary() {
538 let session = Session {
539 version: 1,
540 messages: vec![
541 ConversationMessage::user_text("one ".repeat(200)),
542 ConversationMessage::assistant(vec![ContentBlock::Text {
543 text: "two ".repeat(200),
544 }]),
545 ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
546 ConversationMessage {
547 role: MessageRole::Assistant,
548 blocks: vec![ContentBlock::Text {
549 text: "recent".to_string(),
550 }],
551 usage: None,
552 },
553 ],
554 };
555
556 let result = compact_session(
557 &session,
558 CompactionConfig {
559 preserve_recent_messages: 2,
560 max_estimated_tokens: 1,
561 },
562 );
563
564 assert_eq!(result.removed_message_count, 2);
565 assert_eq!(
566 result.compacted_session.messages[0].role,
567 MessageRole::System
568 );
569 assert!(matches!(
570 &result.compacted_session.messages[0].blocks[0],
571 ContentBlock::Text { text } if text.contains("Summary:")
572 ));
573 assert!(result.formatted_summary.contains("Scope:"));
574 assert!(result.formatted_summary.contains("Key timeline:"));
575 assert!(should_compact(
576 &session,
577 CompactionConfig {
578 preserve_recent_messages: 2,
579 max_estimated_tokens: 1,
580 }
581 ));
582 assert!(
583 estimate_session_tokens(&result.compacted_session) < estimate_session_tokens(&session)
584 );
585 }
586
587 #[test]
588 fn keeps_previous_compacted_context_when_compacting_again() {
589 let initial_session = Session {
590 version: 1,
591 messages: vec![
592 ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
593 ConversationMessage::assistant(vec![ContentBlock::Text {
594 text: "I will inspect the compact flow.".to_string(),
595 }]),
596 ConversationMessage::user_text(
597 "Also update rust/crates/runtime/src/conversation.rs",
598 ),
599 ConversationMessage::assistant(vec![ContentBlock::Text {
600 text: "Next: preserve prior summary context during auto compact.".to_string(),
601 }]),
602 ],
603 };
604 let config = CompactionConfig {
605 preserve_recent_messages: 2,
606 max_estimated_tokens: 1,
607 };
608
609 let first = compact_session(&initial_session, config);
610 let mut follow_up_messages = first.compacted_session.messages.clone();
611 follow_up_messages.extend([
612 ConversationMessage::user_text("Please add regression tests for compaction."),
613 ConversationMessage::assistant(vec![ContentBlock::Text {
614 text: "Working on regression coverage now.".to_string(),
615 }]),
616 ]);
617
618 let second = compact_session(
619 &Session {
620 version: 1,
621 messages: follow_up_messages,
622 },
623 config,
624 );
625
626 assert!(second
627 .formatted_summary
628 .contains("Previously compacted context:"));
629 assert!(second
630 .formatted_summary
631 .contains("Scope: 2 earlier messages compacted"));
632 assert!(second
633 .formatted_summary
634 .contains("Newly compacted context:"));
635 assert!(second
636 .formatted_summary
637 .contains("Also update rust/crates/runtime/src/conversation.rs"));
638 assert!(matches!(
639 &second.compacted_session.messages[0].blocks[0],
640 ContentBlock::Text { text }
641 if text.contains("Previously compacted context:")
642 && text.contains("Newly compacted context:")
643 ));
644 assert!(matches!(
645 &second.compacted_session.messages[1].blocks[0],
646 ContentBlock::Text { text } if text.contains("Please add regression tests for compaction.")
647 ));
648 }
649
650 #[test]
651 fn ignores_existing_compacted_summary_when_deciding_to_recompact() {
652 let summary = "<summary>Conversation summary:\n- Scope: earlier work preserved.\n- Key timeline:\n - user: large preserved context\n</summary>";
653 let session = Session {
654 version: 1,
655 messages: vec![
656 ConversationMessage {
657 role: MessageRole::System,
658 blocks: vec![ContentBlock::Text {
659 text: get_compact_continuation_message(summary, true, true),
660 }],
661 usage: None,
662 },
663 ConversationMessage::user_text("tiny"),
664 ConversationMessage::assistant(vec![ContentBlock::Text {
665 text: "recent".to_string(),
666 }]),
667 ],
668 };
669
670 assert!(!should_compact(
671 &session,
672 CompactionConfig {
673 preserve_recent_messages: 2,
674 max_estimated_tokens: 1,
675 }
676 ));
677 }
678
679 #[test]
680 fn truncates_long_blocks_in_summary() {
681 let summary = super::summarize_block(&ContentBlock::Text {
682 text: "x".repeat(400),
683 });
684 assert!(summary.ends_with('…'));
685 assert!(summary.chars().count() <= 161);
686 }
687
688 #[test]
689 fn extracts_key_files_from_message_content() {
690 let files = collect_key_files(&[ConversationMessage::user_text(
691 "Update rust/crates/runtime/src/compact.rs and rust/crates/tools/src/lib.rs next.",
692 )]);
693 assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
694 assert!(files.contains(&"rust/crates/tools/src/lib.rs".to_string()));
695 }
696
697 #[test]
698 fn infers_pending_work_from_recent_messages() {
699 let pending = infer_pending_work(&[
700 ConversationMessage::user_text("done"),
701 ConversationMessage::assistant(vec![ContentBlock::Text {
702 text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
703 }]),
704 ]);
705 assert_eq!(pending.len(), 1);
706 assert!(pending[0].contains("Next: update tests"));
707 }
708}