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 { .. } | ContentBlock::Thinking { .. } => 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 ContentBlock::Thinking { text, .. } => format!("[thinking: {}]", &text[..text.len().min(80)]),
215 };
216 truncate_summary(&raw, 160)
217}
218
219fn collect_recent_role_summaries(
220 messages: &[ConversationMessage],
221 role: MessageRole,
222 limit: usize,
223) -> Vec<String> {
224 messages
225 .iter()
226 .filter(|message| message.role == role)
227 .rev()
228 .filter_map(|message| first_text_block(message))
229 .take(limit)
230 .map(|text| truncate_summary(text, 160))
231 .collect::<Vec<_>>()
232 .into_iter()
233 .rev()
234 .collect()
235}
236
237fn infer_pending_work(messages: &[ConversationMessage]) -> Vec<String> {
238 messages
239 .iter()
240 .rev()
241 .filter_map(first_text_block)
242 .filter(|text| {
243 let lowered = text.to_ascii_lowercase();
244 lowered.contains("todo")
245 || lowered.contains("next")
246 || lowered.contains("pending")
247 || lowered.contains("follow up")
248 || lowered.contains("remaining")
249 })
250 .take(3)
251 .map(|text| truncate_summary(text, 160))
252 .collect::<Vec<_>>()
253 .into_iter()
254 .rev()
255 .collect()
256}
257
258fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
259 let mut files = messages
260 .iter()
261 .flat_map(|message| message.blocks.iter())
262 .filter_map(|block| match block {
263 ContentBlock::Text { text } => Some(text.as_str()),
264 ContentBlock::ToolUse { input, .. } => Some(input.as_str()),
265 ContentBlock::ToolResult { output, .. } => Some(output.as_str()),
266 ContentBlock::Image { .. } | ContentBlock::Thinking { .. } => None,
267 })
268 .flat_map(extract_file_candidates)
269 .collect::<Vec<_>>();
270 files.sort();
271 files.dedup();
272 files.into_iter().take(8).collect()
273}
274
275fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
276 messages
277 .iter()
278 .rev()
279 .filter_map(first_text_block)
280 .find(|text| !text.trim().is_empty())
281 .map(|text| truncate_summary(text, 200))
282}
283
284fn first_text_block(message: &ConversationMessage) -> Option<&str> {
285 message.blocks.iter().find_map(|block| match block {
286 ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
287 ContentBlock::ToolUse { .. }
288 | ContentBlock::ToolResult { .. }
289 | ContentBlock::Text { .. }
290 | ContentBlock::Image { .. }
291 | ContentBlock::Thinking { .. } => None,
292 })
293}
294
295fn has_interesting_extension(candidate: &str) -> bool {
296 std::path::Path::new(candidate)
297 .extension()
298 .and_then(|extension| extension.to_str())
299 .is_some_and(|extension| {
300 ["rs", "ts", "tsx", "js", "json", "md"]
301 .iter()
302 .any(|expected| extension.eq_ignore_ascii_case(expected))
303 })
304}
305
306fn extract_file_candidates(content: &str) -> Vec<String> {
307 content
308 .split_whitespace()
309 .filter_map(|token| {
310 let candidate = token.trim_matches(|char: char| {
311 matches!(char, ',' | '.' | ':' | ';' | ')' | '(' | '"' | '\'' | '`')
312 });
313 if candidate.contains('/') && has_interesting_extension(candidate) {
314 Some(candidate.to_string())
315 } else {
316 None
317 }
318 })
319 .collect()
320}
321
322fn truncate_summary(content: &str, max_chars: usize) -> String {
323 if content.chars().count() <= max_chars {
324 return content.to_string();
325 }
326 let mut truncated = content.chars().take(max_chars).collect::<String>();
327 truncated.push('…');
328 truncated
329}
330
331fn estimate_message_tokens(message: &ConversationMessage) -> usize {
332 message
333 .blocks
334 .iter()
335 .map(|block| match block {
336 ContentBlock::Text { text } => text.len() / 4 + 1,
337 ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
338 ContentBlock::ToolResult { tool_name, output, .. } => (tool_name.len() + output.len()) / 4 + 1,
339 ContentBlock::Image { data, .. } => data.len() / 4 + 1,
340 ContentBlock::Thinking { text, thought_signature } => {
341 text.len() / 4 + thought_signature.as_ref().map_or(0, |s| s.len() / 4) + 1
342 }
343 })
344 .sum()
345}
346
347fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
348 let start = format!("<{tag}>");
349 let end = format!("</{tag}>");
350 let start_index = content.find(&start)? + start.len();
351 let end_index = content[start_index..].find(&end)? + start_index;
352 Some(content[start_index..end_index].to_string())
353}
354
355fn strip_tag_block(content: &str, tag: &str) -> String {
356 let start = format!("<{tag}>");
357 let end = format!("</{tag}>");
358 if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
359 let end_index = end_index_rel + end.len();
360 let mut stripped = String::new();
361 stripped.push_str(&content[..start_index]);
362 stripped.push_str(&content[end_index..]);
363 stripped
364 } else {
365 content.to_string()
366 }
367}
368
369fn collapse_blank_lines(content: &str) -> String {
370 let mut result = String::new();
371 let mut last_blank = false;
372 for line in content.lines() {
373 let is_blank = line.trim().is_empty();
374 if is_blank && last_blank {
375 continue;
376 }
377 result.push_str(line);
378 result.push('\n');
379 last_blank = is_blank;
380 }
381 result
382}
383
384#[cfg(test)]
385mod tests {
386 use super::{
387 collect_key_files, compact_session, estimate_session_tokens, format_compact_summary,
388 infer_pending_work, should_compact, CompactionConfig,
389 };
390 use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
391
392 #[test]
393 fn formats_compact_summary_like_upstream() {
394 let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
395 assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
396 }
397
398 #[test]
399 fn leaves_small_sessions_unchanged() {
400 let session = Session {
401 version: 1,
402 messages: vec![ConversationMessage::user_text("hello")],
403 };
404
405 let result = compact_session(&session, CompactionConfig::default());
406 assert_eq!(result.removed_message_count, 0);
407 assert_eq!(result.compacted_session, session);
408 assert!(result.summary.is_empty());
409 assert!(result.formatted_summary.is_empty());
410 }
411
412 #[test]
413 fn compacts_older_messages_into_a_system_summary() {
414 let session = Session {
415 version: 1,
416 messages: vec![
417 ConversationMessage::user_text("one ".repeat(200)),
418 ConversationMessage::assistant(vec![ContentBlock::Text {
419 text: "two ".repeat(200),
420 }]),
421 ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
422 ConversationMessage {
423 role: MessageRole::Assistant,
424 blocks: vec![ContentBlock::Text {
425 text: "recent".to_string(),
426 }],
427 usage: None,
428 },
429 ],
430 };
431
432 let result = compact_session(
433 &session,
434 CompactionConfig {
435 preserve_recent_messages: 2,
436 max_estimated_tokens: 1,
437 },
438 );
439
440 assert_eq!(result.removed_message_count, 2);
441 assert_eq!(
442 result.compacted_session.messages[0].role,
443 MessageRole::System
444 );
445 assert!(matches!(
446 &result.compacted_session.messages[0].blocks[0],
447 ContentBlock::Text { text } if text.contains("Summary:")
448 ));
449 assert!(result.formatted_summary.contains("Scope:"));
450 assert!(result.formatted_summary.contains("Key timeline:"));
451 assert!(should_compact(
452 &session,
453 CompactionConfig {
454 preserve_recent_messages: 2,
455 max_estimated_tokens: 1,
456 }
457 ));
458 assert!(
459 estimate_session_tokens(&result.compacted_session) < estimate_session_tokens(&session)
460 );
461 }
462
463 #[test]
464 fn truncates_long_blocks_in_summary() {
465 let summary = super::summarize_block(&ContentBlock::Text {
466 text: "x".repeat(400),
467 });
468 assert!(summary.ends_with('…'));
469 assert!(summary.chars().count() <= 161);
470 }
471
472 #[test]
473 fn extracts_key_files_from_message_content() {
474 let files = collect_key_files(&[ConversationMessage::user_text(
475 "Update rust/crates/runtime/src/compact.rs and rust/crates/rusty-ternlang-cli/src/main.rs next.",
476 )]);
477 assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
478 assert!(files.contains(&"rust/crates/rusty-ternlang-cli/src/main.rs".to_string()));
479 }
480
481 #[test]
482 fn infers_pending_work_from_recent_messages() {
483 let pending = infer_pending_work(&[
484 ConversationMessage::user_text("done"),
485 ConversationMessage::assistant(vec![ContentBlock::Text {
486 text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
487 }]),
488 ]);
489 assert_eq!(pending.len(), 1);
490 assert!(pending[0].contains("Next: update tests"));
491 }
492}