codei_session/
compact.rs

1use codei_config::AgentConfig;
2use codei_llm::Message;
3
4use crate::model::{to_llm_messages, Role, Session, StoredMessage};
5
6/// Token budget for context window management.
7#[derive(Debug, Clone, Copy)]
8pub struct TokenBudget {
9    pub max_tokens: u32,
10    pub compaction_threshold: f32,
11    pub keep_messages: usize,
12}
13
14impl TokenBudget {
15    pub fn limit(&self) -> u32 {
16        ((self.max_tokens as f32) * self.compaction_threshold) as u32
17    }
18
19    pub fn from_agent(agent: &AgentConfig) -> Self {
20        Self {
21            max_tokens: agent.context_window_tokens,
22            compaction_threshold: agent.compaction_threshold,
23            keep_messages: agent.compaction_keep_messages.max(2) as usize,
24        }
25    }
26}
27
28/// Rough token estimate (chars / 4).
29pub fn estimate_tokens(messages: &[Message]) -> u32 {
30    messages
31        .iter()
32        .map(|m| {
33            let content_len = m.content.as_ref().map(|c| c.len()).unwrap_or(0);
34            let tools_len = m
35                .tool_calls
36                .as_ref()
37                .map(|calls| calls.iter().map(|c| c.arguments.len()).sum::<usize>())
38                .unwrap_or(0);
39            (content_len + tools_len) as u32 / 4 + 8
40        })
41        .sum()
42}
43
44/// Whether the session should be compacted based on estimated token usage.
45pub fn should_compact_session(session: &Session, system_prompt: &str, agent: &AgentConfig) -> bool {
46    let budget = TokenBudget::from_agent(agent);
47    let tokens = estimate_tokens(&to_llm_messages(session, system_prompt));
48    tokens > budget.limit() && session.messages.len() > budget.keep_messages
49}
50
51/// Format stored messages into a transcript for LLM summarization.
52pub fn format_transcript(messages: &[StoredMessage]) -> String {
53    let mut out = String::new();
54    for msg in messages {
55        let role = match msg.role {
56            Role::User => "User",
57            Role::Assistant => "Assistant",
58            Role::Tool => "Tool",
59            Role::System => continue,
60        };
61        let Some(text) = msg.text() else { continue };
62        let trimmed = text.trim();
63        if trimmed.is_empty() {
64            continue;
65        }
66        out.push_str(role);
67        out.push_str(": ");
68        out.push_str(&truncate_chars(trimmed, 6_000));
69        out.push_str("\n\n");
70    }
71    out
72}
73
74fn truncate_chars(text: &str, max_chars: usize) -> String {
75    if text.chars().count() <= max_chars {
76        return text.to_string();
77    }
78    let truncated: String = text.chars().take(max_chars).collect();
79    format!("{truncated}…")
80}
81
82/// Truncate older messages when over budget. Keeps system prompt and recent turns.
83/// Used as a last-resort guard when building the LLM request.
84pub fn compact_messages(messages: Vec<Message>, budget: TokenBudget) -> Vec<Message> {
85    if messages.is_empty() || estimate_tokens(&messages) <= budget.limit() {
86        return messages;
87    }
88
89    const MIN_KEEP: usize = 2;
90    let keep_recent = budget.keep_messages.max(MIN_KEEP);
91    if messages.len() <= keep_recent + 1 {
92        return messages;
93    }
94
95    let system = messages.first().cloned();
96    let mut rest: Vec<Message> = messages.into_iter().skip(1).collect();
97    if rest.len() <= keep_recent {
98        return reconstruct(system, rest, 0);
99    }
100
101    let omitted = rest.len() - keep_recent;
102    rest.drain(0..omitted);
103    reconstruct(system, rest, omitted)
104}
105
106fn reconstruct(system: Option<Message>, kept: Vec<Message>, omitted: usize) -> Vec<Message> {
107    let mut out = Vec::new();
108    if let Some(sys) = system {
109        out.push(sys);
110    }
111    if omitted > 0 {
112        out.push(Message::user(format!(
113            "[Context compacted: {omitted} earlier messages omitted to fit the context window]"
114        )));
115    }
116    out.extend(kept);
117    out
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use codei_llm::Message;
124
125    #[test]
126    fn compacts_when_over_budget() {
127        let budget = TokenBudget {
128            max_tokens: 100,
129            compaction_threshold: 0.5,
130            keep_messages: 4,
131        };
132        let messages: Vec<Message> = std::iter::once(Message::system("sys"))
133            .chain((0..20).map(|i| Message::user(format!("message {i} {}", "x".repeat(50)))))
134            .collect();
135        let compacted = compact_messages(messages, budget);
136        assert!(compacted.len() < 21);
137        assert!(compacted[0].content.as_deref() == Some("sys"));
138    }
139
140    #[test]
141    fn should_compact_when_over_threshold() {
142        use codei_config::AgentConfig;
143
144        let mut session = Session::new(std::path::PathBuf::from("/tmp"));
145        for i in 0..30 {
146            session.push_user(format!("message {i} {}", "x".repeat(80)));
147        }
148        let agent = AgentConfig {
149            context_window_tokens: 200,
150            compaction_threshold: 0.5,
151            compaction_keep_messages: 6,
152            ..Default::default()
153        };
154        assert!(should_compact_session(&session, "system prompt", &agent));
155    }
156
157    #[test]
158    fn compact_with_summary_inserts_summary_message() {
159        let mut session = Session::new(std::path::PathBuf::from("/tmp"));
160        for i in 0..5 {
161            session.push_user(format!("msg {i}"));
162        }
163        session.compact_with_summary(2, "summary text".into());
164        assert_eq!(session.messages.len(), 3);
165        assert!(session.messages[0].text().unwrap().contains("summary text"));
166        assert!(session.messages[1].text().unwrap().contains("msg 3"));
167    }
168}
codei_session/compact.rs

codei_session/
compact.rs
