Skip to main content

machi_compaction/
max_messages.rs

1//! Keep at most N messages, preserving leading system and tail.
2
3use machi_types::{ErrorCode, MachiError, Message};
4
5use crate::select::{apply_range, select_compaction_range, tool_pair_invariant_holds};
6use crate::strategy::{CompactionOutcome, CompactionStrategy};
7
8/// Drop oldest non-system messages until `max` remains.
9#[derive(Debug, Clone, Copy)]
10pub struct MaxMessages {
11    /// Maximum messages retained (including system).
12    pub max: usize,
13}
14
15impl MaxMessages {
16    /// Construct with a positive max.
17    ///
18    /// # Errors
19    ///
20    /// Returns error when `max == 0`.
21    pub fn new(max: usize) -> Result<Self, MachiError> {
22        if max == 0 {
23            return Err(MachiError::new(
24                ErrorCode::CompactionFailed,
25                "MaxMessages max must be >= 1",
26            ));
27        }
28        Ok(Self { max })
29    }
30}
31
32impl CompactionStrategy for MaxMessages {
33    fn name(&self) -> &'static str {
34        "max_messages"
35    }
36
37    fn should_compact(&self, messages: &[Message], _token_estimate: u64) -> bool {
38        messages.len() > self.max
39    }
40
41    fn compact(&self, messages: Vec<Message>) -> Result<CompactionOutcome, MachiError> {
42        if messages.len() <= self.max {
43            return Ok(CompactionOutcome {
44                messages,
45                changed: false,
46                strategy: self.name(),
47            });
48        }
49        let before = messages.len();
50        let compacted = compact_max_messages(messages, self.max);
51        let changed = compacted.len() != before;
52        Ok(CompactionOutcome {
53            messages: compacted,
54            changed,
55            strategy: self.name(),
56        })
57    }
58}
59
60/// Shared algorithm used by runtime `VecConversationState` and this strategy.
61///
62/// Uses [`select_compaction_range`] so tool-result runs are never split mid-pair.
63/// When no safe split exists, returns the input unchanged (same as other strategies).
64#[must_use]
65pub fn compact_max_messages(messages: Vec<Message>, max_messages: usize) -> Vec<Message> {
66    if max_messages == 0 || messages.len() <= max_messages {
67        return messages;
68    }
69    let Some(range) = select_compaction_range(&messages, max_messages) else {
70        return messages;
71    };
72    let out = apply_range(messages, range, None);
73    debug_assert!(
74        tool_pair_invariant_holds(&out),
75        "compact_max_messages must preserve tool-pair invariant"
76    );
77    out
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn keeps_system_and_tail() {
86        let s = MaxMessages::new(3).expect("max");
87        let out = s
88            .compact(vec![
89                Message::system("sys"),
90                Message::user("1"),
91                Message::user("2"),
92                Message::user("3"),
93                Message::user("4"),
94            ])
95            .expect("compact");
96        assert!(out.changed);
97        assert_eq!(out.messages.len(), 3);
98        assert_eq!(
99            out.messages.first().map(Message::text).as_deref(),
100            Some("sys")
101        );
102        assert_eq!(out.messages.get(2).map(Message::text).as_deref(), Some("4"));
103    }
104}