Skip to main content

codei_agent/
compact.rs

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