machi_compaction/
token_threshold.rs

1//! Compact when estimated tokens exceed a threshold.
2
3use machi_types::{ErrorCode, MachiError, Message};
4
5use crate::max_messages::compact_max_messages;
6use crate::strategy::{CompactionOutcome, CompactionStrategy};
7
8/// Drop oldest non-system messages when `token_estimate` exceeds `max_tokens`.
9///
10/// Uses the same tail-preserving algorithm as [`crate::MaxMessages`], keeping
11/// at most `keep_messages` after the optional leading system message.
12#[derive(Debug, Clone, Copy)]
13pub struct TokenThreshold {
14    /// Token estimate that triggers compaction (must be >= 1).
15    pub max_tokens: u64,
16    /// Message count retained after compaction (must be >= 1).
17    pub keep_messages: usize,
18}
19
20impl TokenThreshold {
21    /// Construct a token-threshold strategy.
22    ///
23    /// # Errors
24    ///
25    /// Returns error when `max_tokens == 0` or `keep_messages == 0`.
26    pub fn new(max_tokens: u64, keep_messages: usize) -> Result<Self, MachiError> {
27        if max_tokens == 0 {
28            return Err(MachiError::new(
29                ErrorCode::CompactionFailed,
30                "TokenThreshold max_tokens must be >= 1",
31            ));
32        }
33        if keep_messages == 0 {
34            return Err(MachiError::new(
35                ErrorCode::CompactionFailed,
36                "TokenThreshold keep_messages must be >= 1",
37            ));
38        }
39        Ok(Self {
40            max_tokens,
41            keep_messages,
42        })
43    }
44}
45
46impl CompactionStrategy for TokenThreshold {
47    fn name(&self) -> &'static str {
48        "token_threshold"
49    }
50
51    fn should_compact(&self, messages: &[Message], token_estimate: u64) -> bool {
52        token_estimate > self.max_tokens && messages.len() > self.keep_messages
53    }
54
55    fn compact(&self, messages: Vec<Message>) -> Result<CompactionOutcome, MachiError> {
56        if messages.len() <= self.keep_messages {
57            return Ok(CompactionOutcome {
58                messages,
59                changed: false,
60                strategy: self.name(),
61            });
62        }
63        let compacted = compact_max_messages(messages, self.keep_messages);
64        Ok(CompactionOutcome {
65            messages: compacted,
66            changed: true,
67            strategy: self.name(),
68        })
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use machi_types::Message;
75
76    use super::*;
77
78    #[test]
79    fn triggers_on_token_estimate() {
80        let s = TokenThreshold::new(10, 3).expect("new");
81        let msgs = vec![
82            Message::system("sys"),
83            Message::user("1"),
84            Message::user("2"),
85            Message::user("3"),
86            Message::user("4"),
87        ];
88        assert!(!s.should_compact(&msgs, 5));
89        assert!(s.should_compact(&msgs, 11));
90        let out = s.compact(msgs).expect("compact");
91        assert!(out.changed);
92        assert_eq!(out.messages.len(), 3);
93        assert_eq!(
94            out.messages.first().map(Message::text).as_deref(),
95            Some("sys")
96        );
97    }
98
99    #[test]
100    fn rejects_zero() {
101        assert!(TokenThreshold::new(0, 3).is_err());
102        assert!(TokenThreshold::new(10, 0).is_err());
103    }
104}
machi_compaction/token_threshold.rs

machi_compaction/
token_threshold.rs
