1#![allow(dead_code)]
2#[cfg(not(feature = "archive-compress-zstd"))]
5use crate::compression::NoCompressor;
6use crate::compression::SegmentCompressor;
7#[cfg(feature = "archive-compress-zstd")]
8use crate::compression::ZstdCompressor;
9use crate::segment_chain::{ChainMetadata, SegmentMeta, SegmentStatus, chain_path};
10use crate::store::SessionStore;
11use crate::{SessionEntry, SessionError};
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14
15pub struct CompactionEngine {
16 store: Arc<dyn SessionStore>,
17 compressor: Box<dyn SegmentCompressor>,
18 rules: CompactionRules,
19}
20
21#[derive(Clone)]
22pub struct CompactionRules {
23 pub tool_result_threshold_turns: usize,
24 pub max_tool_result_chars: usize,
25 pub remove_old_thinking: bool,
26}
27
28impl Default for CompactionRules {
29 fn default() -> Self {
30 Self {
31 tool_result_threshold_turns: 20,
32 max_tool_result_chars: 4000,
33 remove_old_thinking: true,
34 }
35 }
36}
37
38impl CompactionEngine {
39 pub fn new(store: Arc<dyn SessionStore>) -> Self {
40 Self {
41 store,
42 #[cfg(feature = "archive-compress-zstd")]
43 compressor: Box::new(ZstdCompressor::default()),
44 #[cfg(not(feature = "archive-compress-zstd"))]
45 compressor: Box::new(NoCompressor),
46 rules: CompactionRules::default(),
47 }
48 }
49
50 pub fn with_compressor(mut self, compressor: Box<dyn SegmentCompressor>) -> Self {
51 self.compressor = compressor;
52 self
53 }
54
55 pub fn with_rules(mut self, rules: CompactionRules) -> Self {
56 self.rules = rules;
57 self
58 }
59
60 pub fn should_compact(&self, head_path: &Path, max_entries: usize) -> bool {
61 if let Ok(entries) = self.store.read_entries(head_path) {
62 entries.len() > max_entries
63 } else {
64 false
65 }
66 }
67
68 pub fn compact_segment(
69 &self,
70 head_path: &Path,
71 session_dir: &Path,
72 max_entries: usize,
73 ) -> Result<CompactionResult, SessionError> {
74 let entries = self.store.read_entries(head_path)?;
75 if entries.len() <= max_entries {
76 return Ok(CompactionResult::Skipped);
77 }
78
79 let segment_id = format!("s-{}", uuid::Uuid::new_v4());
80 let archive_path = session_dir.join(format!("{segment_id}.tlog"));
81 let compressed_path = session_dir.join(format!("{segment_id}.tlog.zst"));
82
83 let original_bytes = std::fs::metadata(head_path).map(|m| m.len()).unwrap_or(0);
84
85 let archived_path = self.archive_head(head_path, &archive_path, &compressed_path)?;
86
87 let compacted = self.apply_rules(&entries, max_entries);
88
89 std::fs::write(head_path, b"")?;
90 for entry in &compacted {
91 self.store.append_entry(head_path, entry)?;
92 }
93
94 let _record_count = entries.len();
95 let archived_bytes = std::fs::metadata(&archived_path)
96 .map(|m| m.len())
97 .unwrap_or(0);
98
99 self.update_chain(
100 session_dir,
101 &segment_id,
102 _record_count,
103 original_bytes,
104 archived_bytes,
105 )?;
106
107 Ok(CompactionResult::Compacted {
108 segment_id,
109 original_count: _record_count,
110 compacted_count: compacted.len(),
111 original_bytes,
112 archived_bytes,
113 })
114 }
115
116 fn archive_head(
117 &self,
118 head_path: &Path,
119 archive_path: &Path,
120 compressed_path: &Path,
121 ) -> Result<PathBuf, SessionError> {
122 let raw = std::fs::read(head_path)?;
123 let compressed = self
124 .compressor
125 .compress(&raw)
126 .map_err(|e| SessionError::ParseError(e.to_string()))?;
127 if self.compressor.format_tag() == "none" {
128 std::fs::write(archive_path, &raw)?;
129 Ok(archive_path.to_path_buf())
130 } else {
131 std::fs::write(compressed_path, &compressed)?;
132 Ok(compressed_path.to_path_buf())
133 }
134 }
135
136 fn apply_rules(&self, entries: &[SessionEntry], keep_recent: usize) -> Vec<SessionEntry> {
137 if entries.len() <= keep_recent {
138 return entries.to_vec();
139 }
140
141 let split = entries.len() - keep_recent;
142 let mut result = Vec::new();
143
144 result.push(SessionEntry {
145 id: uuid::Uuid::new_v4().to_string(),
146 parent_id: None,
147 timestamp: chrono::Utc::now(),
148 role: "system".into(),
149 content: format!(
150 "[Compaction summary: {} earlier entries summarized, {} recent entries preserved]",
151 split, keep_recent
152 ),
153 metadata: crate::SessionMetadata::default(),
154 });
155
156 for entry in &entries[split..] {
157 result.push(entry.clone());
158 }
159
160 result
161 }
162
163 fn update_chain(
164 &self,
165 session_dir: &Path,
166 segment_id: &str,
167 _record_count: usize,
168 _orig_bytes: u64,
169 archived_bytes: u64,
170 ) -> Result<(), SessionError> {
171 let chain_file = chain_path(session_dir);
172 let mut chain = ChainMetadata::read(&chain_file)
173 .map_err(|e| SessionError::ParseError(e.to_string()))?;
174
175 let previous_archive = chain
176 .head_segment()
177 .and_then(|head| head.prev_segment_id.clone());
178 let now = chrono::Utc::now().timestamp_millis();
179 chain
180 .segments
181 .retain(|segment| segment.status != SegmentStatus::Active);
182 chain.segments.push(SegmentMeta {
183 segment_id: segment_id.to_string(),
184 status: SegmentStatus::Compressed,
185 prev_segment_id: previous_archive,
186 record_count: _record_count,
187 orig_bytes: _orig_bytes,
188 archived_bytes: Some(archived_bytes),
189 created_ts: now,
190 archived_ts: Some(now),
191 archive_format: Some(self.compressor.format_tag().to_string()),
192 ref_count: 0,
193 });
194 chain.segments.push(SegmentMeta {
195 segment_id: "head".to_string(),
196 status: SegmentStatus::Active,
197 prev_segment_id: Some(segment_id.to_string()),
198 record_count: 0,
199 orig_bytes: 0,
200 archived_bytes: None,
201 created_ts: now,
202 archived_ts: None,
203 archive_format: None,
204 ref_count: 0,
205 });
206
207 chain
208 .write(&chain_file)
209 .map_err(|e| SessionError::ParseError(e.to_string()))?;
210 Ok(())
211 }
212}
213
214#[derive(Debug)]
215pub enum CompactionResult {
216 Skipped,
217 Compacted {
218 segment_id: String,
219 original_count: usize,
220 compacted_count: usize,
221 original_bytes: u64,
222 archived_bytes: u64,
223 },
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229 use crate::{CompactTextSessionStore, SessionStore};
230
231 #[test]
232 fn should_compact_returns_false_under_threshold() {
233 let store = Arc::new(CompactTextSessionStore);
234 let dir = std::env::temp_dir().join("compact_test_threshold");
235 let _ = std::fs::remove_dir_all(&dir);
236 std::fs::create_dir_all(&dir).expect("operation should succeed");
237 let path = dir.join("head.tlog");
238
239 let entry = SessionEntry {
240 id: uuid::Uuid::new_v4().to_string(),
241 parent_id: None,
242 timestamp: chrono::Utc::now(),
243 role: "user".into(),
244 content: "hello".into(),
245 metadata: crate::SessionMetadata::default(),
246 };
247 store
248 .append_entry(&path, &entry)
249 .expect("operation should succeed");
250
251 let engine = CompactionEngine::new(store);
252 assert!(!engine.should_compact(&path, 10));
253
254 std::fs::remove_dir_all(&dir).ok();
255 }
256
257 #[test]
258 fn should_compact_returns_true_over_threshold() {
259 let store = Arc::new(CompactTextSessionStore);
260 let dir = std::env::temp_dir().join("compact_test_over");
261 let _ = std::fs::remove_dir_all(&dir);
262 std::fs::create_dir_all(&dir).expect("operation should succeed");
263 let path = dir.join("head.tlog");
264
265 for _ in 0..5 {
266 let entry = SessionEntry {
267 id: uuid::Uuid::new_v4().to_string(),
268 parent_id: None,
269 timestamp: chrono::Utc::now(),
270 role: "user".into(),
271 content: "test entry".into(),
272 metadata: crate::SessionMetadata::default(),
273 };
274 store
275 .append_entry(&path, &entry)
276 .expect("operation should succeed");
277 }
278
279 let engine = CompactionEngine::new(store);
280 assert!(engine.should_compact(&path, 3));
281
282 std::fs::remove_dir_all(&dir).ok();
283 }
284
285 #[test]
286 fn compact_freezes_and_archives() {
287 let store = Arc::new(CompactTextSessionStore);
288 let dir = std::env::temp_dir().join("compact_test_freeze");
289 let _ = std::fs::remove_dir_all(&dir);
290 std::fs::create_dir_all(&dir).expect("operation should succeed");
291 let path = dir.join("head.tlog");
292
293 for i in 0..10 {
294 let entry = SessionEntry {
295 id: uuid::Uuid::new_v4().to_string(),
296 parent_id: None,
297 timestamp: chrono::Utc::now(),
298 role: "user".into(),
299 content: format!("entry {i}"),
300 metadata: crate::SessionMetadata::default(),
301 };
302 store
303 .append_entry(&path, &entry)
304 .expect("operation should succeed");
305 }
306
307 let engine = CompactionEngine::new(store);
308 let result = engine
309 .compact_segment(&path, &dir, 3)
310 .expect("operation should succeed");
311
312 match result {
313 CompactionResult::Compacted {
314 original_count,
315 compacted_count,
316 ..
317 } => {
318 assert_eq!(original_count, 10);
319 assert!(compacted_count < original_count);
320 }
321 CompactionResult::Skipped => panic!("should have compacted"),
322 }
323
324 std::fs::remove_dir_all(&dir).ok();
325 }
326
327 #[test]
328 fn compact_skips_when_under_threshold() {
329 let store = Arc::new(CompactTextSessionStore);
330 let dir = std::env::temp_dir().join("compact_test_skip");
331 let _ = std::fs::remove_dir_all(&dir);
332 std::fs::create_dir_all(&dir).expect("operation should succeed");
333 let path = dir.join("head.tlog");
334
335 let entry = SessionEntry {
336 id: uuid::Uuid::new_v4().to_string(),
337 parent_id: None,
338 timestamp: chrono::Utc::now(),
339 role: "user".into(),
340 content: "single entry".into(),
341 metadata: crate::SessionMetadata::default(),
342 };
343 store
344 .append_entry(&path, &entry)
345 .expect("operation should succeed");
346
347 let engine = CompactionEngine::new(store);
348 let result = engine
349 .compact_segment(&path, &dir, 10)
350 .expect("operation should succeed");
351 assert!(matches!(result, CompactionResult::Skipped));
352
353 std::fs::remove_dir_all(&dir).ok();
354 }
355}