codei_session/
compact.rs

1use codei_llm::Message;
2
3/// Token budget for context window management.
4#[derive(Debug, Clone, Copy)]
5pub struct TokenBudget {
6    pub max_tokens: u32,
7    pub compaction_threshold: f32,
8}
9
10impl TokenBudget {
11    pub fn limit(&self) -> u32 {
12        ((self.max_tokens as f32) * self.compaction_threshold) as u32
13    }
14}
15
16/// Rough token estimate (chars / 4).
17pub fn estimate_tokens(messages: &[Message]) -> u32 {
18    messages
19        .iter()
20        .map(|m| {
21            let content_len = m.content.as_ref().map(|c| c.len()).unwrap_or(0);
22            let tools_len = m
23                .tool_calls
24                .as_ref()
25                .map(|calls| calls.iter().map(|c| c.arguments.len()).sum::<usize>())
26                .unwrap_or(0);
27            (content_len + tools_len) as u32 / 4 + 8
28        })
29        .sum()
30}
31
32/// Truncate older messages when over budget. Keeps system prompt and recent turns.
33pub fn compact_messages(messages: Vec<Message>, budget: TokenBudget) -> Vec<Message> {
34    if messages.is_empty() || estimate_tokens(&messages) <= budget.limit() {
35        return messages;
36    }
37
38    const KEEP_RECENT: usize = 12;
39    if messages.len() <= KEEP_RECENT + 1 {
40        return messages;
41    }
42
43    let system = messages.first().cloned();
44    let mut rest: Vec<Message> = messages.into_iter().skip(1).collect();
45    if rest.len() <= KEEP_RECENT {
46        return reconstruct(system, rest, 0);
47    }
48
49    let omitted = rest.len() - KEEP_RECENT;
50    rest.drain(0..omitted);
51    reconstruct(system, rest, omitted)
52}
53
54fn reconstruct(system: Option<Message>, kept: Vec<Message>, omitted: usize) -> Vec<Message> {
55    let mut out = Vec::new();
56    if let Some(sys) = system {
57        out.push(sys);
58    }
59    if omitted > 0 {
60        out.push(Message::user(format!(
61            "[Context compacted: {omitted} earlier messages omitted to fit the context window]"
62        )));
63    }
64    out.extend(kept);
65    out
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use codei_llm::Message;
72
73    #[test]
74    fn compacts_when_over_budget() {
75        let budget = TokenBudget {
76            max_tokens: 100,
77            compaction_threshold: 0.5,
78        };
79        let messages: Vec<Message> = std::iter::once(Message::system("sys"))
80            .chain((0..20).map(|i| Message::user(format!("message {i} {}", "x".repeat(50)))))
81            .collect();
82        let compacted = compact_messages(messages, budget);
83        assert!(compacted.len() < 21);
84        assert!(compacted[0].content.as_deref() == Some("sys"));
85    }
86}
codei_session/compact.rs

codei_session/
compact.rs
