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 { .. } => 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::ToolUse { name, input, .. } => format!("tool_use {name}({input})"),
269 ContentBlock::ToolResult {
270 tool_name,
271 output,
272 is_error,
273 ..
274 } => format!(
275 "tool_result {tool_name}: {}{output}",
276 if *is_error { "error " } else { "" }
277 ),
278 };
279 truncate_summary(&raw, 160)
280}
281
282fn collect_recent_role_summaries(
283 messages: &[ConversationMessage],
284 role: MessageRole,
285 limit: usize,
286) -> Vec<String> {
287 messages
288 .iter()
289 .filter(|message| message.role == role)
290 .rev()
291 .filter_map(|message| first_text_block(message))
292 .take(limit)
293 .map(|text| truncate_summary(text, 160))
294 .collect::<Vec<_>>()
295 .into_iter()
296 .rev()
297 .collect()
298}
299
300fn infer_pending_work(messages: &[ConversationMessage]) -> Vec<String> {
301 messages
302 .iter()
303 .rev()
304 .filter_map(first_text_block)
305 .filter(|text| {
306 let lowered = text.to_ascii_lowercase();
307 lowered.contains("todo")
308 || lowered.contains("next")
309 || lowered.contains("pending")
310 || lowered.contains("follow up")
311 || lowered.contains("remaining")
312 })
313 .take(3)
314 .map(|text| truncate_summary(text, 160))
315 .collect::<Vec<_>>()
316 .into_iter()
317 .rev()
318 .collect()
319}
320
321fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
322 let mut files = messages
323 .iter()
324 .flat_map(|message| message.blocks.iter())
325 .map(|block| match block {
326 ContentBlock::Text { text } => text.as_str(),
327 ContentBlock::ToolUse { input, .. } => input.as_str(),
328 ContentBlock::ToolResult { output, .. } => output.as_str(),
329 })
330 .flat_map(extract_file_candidates)
331 .collect::<Vec<_>>();
332 files.sort();
333 files.dedup();
334 files.into_iter().take(8).collect()
335}
336
337fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
338 messages
339 .iter()
340 .rev()
341 .filter_map(first_text_block)
342 .find(|text| !text.trim().is_empty())
343 .map(|text| truncate_summary(text, 200))
344}
345
346fn first_text_block(message: &ConversationMessage) -> Option<&str> {
347 message.blocks.iter().find_map(|block| match block {
348 ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
349 ContentBlock::ToolUse { .. }
350 | ContentBlock::ToolResult { .. }
351 | ContentBlock::Text { .. } => None,
352 })
353}
354
355fn has_interesting_extension(candidate: &str) -> bool {
356 std::path::Path::new(candidate)
357 .extension()
358 .and_then(|extension| extension.to_str())
359 .is_some_and(|extension| {
360 ["rs", "ts", "tsx", "js", "json", "md"]
361 .iter()
362 .any(|expected| extension.eq_ignore_ascii_case(expected))
363 })
364}
365
366fn extract_file_candidates(content: &str) -> Vec<String> {
367 content
368 .split_whitespace()
369 .filter_map(|token| {
370 let candidate = token.trim_matches(|char: char| {
371 matches!(char, ',' | '.' | ':' | ';' | ')' | '(' | '"' | '\'' | '`')
372 });
373 if (candidate.contains('/') || candidate.contains('\\'))
374 && has_interesting_extension(candidate)
375 {
376 Some(candidate.to_string())
377 } else {
378 None
379 }
380 })
381 .collect()
382}
383
384fn truncate_summary(content: &str, max_chars: usize) -> String {
385 if content.chars().count() <= max_chars {
386 return content.to_string();
387 }
388 let mut truncated = content.chars().take(max_chars).collect::<String>();
389 truncated.push('…');
390 truncated
391}
392
393fn estimate_message_tokens(message: &ConversationMessage) -> usize {
394 message
395 .blocks
396 .iter()
397 .map(|block| match block {
398 ContentBlock::Text { text } => text.len() / 4 + 1,
399 ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
400 ContentBlock::ToolResult {
401 tool_name, output, ..
402 } => (tool_name.len() + output.len()) / 4 + 1,
403 })
404 .sum()
405}
406
407fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
408 let start = format!("<{tag}>");
409 let end = format!("</{tag}>");
410 let start_index = content.find(&start)? + start.len();
411 let end_index = content[start_index..].find(&end)? + start_index;
412 Some(content[start_index..end_index].to_string())
413}
414
415fn strip_tag_block(content: &str, tag: &str) -> String {
416 let start = format!("<{tag}>");
417 let end = format!("</{tag}>");
418 if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
419 let end_index = end_index_rel + end.len();
420 let mut stripped = String::new();
421 stripped.push_str(&content[..start_index]);
422 stripped.push_str(&content[end_index..]);
423 stripped
424 } else {
425 content.to_string()
426 }
427}
428
429fn collapse_blank_lines(content: &str) -> String {
430 let mut result = String::new();
431 let mut last_blank = false;
432 for line in content.lines() {
433 let is_blank = line.trim().is_empty();
434 if is_blank && last_blank {
435 continue;
436 }
437 result.push_str(line);
438 result.push('\n');
439 last_blank = is_blank;
440 }
441 result
442}
443
444fn extract_existing_compacted_summary(message: &ConversationMessage) -> Option<String> {
445 if message.role != MessageRole::System {
446 return None;
447 }
448
449 let text = first_text_block(message)?;
450 let summary = text.strip_prefix(COMPACT_CONTINUATION_PREAMBLE)?;
451 let summary = summary
452 .split_once(&format!("\n\n{COMPACT_RECENT_MESSAGES_NOTE}"))
453 .map_or(summary, |(value, _)| value);
454 let summary = summary
455 .split_once(&format!("\n{COMPACT_DIRECT_RESUME_INSTRUCTION}"))
456 .map_or(summary, |(value, _)| value);
457 Some(summary.trim().to_string())
458}
459
460fn extract_summary_highlights(summary: &str) -> Vec<String> {
461 let mut lines = Vec::new();
462 let mut in_timeline = false;
463
464 for line in format_compact_summary(summary).lines() {
465 let trimmed = line.trim_end();
466 if trimmed.is_empty() || trimmed == "Summary:" || trimmed == "Conversation summary:" {
467 continue;
468 }
469 if trimmed == "- Key timeline:" {
470 in_timeline = true;
471 continue;
472 }
473 if in_timeline {
474 continue;
475 }
476 lines.push(trimmed.to_string());
477 }
478
479 lines
480}
481
482fn extract_summary_timeline(summary: &str) -> Vec<String> {
483 let mut lines = Vec::new();
484 let mut in_timeline = false;
485
486 for line in format_compact_summary(summary).lines() {
487 let trimmed = line.trim_end();
488 if trimmed == "- Key timeline:" {
489 in_timeline = true;
490 continue;
491 }
492 if !in_timeline {
493 continue;
494 }
495 if trimmed.is_empty() {
496 break;
497 }
498 lines.push(trimmed.to_string());
499 }
500
501 lines
502}
503
504#[cfg(test)]
505mod tests {
506 use super::{
507 collect_key_files, compact_session, estimate_session_tokens, format_compact_summary,
508 get_compact_continuation_message, infer_pending_work, should_compact, CompactionConfig,
509 };
510 use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
511
512 #[test]
513 fn formats_compact_summary_like_upstream() {
514 let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
515 assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
516 }
517
518 #[test]
519 fn leaves_small_sessions_unchanged() {
520 let session = Session {
521 version: 1,
522 messages: vec![ConversationMessage::user_text("hello")],
523 };
524
525 let result = compact_session(&session, CompactionConfig::default());
526 assert_eq!(result.removed_message_count, 0);
527 assert_eq!(result.compacted_session, session);
528 assert!(result.summary.is_empty());
529 assert!(result.formatted_summary.is_empty());
530 }
531
532 #[test]
533 fn compacts_older_messages_into_a_system_summary() {
534 let session = Session {
535 version: 1,
536 messages: vec![
537 ConversationMessage::user_text("one ".repeat(200)),
538 ConversationMessage::assistant(vec![ContentBlock::Text {
539 text: "two ".repeat(200),
540 }]),
541 ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
542 ConversationMessage {
543 role: MessageRole::Assistant,
544 blocks: vec![ContentBlock::Text {
545 text: "recent".to_string(),
546 }],
547 usage: None,
548 },
549 ],
550 };
551
552 let result = compact_session(
553 &session,
554 CompactionConfig {
555 preserve_recent_messages: 2,
556 max_estimated_tokens: 1,
557 },
558 );
559
560 assert_eq!(result.removed_message_count, 2);
561 assert_eq!(
562 result.compacted_session.messages[0].role,
563 MessageRole::System
564 );
565 assert!(matches!(
566 &result.compacted_session.messages[0].blocks[0],
567 ContentBlock::Text { text } if text.contains("Summary:")
568 ));
569 assert!(result.formatted_summary.contains("Scope:"));
570 assert!(result.formatted_summary.contains("Key timeline:"));
571 assert!(should_compact(
572 &session,
573 CompactionConfig {
574 preserve_recent_messages: 2,
575 max_estimated_tokens: 1,
576 }
577 ));
578 assert!(
579 estimate_session_tokens(&result.compacted_session) < estimate_session_tokens(&session)
580 );
581 }
582
583 #[test]
584 fn keeps_previous_compacted_context_when_compacting_again() {
585 let initial_session = Session {
586 version: 1,
587 messages: vec![
588 ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
589 ConversationMessage::assistant(vec![ContentBlock::Text {
590 text: "I will inspect the compact flow.".to_string(),
591 }]),
592 ConversationMessage::user_text(
593 "Also update rust/crates/runtime/src/conversation.rs",
594 ),
595 ConversationMessage::assistant(vec![ContentBlock::Text {
596 text: "Next: preserve prior summary context during auto compact.".to_string(),
597 }]),
598 ],
599 };
600 let config = CompactionConfig {
601 preserve_recent_messages: 2,
602 max_estimated_tokens: 1,
603 };
604
605 let first = compact_session(&initial_session, config);
606 let mut follow_up_messages = first.compacted_session.messages.clone();
607 follow_up_messages.extend([
608 ConversationMessage::user_text("Please add regression tests for compaction."),
609 ConversationMessage::assistant(vec![ContentBlock::Text {
610 text: "Working on regression coverage now.".to_string(),
611 }]),
612 ]);
613
614 let second = compact_session(
615 &Session {
616 version: 1,
617 messages: follow_up_messages,
618 },
619 config,
620 );
621
622 assert!(second
623 .formatted_summary
624 .contains("Previously compacted context:"));
625 assert!(second
626 .formatted_summary
627 .contains("Scope: 2 earlier messages compacted"));
628 assert!(second
629 .formatted_summary
630 .contains("Newly compacted context:"));
631 assert!(second
632 .formatted_summary
633 .contains("Also update rust/crates/runtime/src/conversation.rs"));
634 assert!(matches!(
635 &second.compacted_session.messages[0].blocks[0],
636 ContentBlock::Text { text }
637 if text.contains("Previously compacted context:")
638 && text.contains("Newly compacted context:")
639 ));
640 assert!(matches!(
641 &second.compacted_session.messages[1].blocks[0],
642 ContentBlock::Text { text } if text.contains("Please add regression tests for compaction.")
643 ));
644 }
645
646 #[test]
647 fn ignores_existing_compacted_summary_when_deciding_to_recompact() {
648 let summary = "<summary>Conversation summary:\n- Scope: earlier work preserved.\n- Key timeline:\n - user: large preserved context\n</summary>";
649 let session = Session {
650 version: 1,
651 messages: vec![
652 ConversationMessage {
653 role: MessageRole::System,
654 blocks: vec![ContentBlock::Text {
655 text: get_compact_continuation_message(summary, true, true),
656 }],
657 usage: None,
658 },
659 ConversationMessage::user_text("tiny"),
660 ConversationMessage::assistant(vec![ContentBlock::Text {
661 text: "recent".to_string(),
662 }]),
663 ],
664 };
665
666 assert!(!should_compact(
667 &session,
668 CompactionConfig {
669 preserve_recent_messages: 2,
670 max_estimated_tokens: 1,
671 }
672 ));
673 }
674
675 #[test]
676 fn truncates_long_blocks_in_summary() {
677 let summary = super::summarize_block(&ContentBlock::Text {
678 text: "x".repeat(400),
679 });
680 assert!(summary.ends_with('…'));
681 assert!(summary.chars().count() <= 161);
682 }
683
684 #[test]
685 fn extracts_key_files_from_message_content() {
686 let files = collect_key_files(&[ConversationMessage::user_text(
687 "Update rust/crates/runtime/src/compact.rs and rust/crates/tools/src/lib.rs next.",
688 )]);
689 assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
690 assert!(files.contains(&"rust/crates/tools/src/lib.rs".to_string()));
691 }
692
693 #[test]
694 fn infers_pending_work_from_recent_messages() {
695 let pending = infer_pending_work(&[
696 ConversationMessage::user_text("done"),
697 ConversationMessage::assistant(vec![ContentBlock::Text {
698 text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
699 }]),
700 ]);
701 assert_eq!(pending.len(), 1);
702 assert!(pending[0].contains("Next: update tests"));
703 }
704}