Skip to main content

codei_agent/
compact.rs

1use codei_llm::{collect_response, ChatRequest, LlmProvider, Message};
2use codei_session::{
3    cap_output_tokens, format_transcript, should_compact_session, Session, SessionStore,
4};
5use tracing::warn;
6
7use crate::error::AgentError;
8use crate::loop_::AgentLoop;
9
10const SUMMARIZE_SYSTEM: &str = r#"You compress conversation history for a coding assistant context window.
11
12Produce a dense summary that preserves:
13- User goals, requirements, and constraints
14- Key decisions, conclusions, and agreed approaches
15- Important file paths, symbols, commands, and code changes
16- Unresolved tasks, errors, and open questions
17
18Do not invent facts. Omit small talk. Use the same language as the conversation."#;
19
20impl AgentLoop {
21    /// Summarize older session messages with the LLM and keep recent turns intact.
22    pub async fn compact_session(
23        &self,
24        session: &mut Session,
25        store: &SessionStore,
26    ) -> Result<(), AgentError> {
27        let agent_cfg = &self.config().config.agent;
28        let keep = agent_cfg.compaction_keep_messages.max(2) as usize;
29        if session.messages.len() <= keep {
30            return Ok(());
31        }
32
33        let split = session.messages.len() - keep;
34        let transcript = format_transcript(&session.messages[..split]);
35        if transcript.trim().is_empty() {
36            session.compact(keep);
37            store.save(session)?;
38            return Ok(());
39        }
40
41        let model = self.model().read().expect("model lock poisoned").clone();
42        let provider = self
43            .provider()
44            .read()
45            .expect("provider lock poisoned")
46            .clone();
47        let language = &self.config().config.defaults.language;
48        let max_tokens = cap_output_tokens(
49            &[
50                Message::system(SUMMARIZE_SYSTEM),
51                Message::user(format!(
52                "Summarize the conversation below. Write the summary in {language}.\n\n{transcript}"
53            )),
54            ],
55            None,
56            agent_cfg
57                .compaction_summary_max_tokens
58                .min(self.config().config.defaults.max_tokens),
59            agent_cfg.context_window_tokens,
60        );
61
62        let summary = match summarize_transcript(
63            provider.as_ref(),
64            &model,
65            &transcript,
66            language,
67            max_tokens,
68        )
69        .await
70        {
71            Ok(summary) => summary,
72            Err(err) => {
73                warn!(error = %err, "LLM compaction failed; falling back to truncation");
74                session.compact(keep);
75                store.save(session)?;
76                return Ok(());
77            }
78        };
79
80        session.compact_with_summary(keep, summary);
81        store.save(session)?;
82        Ok(())
83    }
84
85    pub(crate) async fn compact_session_if_needed(
86        &self,
87        session: &mut Session,
88        store: &SessionStore,
89    ) -> Result<bool, AgentError> {
90        if !should_compact_session(session, self.system_prompt(), &self.config().config.agent) {
91            return Ok(false);
92        }
93        self.compact_session(session, store).await?;
94        Ok(true)
95    }
96}
97
98async fn summarize_transcript(
99    provider: &dyn LlmProvider,
100    model: &str,
101    transcript: &str,
102    language: &str,
103    max_tokens: u32,
104) -> Result<String, codei_llm::LlmError> {
105    let request = ChatRequest {
106        model: model.to_string(),
107        messages: vec![
108            Message::system(SUMMARIZE_SYSTEM),
109            Message::user(format!(
110                "Summarize the conversation below. Write the summary in {language}.\n\n{transcript}"
111            )),
112        ],
113        tools: None,
114        temperature: Some(0.1),
115        max_tokens: Some(max_tokens.max(256)),
116    };
117
118    let stream = provider.chat(request).await?;
119    let response = collect_response(stream).await?;
120    let summary = response.content.trim().to_string();
121    if summary.is_empty() {
122        return Err(codei_llm::LlmError::Config(
123            "compaction summary was empty".into(),
124        ));
125    }
126    Ok(summary)
127}
128
129#[cfg(test)]
130mod tests {
131    use codei_config::AgentConfig;
132    use codei_session::TokenBudget;
133
134    #[test]
135    fn budget_limit_uses_threshold() {
136        let budget = TokenBudget::from_agent(&AgentConfig {
137            context_window_tokens: 1000,
138            compaction_threshold: 0.8,
139            ..Default::default()
140        });
141        assert_eq!(budget.limit(), 800);
142    }
143}