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(Box::new(Self::Evict(0.6)), Box::new(Self::Retain(10)))
78 }
79}
80
81impl CompactionStrategy {
82 pub fn calculate_eviction_range(
91 &self,
92 messages: &[MessageMeta],
93 retention_window: usize,
94 ) -> Option<EvictionRange> {
95 if messages.len() <= retention_window {
96 return None; }
98
99 let raw_end = self.calculate_raw_end(messages.len(), retention_window);
100
101 let start = Self::find_safe_start(messages);
103
104 if start >= raw_end {
105 return None; }
107
108 let end = Self::adjust_end_for_tool_safety(messages, raw_end, retention_window);
110
111 if start >= end {
112 return None;
113 }
114
115 Some(EvictionRange::new(start, end))
116 }
117
118 fn calculate_raw_end(&self, total: usize, retention_window: usize) -> usize {
120 match self {
121 Self::Evict(fraction) => {
122 let evict_count = (total as f64 * fraction).floor() as usize;
123 total.saturating_sub(retention_window).min(evict_count)
124 }
125 Self::Retain(keep) => total.saturating_sub(*keep.max(&retention_window)),
126 Self::Min(a, b) => {
127 let end_a = a.calculate_raw_end(total, retention_window);
128 let end_b = b.calculate_raw_end(total, retention_window);
129 end_a.min(end_b)
130 }
131 Self::Max(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.max(end_b)
135 }
136 }
137 }
138
139 fn find_safe_start(messages: &[MessageMeta]) -> usize {
141 messages
142 .iter()
143 .position(|m| m.role == MessageRole::Assistant)
144 .unwrap_or(0)
145 }
146
147 fn adjust_end_for_tool_safety(
149 messages: &[MessageMeta],
150 mut end: usize,
151 retention_window: usize,
152 ) -> usize {
153 let min_end = messages.len().saturating_sub(retention_window);
154
155 if end > min_end {
157 end = min_end;
158 }
159
160 if end == 0 || end >= messages.len() {
161 return end;
162 }
163
164 let last_evicted = &messages[end - 1];
167
168 if last_evicted.has_tool_call {
169 if let Some(tool_id) = &last_evicted.tool_id {
172 for i in end..messages.len().min(end + 5) {
173 if messages[i].is_tool_result && messages[i].tool_id.as_ref() == Some(tool_id) {
174 end = i + 1;
176 break;
177 }
178 }
179 }
180 }
181
182 let msg_at_end = messages.get(end);
184 if let Some(msg) = msg_at_end
185 && msg.is_tool_result
186 {
187 while end > 0 {
190 let prev = &messages[end - 1];
191 if prev.is_tool_result || prev.has_tool_call {
192 end -= 1;
193 } else {
194 break;
195 }
196 }
197 }
198
199 while end > 0 && end < messages.len() {
201 if messages[end].is_tool_result {
202 end -= 1;
203 } else {
204 break;
205 }
206 }
207
208 end
209 }
210
211 pub fn filter_droppable(messages: &[MessageMeta], range: &EvictionRange) -> Vec<usize> {
214 (range.start..range.end)
215 .filter(|&i| !messages[i].droppable)
216 .collect()
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 fn make_messages(roles: &[(MessageRole, bool, bool)]) -> Vec<MessageMeta> {
225 roles
226 .iter()
227 .enumerate()
228 .map(|(i, (role, has_tool_call, is_tool_result))| MessageMeta {
229 index: i,
230 role: *role,
231 droppable: false,
232 has_tool_call: *has_tool_call,
233 is_tool_result: *is_tool_result,
234 tool_id: if *has_tool_call || *is_tool_result {
235 Some(format!("tool_{}", i))
236 } else {
237 None
238 },
239 token_count: 100,
240 })
241 .collect()
242 }
243
244 #[test]
245 fn test_eviction_range_empty() {
246 let strategy = CompactionStrategy::Retain(10);
247 let messages = make_messages(&[
248 (MessageRole::System, false, false),
249 (MessageRole::User, false, false),
250 (MessageRole::Assistant, false, false),
251 ]);
252
253 let range = strategy.calculate_eviction_range(&messages, 5);
254 assert!(range.is_none());
255 }
256
257 #[test]
258 fn test_eviction_starts_at_assistant() {
259 let strategy = CompactionStrategy::Evict(0.5);
260 let messages = make_messages(&[
261 (MessageRole::System, false, false),
262 (MessageRole::User, false, false),
263 (MessageRole::Assistant, false, false),
264 (MessageRole::User, false, false),
265 (MessageRole::Assistant, false, false),
266 (MessageRole::User, false, false),
267 (MessageRole::Assistant, false, false),
268 ]);
269
270 let range = strategy.calculate_eviction_range(&messages, 2);
271 assert!(range.is_some());
272 let range = range.unwrap();
273 assert_eq!(range.start, 2);
275 }
276
277 #[test]
278 fn test_tool_call_result_adjacency() {
279 let mut messages = make_messages(&[
280 (MessageRole::System, false, false),
281 (MessageRole::User, false, false),
282 (MessageRole::Assistant, true, false), (MessageRole::Tool, false, true), (MessageRole::Assistant, false, false),
285 (MessageRole::User, false, false),
286 (MessageRole::Assistant, false, false),
287 ]);
288
289 messages[2].tool_id = Some("call_1".to_string());
291 messages[3].tool_id = Some("call_1".to_string());
292
293 let strategy = CompactionStrategy::Retain(2);
294 let range = strategy.calculate_eviction_range(&messages, 2);
295
296 if let Some(range) = range {
298 if range.end > 2 && range.end <= 3 {
300 panic!("Eviction split tool call from result!");
301 }
302 }
303 }
304
305 #[test]
306 fn test_filter_droppable() {
307 let mut messages = make_messages(&[
308 (MessageRole::System, false, false),
309 (MessageRole::User, false, false),
310 (MessageRole::Assistant, false, false),
311 (MessageRole::User, false, false), (MessageRole::Assistant, false, false),
313 ]);
314 messages[3].droppable = true;
315
316 let range = EvictionRange::new(0, 5);
317 let non_droppable = CompactionStrategy::filter_droppable(&messages, &range);
318
319 assert_eq!(non_droppable.len(), 4);
320 assert!(!non_droppable.contains(&3));
321 }
322
323 #[test]
324 fn test_min_strategy() {
325 let strategy = CompactionStrategy::Min(
326 Box::new(CompactionStrategy::Evict(0.8)),
327 Box::new(CompactionStrategy::Retain(5)),
328 );
329
330 let messages = make_messages(&vec![(MessageRole::Assistant, false, false); 10]);
336
337 let range = strategy.calculate_eviction_range(&messages, 3);
338 assert!(range.is_some());
339 }
341}