Skip to main content

motosan_agent_loop/session/
compaction.rs

1//! Stateful compaction — persist compaction decisions as
2//! `SessionEntry::Custom { kind: "compaction", ... }` markers.
3
4use std::sync::Arc;
5
6use crate::error::Result;
7use crate::llm::LlmClient;
8use crate::message::{Message, Role};
9use crate::session::entry::{EntryId, StoredEntry};
10
11pub trait CompactionStrategy: Send + Sync {
12    fn should_compact(&self, messages: &[Message]) -> bool;
13    fn select_cutoff(&self, entries: &[StoredEntry]) -> Option<EntryId>;
14    fn summarizer_system_prompt(&self) -> &str {
15        "You are a conversation summarizer. Produce a concise summary of the following conversation history, preserving all important context, decisions, and information needed to continue the conversation."
16    }
17}
18
19#[derive(Debug, Clone)]
20pub struct ThresholdStrategy {
21    pub threshold: f32,
22    pub max_context_tokens: usize,
23    pub keep_turns: usize,
24}
25
26impl Default for ThresholdStrategy {
27    fn default() -> Self {
28        Self {
29            threshold: 0.8,
30            max_context_tokens: 200_000,
31            keep_turns: 3,
32        }
33    }
34}
35
36impl CompactionStrategy for ThresholdStrategy {
37    fn should_compact(&self, messages: &[Message]) -> bool {
38        let total: usize = messages.iter().map(|m| m.approx_visible_chars() / 4).sum();
39        let budget = (self.threshold * self.max_context_tokens as f32) as usize;
40        total > budget
41    }
42
43    fn select_cutoff(&self, entries: &[StoredEntry]) -> Option<EntryId> {
44        let mut user_count = 0;
45        for stored in entries.iter().rev() {
46            if let Some(m) = stored.entry.as_message() {
47                if matches!(m.role(), Role::User) {
48                    user_count += 1;
49                    if user_count >= self.keep_turns {
50                        return Some(stored.id.clone());
51                    }
52                }
53            }
54        }
55        None
56    }
57}
58
59#[derive(Debug, Clone)]
60pub struct CompactionResult {
61    pub marker_id: EntryId,
62    pub summary: String,
63    pub first_kept_entry_id: EntryId,
64}
65
66pub(crate) async fn run_compaction(
67    strategy: &dyn CompactionStrategy,
68    summarizer: Arc<dyn LlmClient>,
69    entries: &[StoredEntry],
70    projected_messages: &[Message],
71) -> Result<Option<(EntryId, String)>> {
72    if !strategy.should_compact(projected_messages) {
73        return Ok(None);
74    }
75    let Some(first_kept_id) = strategy.select_cutoff(entries) else {
76        return Ok(None);
77    };
78
79    let cutoff_pos = match entries.iter().position(|se| se.id == first_kept_id) {
80        Some(pos) => pos,
81        None => {
82            tracing::warn!(
83                "select_cutoff returned entry id '{}' not found in entries; \
84                 summarizing full history as fallback",
85                first_kept_id
86            );
87            entries.len()
88        }
89    };
90
91    let mut history_text = String::new();
92    for stored in &entries[..cutoff_pos] {
93        if let Some(m) = stored.entry.as_message() {
94            let role = match m.role() {
95                Role::System => "System",
96                Role::User => "User",
97                Role::Assistant => "Assistant",
98                Role::Tool => "Tool",
99            };
100            history_text.push_str(&format!("[{role}]: {}\n", m.text()));
101        }
102    }
103
104    if history_text.is_empty() {
105        return Ok(None);
106    }
107
108    let summary = summarizer
109        .simple_call(strategy.summarizer_system_prompt(), &history_text)
110        .await?;
111
112    Ok(Some((first_kept_id, summary)))
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::message::Message;
119    use crate::session::entry::SessionEntry;
120
121    fn se(id: &str, entry: SessionEntry) -> StoredEntry {
122        StoredEntry::new(id.to_string(), entry)
123    }
124
125    #[test]
126    fn threshold_strategy_below_budget_returns_false() {
127        let s = ThresholdStrategy::default();
128        let msgs = vec![Message::user("short")];
129        assert!(!s.should_compact(&msgs));
130    }
131
132    #[test]
133    fn threshold_strategy_over_budget_returns_true() {
134        let s = ThresholdStrategy {
135            threshold: 0.8,
136            max_context_tokens: 100,
137            keep_turns: 2,
138        };
139        let big = "x".repeat(400);
140        let msgs = vec![Message::user(&big)];
141        assert!(s.should_compact(&msgs));
142    }
143
144    #[test]
145    fn threshold_strategy_selects_kept_turns_from_tail() {
146        let s = ThresholdStrategy {
147            keep_turns: 2,
148            ..ThresholdStrategy::default()
149        };
150        let entries = vec![
151            se("1", SessionEntry::message(Message::user("u1"))),
152            se("2", SessionEntry::message(Message::assistant("a1"))),
153            se("3", SessionEntry::message(Message::user("u2"))),
154            se("4", SessionEntry::message(Message::assistant("a2"))),
155            se("5", SessionEntry::message(Message::user("u3"))),
156        ];
157        let cutoff = s.select_cutoff(&entries);
158        assert_eq!(cutoff.as_deref(), Some("3"));
159    }
160
161    #[test]
162    fn select_cutoff_over_active_branch_ignores_abandoned_branches() {
163        use crate::session::projection::active_branch;
164
165        // Branch 1: A(user), B(user). Branch 2 forks from A: C(user).
166        // Active leaf = C; active branch = A -> C. keep_turns=1 must select
167        // the last user message on the *active* branch (C), never B.
168        let entries = vec![
169            se("A", SessionEntry::message(Message::user("a"))),
170            se("B", SessionEntry::message(Message::user("b"))),
171            StoredEntry::with_parent(
172                "C".to_string(),
173                Some("A".to_string()),
174                SessionEntry::message(Message::user("c")),
175            ),
176        ];
177        let branch = active_branch(&entries);
178        let strategy = ThresholdStrategy {
179            keep_turns: 1,
180            ..ThresholdStrategy::default()
181        };
182        assert_eq!(strategy.select_cutoff(&branch).as_deref(), Some("C"));
183    }
184}