syncable_cli/agent/compact/
strategy.rs1use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12pub enum MessageRole {
13 System,
14 User,
15 Assistant,
16 Tool,
17}
18
19#[derive(Debug, Clone)]
21pub struct MessageMeta {
22 pub index: usize,
24 pub role: MessageRole,
26 pub droppable: bool,
28 pub has_tool_call: bool,
30 pub is_tool_result: bool,
32 pub tool_id: Option<String>,
34 pub token_count: usize,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
40pub struct EvictionRange {
41 pub start: usize,
43 pub end: usize,
45}
46
47impl EvictionRange {
48 pub fn new(start: usize, end: usize) -> Self {
49 Self { start, end }
50 }
51
52 pub fn len(&self) -> usize {
53 self.end.saturating_sub(self.start)
54 }
55
56 pub fn is_empty(&self) -> bool {
57 self.len() == 0
58 }
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63pub enum CompactionStrategy {
64 Evict(f64),
66 Retain(usize),
68 Min(Box<CompactionStrategy>, Box<CompactionStrategy>),
70 Max(Box<CompactionStrategy>, Box<CompactionStrategy>),
72}
73
74impl Default for CompactionStrategy {
75 fn default() -> Self {
76 Self::Min(
78 Box::new(Self::Evict(0.6)),
79 Box::new(Self::Retain(10)),
80 )
81 }
82}
83
84impl CompactionStrategy {
85 pub fn calculate_eviction_range(
94 &self,
95 messages: &[MessageMeta],
96 retention_window: usize,
97 ) -> Option<EvictionRange> {
98 if messages.len() <= retention_window {
99 return None; }
101
102 let raw_end = self.calculate_raw_end(messages.len(), retention_window);
103
104 let start = Self::find_safe_start(messages);
106
107 if start >= raw_end {
108 return None; }
110
111 let end = Self::adjust_end_for_tool_safety(messages, raw_end, retention_window);
113
114 if start >= end {
115 return None;
116 }
117
118 Some(EvictionRange::new(start, end))
119 }
120
121 fn calculate_raw_end(&self, total: usize, retention_window: usize) -> usize {
123 match self {
124 Self::Evict(fraction) => {
125 let evict_count = (total as f64 * fraction).floor() as usize;
126 total.saturating_sub(retention_window).min(evict_count)
127 }
128 Self::Retain(keep) => {
129 total.saturating_sub(*keep.max(&retention_window))
130 }
131 Self::Min(a, b) => {
132 let end_a = a.calculate_raw_end(total, retention_window);
133 let end_b = b.calculate_raw_end(total, retention_window);
134 end_a.min(end_b)
135 }
136 Self::Max(a, b) => {
137 let end_a = a.calculate_raw_end(total, retention_window);
138 let end_b = b.calculate_raw_end(total, retention_window);
139 end_a.max(end_b)
140 }
141 }
142 }
143
144 fn find_safe_start(messages: &[MessageMeta]) -> usize {
146 messages
147 .iter()
148 .position(|m| m.role == MessageRole::Assistant)
149 .unwrap_or(0)
150 }
151
152 fn adjust_end_for_tool_safety(
154 messages: &[MessageMeta],
155 mut end: usize,
156 retention_window: usize,
157 ) -> usize {
158 let min_end = messages.len().saturating_sub(retention_window);
159
160 if end > min_end {
162 end = min_end;
163 }
164
165 if end == 0 || end >= messages.len() {
166 return end;
167 }
168
169 let last_evicted = &messages[end - 1];
172
173 if last_evicted.has_tool_call {
174 if let Some(tool_id) = &last_evicted.tool_id {
177 for i in end..messages.len().min(end + 5) {
178 if messages[i].is_tool_result
179 && messages[i].tool_id.as_ref() == Some(tool_id)
180 {
181 end = i + 1;
183 break;
184 }
185 }
186 }
187 }
188
189 let msg_at_end = messages.get(end);
191 if let Some(msg) = msg_at_end {
192 if msg.is_tool_result {
193 while end > 0 {
196 let prev = &messages[end - 1];
197 if prev.is_tool_result || prev.has_tool_call {
198 end -= 1;
199 } else {
200 break;
201 }
202 }
203 }
204 }
205
206 while end > 0 && end < messages.len() {
208 if messages[end].is_tool_result {
209 end -= 1;
210 } else {
211 break;
212 }
213 }
214
215 end
216 }
217
218 pub fn filter_droppable(messages: &[MessageMeta], range: &EvictionRange) -> Vec<usize> {
221 (range.start..range.end)
222 .filter(|&i| !messages[i].droppable)
223 .collect()
224 }
225}
226
227#[cfg(test)]
228mod tests {
229 use super::*;
230
231 fn make_messages(roles: &[(MessageRole, bool, bool)]) -> Vec<MessageMeta> {
232 roles
233 .iter()
234 .enumerate()
235 .map(|(i, (role, has_tool_call, is_tool_result))| MessageMeta {
236 index: i,
237 role: *role,
238 droppable: false,
239 has_tool_call: *has_tool_call,
240 is_tool_result: *is_tool_result,
241 tool_id: if *has_tool_call || *is_tool_result {
242 Some(format!("tool_{}", i))
243 } else {
244 None
245 },
246 token_count: 100,
247 })
248 .collect()
249 }
250
251 #[test]
252 fn test_eviction_range_empty() {
253 let strategy = CompactionStrategy::Retain(10);
254 let messages = make_messages(&[
255 (MessageRole::System, false, false),
256 (MessageRole::User, false, false),
257 (MessageRole::Assistant, false, false),
258 ]);
259
260 let range = strategy.calculate_eviction_range(&messages, 5);
261 assert!(range.is_none());
262 }
263
264 #[test]
265 fn test_eviction_starts_at_assistant() {
266 let strategy = CompactionStrategy::Evict(0.5);
267 let messages = make_messages(&[
268 (MessageRole::System, false, false),
269 (MessageRole::User, false, false),
270 (MessageRole::Assistant, false, false),
271 (MessageRole::User, false, false),
272 (MessageRole::Assistant, false, false),
273 (MessageRole::User, false, false),
274 (MessageRole::Assistant, false, false),
275 ]);
276
277 let range = strategy.calculate_eviction_range(&messages, 2);
278 assert!(range.is_some());
279 let range = range.unwrap();
280 assert_eq!(range.start, 2);
282 }
283
284 #[test]
285 fn test_tool_call_result_adjacency() {
286 let mut messages = make_messages(&[
287 (MessageRole::System, false, false),
288 (MessageRole::User, false, false),
289 (MessageRole::Assistant, true, false), (MessageRole::Tool, false, true), (MessageRole::Assistant, false, false),
292 (MessageRole::User, false, false),
293 (MessageRole::Assistant, false, false),
294 ]);
295
296 messages[2].tool_id = Some("call_1".to_string());
298 messages[3].tool_id = Some("call_1".to_string());
299
300 let strategy = CompactionStrategy::Retain(2);
301 let range = strategy.calculate_eviction_range(&messages, 2);
302
303 if let Some(range) = range {
305 if range.end > 2 && range.end <= 3 {
307 panic!("Eviction split tool call from result!");
308 }
309 }
310 }
311
312 #[test]
313 fn test_filter_droppable() {
314 let mut messages = make_messages(&[
315 (MessageRole::System, false, false),
316 (MessageRole::User, false, false),
317 (MessageRole::Assistant, false, false),
318 (MessageRole::User, false, false), (MessageRole::Assistant, false, false),
320 ]);
321 messages[3].droppable = true;
322
323 let range = EvictionRange::new(0, 5);
324 let non_droppable = CompactionStrategy::filter_droppable(&messages, &range);
325
326 assert_eq!(non_droppable.len(), 4);
327 assert!(!non_droppable.contains(&3));
328 }
329
330 #[test]
331 fn test_min_strategy() {
332 let strategy = CompactionStrategy::Min(
333 Box::new(CompactionStrategy::Evict(0.8)),
334 Box::new(CompactionStrategy::Retain(5)),
335 );
336
337 let messages = make_messages(&vec![
343 (MessageRole::Assistant, false, false); 10
344 ]);
345
346 let range = strategy.calculate_eviction_range(&messages, 3);
347 assert!(range.is_some());
348 }
350}