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