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('/') && has_interesting_extension(candidate) {
374 Some(candidate.to_string())
375 } else {
376 None
377 }
378 })
379 .collect()
380}
381
382fn truncate_summary(content: &str, max_chars: usize) -> String {
383 if content.chars().count() <= max_chars {
384 return content.to_string();
385 }
386 let mut truncated = content.chars().take(max_chars).collect::<String>();
387 truncated.push('…');
388 truncated
389}
390
391fn estimate_message_tokens(message: &ConversationMessage) -> usize {
392 message
393 .blocks
394 .iter()
395 .map(|block| match block {
396 ContentBlock::Text { text } => text.len() / 4 + 1,
397 ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
398 ContentBlock::ToolResult {
399 tool_name, output, ..
400 } => (tool_name.len() + output.len()) / 4 + 1,
401 })
402 .sum()
403}
404
405fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
406 let start = format!("<{tag}>");
407 let end = format!("</{tag}>");
408 let start_index = content.find(&start)? + start.len();
409 let end_index = content[start_index..].find(&end)? + start_index;
410 Some(content[start_index..end_index].to_string())
411}
412
413fn strip_tag_block(content: &str, tag: &str) -> String {
414 let start = format!("<{tag}>");
415 let end = format!("</{tag}>");
416 if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
417 let end_index = end_index_rel + end.len();
418 let mut stripped = String::new();
419 stripped.push_str(&content[..start_index]);
420 stripped.push_str(&content[end_index..]);
421 stripped
422 } else {
423 content.to_string()
424 }
425}
426
427fn collapse_blank_lines(content: &str) -> String {
428 let mut result = String::new();
429 let mut last_blank = false;
430 for line in content.lines() {
431 let is_blank = line.trim().is_empty();
432 if is_blank && last_blank {
433 continue;
434 }
435 result.push_str(line);
436 result.push('\n');
437 last_blank = is_blank;
438 }
439 result
440}
441
442fn extract_existing_compacted_summary(message: &ConversationMessage) -> Option<String> {
443 if message.role != MessageRole::System {
444 return None;
445 }
446
447 let text = first_text_block(message)?;
448 let summary = text.strip_prefix(COMPACT_CONTINUATION_PREAMBLE)?;
449 let summary = summary
450 .split_once(&format!("\n\n{COMPACT_RECENT_MESSAGES_NOTE}"))
451 .map_or(summary, |(value, _)| value);
452 let summary = summary
453 .split_once(&format!("\n{COMPACT_DIRECT_RESUME_INSTRUCTION}"))
454 .map_or(summary, |(value, _)| value);
455 Some(summary.trim().to_string())
456}
457
458fn extract_summary_highlights(summary: &str) -> Vec<String> {
459 let mut lines = Vec::new();
460 let mut in_timeline = false;
461
462 for line in format_compact_summary(summary).lines() {
463 let trimmed = line.trim_end();
464 if trimmed.is_empty() || trimmed == "Summary:" || trimmed == "Conversation summary:" {
465 continue;
466 }
467 if trimmed == "- Key timeline:" {
468 in_timeline = true;
469 continue;
470 }
471 if in_timeline {
472 continue;
473 }
474 lines.push(trimmed.to_string());
475 }
476
477 lines
478}
479
480fn extract_summary_timeline(summary: &str) -> Vec<String> {
481 let mut lines = Vec::new();
482 let mut in_timeline = false;
483
484 for line in format_compact_summary(summary).lines() {
485 let trimmed = line.trim_end();
486 if trimmed == "- Key timeline:" {
487 in_timeline = true;
488 continue;
489 }
490 if !in_timeline {
491 continue;
492 }
493 if trimmed.is_empty() {
494 break;
495 }
496 lines.push(trimmed.to_string());
497 }
498
499 lines
500}
501
502#[cfg(test)]
503mod tests {
504 use super::{
505 collect_key_files, compact_session, estimate_session_tokens, format_compact_summary,
506 get_compact_continuation_message, infer_pending_work, should_compact, CompactionConfig,
507 };
508 use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
509
510 #[test]
511 fn formats_compact_summary_like_upstream() {
512 let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
513 assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
514 }
515
516 #[test]
517 fn leaves_small_sessions_unchanged() {
518 let session = Session {
519 version: 1,
520 messages: vec![ConversationMessage::user_text("hello")],
521 };
522
523 let result = compact_session(&session, CompactionConfig::default());
524 assert_eq!(result.removed_message_count, 0);
525 assert_eq!(result.compacted_session, session);
526 assert!(result.summary.is_empty());
527 assert!(result.formatted_summary.is_empty());
528 }
529
530 #[test]
531 fn compacts_older_messages_into_a_system_summary() {
532 let session = Session {
533 version: 1,
534 messages: vec![
535 ConversationMessage::user_text("one ".repeat(200)),
536 ConversationMessage::assistant(vec![ContentBlock::Text {
537 text: "two ".repeat(200),
538 }]),
539 ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
540 ConversationMessage {
541 role: MessageRole::Assistant,
542 blocks: vec![ContentBlock::Text {
543 text: "recent".to_string(),
544 }],
545 usage: None,
546 },
547 ],
548 };
549
550 let result = compact_session(
551 &session,
552 CompactionConfig {
553 preserve_recent_messages: 2,
554 max_estimated_tokens: 1,
555 },
556 );
557
558 assert_eq!(result.removed_message_count, 2);
559 assert_eq!(
560 result.compacted_session.messages[0].role,
561 MessageRole::System
562 );
563 assert!(matches!(
564 &result.compacted_session.messages[0].blocks[0],
565 ContentBlock::Text { text } if text.contains("Summary:")
566 ));
567 assert!(result.formatted_summary.contains("Scope:"));
568 assert!(result.formatted_summary.contains("Key timeline:"));
569 assert!(should_compact(
570 &session,
571 CompactionConfig {
572 preserve_recent_messages: 2,
573 max_estimated_tokens: 1,
574 }
575 ));
576 assert!(
577 estimate_session_tokens(&result.compacted_session) < estimate_session_tokens(&session)
578 );
579 }
580
581 #[test]
582 fn keeps_previous_compacted_context_when_compacting_again() {
583 let initial_session = Session {
584 version: 1,
585 messages: vec![
586 ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
587 ConversationMessage::assistant(vec![ContentBlock::Text {
588 text: "I will inspect the compact flow.".to_string(),
589 }]),
590 ConversationMessage::user_text(
591 "Also update rust/crates/runtime/src/conversation.rs",
592 ),
593 ConversationMessage::assistant(vec![ContentBlock::Text {
594 text: "Next: preserve prior summary context during auto compact.".to_string(),
595 }]),
596 ],
597 };
598 let config = CompactionConfig {
599 preserve_recent_messages: 2,
600 max_estimated_tokens: 1,
601 };
602
603 let first = compact_session(&initial_session, config);
604 let mut follow_up_messages = first.compacted_session.messages.clone();
605 follow_up_messages.extend([
606 ConversationMessage::user_text("Please add regression tests for compaction."),
607 ConversationMessage::assistant(vec![ContentBlock::Text {
608 text: "Working on regression coverage now.".to_string(),
609 }]),
610 ]);
611
612 let second = compact_session(
613 &Session {
614 version: 1,
615 messages: follow_up_messages,
616 },
617 config,
618 );
619
620 assert!(second
621 .formatted_summary
622 .contains("Previously compacted context:"));
623 assert!(second
624 .formatted_summary
625 .contains("Scope: 2 earlier messages compacted"));
626 assert!(second
627 .formatted_summary
628 .contains("Newly compacted context:"));
629 assert!(second
630 .formatted_summary
631 .contains("Also update rust/crates/runtime/src/conversation.rs"));
632 assert!(matches!(
633 &second.compacted_session.messages[0].blocks[0],
634 ContentBlock::Text { text }
635 if text.contains("Previously compacted context:")
636 && text.contains("Newly compacted context:")
637 ));
638 assert!(matches!(
639 &second.compacted_session.messages[1].blocks[0],
640 ContentBlock::Text { text } if text.contains("Please add regression tests for compaction.")
641 ));
642 }
643
644 #[test]
645 fn ignores_existing_compacted_summary_when_deciding_to_recompact() {
646 let summary = "<summary>Conversation summary:\n- Scope: earlier work preserved.\n- Key timeline:\n - user: large preserved context\n</summary>";
647 let session = Session {
648 version: 1,
649 messages: vec![
650 ConversationMessage {
651 role: MessageRole::System,
652 blocks: vec![ContentBlock::Text {
653 text: get_compact_continuation_message(summary, true, true),
654 }],
655 usage: None,
656 },
657 ConversationMessage::user_text("tiny"),
658 ConversationMessage::assistant(vec![ContentBlock::Text {
659 text: "recent".to_string(),
660 }]),
661 ],
662 };
663
664 assert!(!should_compact(
665 &session,
666 CompactionConfig {
667 preserve_recent_messages: 2,
668 max_estimated_tokens: 1,
669 }
670 ));
671 }
672
673 #[test]
674 fn truncates_long_blocks_in_summary() {
675 let summary = super::summarize_block(&ContentBlock::Text {
676 text: "x".repeat(400),
677 });
678 assert!(summary.ends_with('…'));
679 assert!(summary.chars().count() <= 161);
680 }
681
682 #[test]
683 fn extracts_key_files_from_message_content() {
684 let files = collect_key_files(&[ConversationMessage::user_text(
685 "Update rust/crates/runtime/src/compact.rs and rust/crates/tools/src/lib.rs next.",
686 )]);
687 assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
688 assert!(files.contains(&"rust/crates/tools/src/lib.rs".to_string()));
689 }
690
691 #[test]
692 fn infers_pending_work_from_recent_messages() {
693 let pending = infer_pending_work(&[
694 ConversationMessage::user_text("done"),
695 ConversationMessage::assistant(vec![ContentBlock::Text {
696 text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
697 }]),
698 ]);
699 assert_eq!(pending.len(), 1);
700 assert!(pending[0].contains("Next: update tests"));
701 }
702}