1use zeph_context::summarization::{SummarizationDeps, summarize_structured};
9use zeph_llm::provider::{Message, MessagePart, Role};
10
11use crate::condenser::{CondensationResult, Condenser, validate_non_overlap};
12use crate::error::SessionError;
13use crate::event::{SessionEvent, SessionEventEnvelope};
14use crate::replay::{ReconstructedState, ReplayEngine};
15
16const CONDENSATION_GUIDELINES: &str = "Summarize the conversation so far, preserving the \
17 session's intent, files modified, decisions made, open questions, and next steps. This \
18 summary durably replaces the condensed portion of the event log — be precise, do not invent \
19 details.";
20
21pub struct LlmCondenser {
27 deps: SummarizationDeps,
28 threshold: f64,
31 keep_recent: usize,
34}
35
36impl LlmCondenser {
37 #[must_use]
40 pub fn new(deps: SummarizationDeps, threshold: f64, keep_recent: usize) -> Self {
41 Self {
42 deps,
43 threshold,
44 keep_recent,
45 }
46 }
47}
48
49impl Condenser for LlmCondenser {
50 async fn should_condense(&self, state: &ReconstructedState, budget_used_fraction: f64) -> bool {
51 budget_used_fraction >= self.threshold && state.messages.len() > self.keep_recent
52 }
53
54 #[tracing::instrument(
55 name = "session.condenser.condense",
56 skip_all,
57 level = "info",
58 fields(event_count = events.len(), last_condensed_seq)
59 )]
60 async fn condense(
61 &self,
62 events: &[SessionEventEnvelope],
63 last_condensed_seq: u64,
64 ) -> Result<CondensationResult, SessionError> {
65 let uncondensed: Vec<SessionEventEnvelope> = events
72 .iter()
73 .filter(|e| e.seq > last_condensed_seq)
74 .cloned()
75 .collect();
76
77 let boundaries: Vec<usize> = uncondensed
81 .iter()
82 .enumerate()
83 .filter(|(_, e)| {
84 matches!(
85 e.kind,
86 SessionEvent::UserMessage { .. } | SessionEvent::AssistantMessage { .. }
87 )
88 })
89 .map(|(i, _)| i)
90 .collect();
91
92 if boundaries.len() <= self.keep_recent {
93 return Err(SessionError::CondensationOverlap(format!(
94 "not enough events to condense: {} message(s) available, keep_recent={}",
95 boundaries.len(),
96 self.keep_recent
97 )));
98 }
99
100 let cutoff = boundaries[boundaries.len() - self.keep_recent];
101 let to_condense = &uncondensed[..cutoff];
102 let lo = to_condense
103 .first()
104 .map_or(last_condensed_seq + 1, |e| e.seq);
105 let hi = to_condense.last().map_or(last_condensed_seq, |e| e.seq);
106 validate_non_overlap(last_condensed_seq, (lo, hi))?;
107
108 let folded = ReplayEngine::fold(to_condense.to_vec(), None);
109 let tokens_before: usize = folded
110 .messages
111 .iter()
112 .map(|m| self.deps.token_counter.count_message_tokens(m))
113 .sum();
114
115 let summary = summarize_structured(&self.deps, &folded.messages, CONDENSATION_GUIDELINES)
116 .await
117 .map_err(SessionError::Llm)?;
118
119 let summary_message = Message::from_parts(
120 Role::System,
121 vec![MessagePart::Summary {
122 text: summary.to_markdown(),
123 }],
124 );
125 let tokens_after = self
126 .deps
127 .token_counter
128 .count_message_tokens(&summary_message);
129
130 Ok(CondensationResult {
131 replaced_range: (lo, hi),
132 summary,
133 tokens_before: u32::try_from(tokens_before).unwrap_or(u32::MAX),
134 tokens_after: u32::try_from(tokens_after).unwrap_or(u32::MAX),
135 })
136 }
137}
138
139#[cfg(test)]
140mod tests {
141 use std::time::Duration;
142
143 use zeph_llm::any::AnyProvider;
144 use zeph_llm::mock::MockProvider;
145
146 use super::*;
147 use crate::event::SessionEventEnvelope;
148
149 struct WordCountTokenCounter;
150 impl zeph_context::summarization::MessageTokenCounter for WordCountTokenCounter {
151 fn count_message_tokens(&self, msg: &Message) -> usize {
152 msg.content.split_whitespace().count().max(1)
153 }
154 }
155
156 fn deps() -> SummarizationDeps {
157 let summary_json = serde_json::json!({
158 "session_intent": "implement feature X",
159 "files_modified": ["src/lib.rs"],
160 "decisions_made": ["used approach A"],
161 "open_questions": [],
162 "next_steps": ["write tests"],
163 })
164 .to_string();
165 SummarizationDeps {
166 provider: AnyProvider::Mock(MockProvider::with_responses(vec![summary_json])),
167 llm_timeout: Duration::from_secs(5),
168 token_counter: std::sync::Arc::new(WordCountTokenCounter),
169 structured_summaries: true,
170 on_anchored_summary: None,
171 }
172 }
173
174 fn envelope(seq: u64, kind: SessionEvent) -> SessionEventEnvelope {
175 SessionEventEnvelope::new(seq, None, None, kind)
176 }
177
178 fn user_msg(seq: u64, text: &str) -> SessionEventEnvelope {
179 envelope(
180 seq,
181 SessionEvent::UserMessage {
182 text: text.to_owned(),
183 image_refs: vec![],
184 },
185 )
186 }
187
188 fn assistant_msg(seq: u64, text: &str) -> SessionEventEnvelope {
189 envelope(
190 seq,
191 SessionEvent::AssistantMessage {
192 parts: vec![MessagePart::Text {
193 text: text.to_owned(),
194 }],
195 },
196 )
197 }
198
199 #[tokio::test]
200 async fn should_condense_respects_threshold_and_keep_recent() {
201 let condenser = LlmCondenser::new(deps(), 0.8, 4);
202 let mut state = ReconstructedState::default();
203 assert!(
204 !condenser.should_condense(&state, 0.9).await,
205 "too few messages"
206 );
207
208 state.messages = vec![
209 Message::from_legacy(Role::User, "a"),
210 Message::from_legacy(Role::User, "b"),
211 Message::from_legacy(Role::User, "c"),
212 Message::from_legacy(Role::User, "d"),
213 Message::from_legacy(Role::User, "e"),
214 ];
215 assert!(
216 !condenser.should_condense(&state, 0.5).await,
217 "below threshold"
218 );
219 assert!(
220 condenser.should_condense(&state, 0.8).await,
221 "at threshold, enough messages"
222 );
223 }
224
225 #[tokio::test]
226 async fn condense_rejects_when_not_enough_events() {
227 let condenser = LlmCondenser::new(deps(), 0.8, 4);
228 let events = vec![user_msg(0, "a"), assistant_msg(1, "b")];
229 let err = condenser.condense(&events, 0).await.unwrap_err();
230 assert!(matches!(err, SessionError::CondensationOverlap(_)));
231 }
232
233 #[tokio::test]
234 async fn condense_computes_replaced_range_keeping_recent_tail() {
235 let condenser = LlmCondenser::new(deps(), 0.8, 1);
236 let events = vec![
240 user_msg(1, "first question"),
241 assistant_msg(2, "first answer"),
242 user_msg(3, "second question"),
243 ];
244 let result = condenser.condense(&events, 0).await.unwrap();
246 assert_eq!(result.replaced_range, (1, 2));
247 assert!(result.tokens_before > 0);
248 }
249
250 #[tokio::test]
257 async fn condense_twice_on_growing_log_advances_past_prior_range() {
258 let first_condenser = LlmCondenser::new(deps(), 0.8, 1);
259 let mut events = vec![
260 user_msg(1, "q1"),
261 assistant_msg(2, "a1"),
262 user_msg(3, "q2"),
263 assistant_msg(4, "a2"),
264 user_msg(5, "q3"),
265 assistant_msg(6, "a3"),
266 user_msg(7, "q4"),
267 ];
268 let first_result = first_condenser.condense(&events, 0).await.unwrap();
269 assert_eq!(first_result.replaced_range, (1, 6));
270
271 events.push(assistant_msg(8, "a4"));
274 events.push(user_msg(9, "q5"));
275 events.push(assistant_msg(10, "a5"));
276
277 let second_condenser = LlmCondenser::new(deps(), 0.8, 1);
278 let second_result = second_condenser
279 .condense(&events, first_result.replaced_range.1)
280 .await
281 .expect("second condensation on a growing log must not re-attempt the first's range");
282 assert_eq!(
283 second_result.replaced_range,
284 (7, 9),
285 "second condensation must start strictly after the first's replaced_range.1"
286 );
287 }
288}