kiss_coding/
compaction.rs

1//! Compaction: summarize older context when the window fills, keeping a
2//! recent tail. Also branch summarization support for /tree navigation.
3
4use kiss_agent::{AgentMessage, convert_to_llm};
5use kiss_ai::{ContentBlock, Message, Model, StreamOptions, ThinkingLevel, UserContent};
6use serde_json::json;
7
8const TOOL_RESULT_SERIALIZE_CAP: usize = 2000;
9
10/// chars/4 heuristic, matching pi's estimator for messages without usage.
11pub fn estimate_tokens_text(text: &str) -> u64 {
12    (text.len() as u64).div_ceil(4)
13}
14
15pub fn estimate_message_tokens(message: &AgentMessage) -> u64 {
16    let text = match message {
17        AgentMessage::User(u) => u.content.as_text(),
18        AgentMessage::Assistant(a) => a
19            .content
20            .iter()
21            .map(|c| match c {
22                ContentBlock::Text { text, .. } => text.len(),
23                ContentBlock::Thinking { thinking, .. } => thinking.len(),
24                ContentBlock::ToolCall(tc) => tc.arguments.to_string().len() + tc.name.len(),
25                ContentBlock::Image { .. } => 1600,
26            })
27            .sum::<usize>()
28            .to_string(),
29        AgentMessage::ToolResult(t) => t
30            .content
31            .iter()
32            .map(|c| match c {
33                ContentBlock::Text { text, .. } => text.len(),
34                _ => 1600,
35            })
36            .sum::<usize>()
37            .to_string(),
38        AgentMessage::BashExecution(b) => (b.command.len() + b.output.len()).to_string(),
39        AgentMessage::Custom(c) => c.content.as_text(),
40        AgentMessage::BranchSummary(b) => b.summary.clone(),
41        AgentMessage::CompactionSummary(c) => c.summary.clone(),
42    };
43    match message {
44        AgentMessage::Assistant(_) | AgentMessage::ToolResult(_) => text
45            .parse::<u64>()
46            .map(|chars| chars.div_ceil(4))
47            .unwrap_or_else(|_| estimate_tokens_text(&text)),
48        _ => estimate_tokens_text(&text),
49    }
50}
51
52/// Estimated context size: prefer the last assistant usage (input+output+
53/// cache) and add estimates for everything after it.
54pub fn estimate_context_tokens(messages: &[AgentMessage]) -> u64 {
55    let last_assistant = messages
56        .iter()
57        .rposition(|m| matches!(m, AgentMessage::Assistant(_)));
58    match last_assistant {
59        Some(pos) => {
60            let AgentMessage::Assistant(a) = &messages[pos] else {
61                unreachable!()
62            };
63            let base = a.usage.input + a.usage.output + a.usage.cache_read + a.usage.cache_write;
64            let tail: u64 = messages[pos + 1..]
65                .iter()
66                .map(estimate_message_tokens)
67                .sum();
68            base + tail
69        }
70        None => messages.iter().map(estimate_message_tokens).sum(),
71    }
72}
73
74/// Serialize a conversation to labeled text for the summary prompt. Tool
75/// results are capped so summarization stays cheap.
76pub fn serialize_conversation(messages: &[Message]) -> String {
77    let mut out = String::new();
78    for message in messages {
79        match message {
80            Message::User(u) => {
81                out.push_str(&format!("[User]: {}\n", u.content.as_text()));
82            }
83            Message::Assistant(a) => {
84                let mut tool_calls: Vec<String> = Vec::new();
85                for block in &a.content {
86                    match block {
87                        ContentBlock::Thinking { thinking, .. } => {
88                            if !thinking.is_empty() {
89                                out.push_str(&format!("[Assistant thinking]: {thinking}\n"));
90                            }
91                        }
92                        ContentBlock::Text { text, .. } => {
93                            if !text.is_empty() {
94                                out.push_str(&format!("[Assistant]: {text}\n"));
95                            }
96                        }
97                        ContentBlock::ToolCall(tc) => {
98                            let args = tc
99                                .arguments
100                                .as_object()
101                                .map(|o| {
102                                    o.iter()
103                                        .map(|(k, v)| format!("{k}={v}"))
104                                        .collect::<Vec<_>>()
105                                        .join(", ")
106                                })
107                                .unwrap_or_else(|| tc.arguments.to_string());
108                            tool_calls.push(format!("{}({args})", tc.name));
109                        }
110                        ContentBlock::Image { .. } => {}
111                    }
112                }
113                if !tool_calls.is_empty() {
114                    out.push_str(&format!(
115                        "[Assistant tool calls]: {}\n",
116                        tool_calls.join("; ")
117                    ));
118                }
119            }
120            Message::ToolResult(t) => {
121                let text: String = t
122                    .content
123                    .iter()
124                    .filter_map(|c| match c {
125                        ContentBlock::Text { text, .. } => Some(text.as_str()),
126                        _ => None,
127                    })
128                    .collect::<Vec<_>>()
129                    .join("\n");
130                let capped = if text.chars().count() > TOOL_RESULT_SERIALIZE_CAP {
131                    let kept: String = text.chars().take(TOOL_RESULT_SERIALIZE_CAP).collect();
132                    let dropped = text.chars().count() - TOOL_RESULT_SERIALIZE_CAP;
133                    format!("{kept}\n[... {dropped} characters truncated]")
134                } else {
135                    text
136                };
137                out.push_str(&format!("[Tool result]: {capped}\n"));
138            }
139        }
140    }
141    out
142}
143
144/// Extract read/modified file paths from tool calls in the messages.
145pub fn extract_file_ops(messages: &[AgentMessage]) -> (Vec<String>, Vec<String>) {
146    let mut read: Vec<String> = Vec::new();
147    let mut modified: Vec<String> = Vec::new();
148    for message in messages {
149        let AgentMessage::Assistant(a) = message else {
150            continue;
151        };
152        for tc in a.tool_calls() {
153            let Some(path) = tc.arguments["path"].as_str() else {
154                continue;
155            };
156            match tc.name.as_str() {
157                "read" => read.push(path.to_string()),
158                "edit" | "write" => modified.push(path.to_string()),
159                _ => {}
160            }
161        }
162    }
163    read.sort();
164    read.dedup();
165    modified.sort();
166    modified.dedup();
167    (read, modified)
168}
169
170pub const SUMMARY_FORMAT: &str = "## Goal\n[What the user is trying to accomplish]\n\n## Constraints & Preferences\n- [Requirements mentioned by user]\n\n## Progress\n### Done\n- [x] [Completed tasks]\n\n### In Progress\n- [ ] [Current work]\n\n### Blocked\n- [Issues, if any]\n\n## Key Decisions\n- **[Decision]**: [Rationale]\n\n## Next Steps\n1. [What should happen next]\n\n## Critical Context\n- [Data needed to continue]";
171
172/// Cut-point decision for a compaction pass.
173#[derive(Debug, Clone, PartialEq)]
174pub struct CompactionPlan {
175    /// Messages to summarize (complete turns before the cut).
176    pub to_summarize: Vec<AgentMessage>,
177    /// Messages kept verbatim after the cut (the retained tail).
178    pub kept: Vec<AgentMessage>,
179    /// Split-turn prefix that must be summarized separately.
180    pub turn_prefix: Vec<AgentMessage>,
181    pub is_split_turn: bool,
182    pub tokens_before: u64,
183}
184
185fn is_valid_cut(message: &AgentMessage) -> bool {
186    // Never cut at tool results: they must stay with their tool call.
187    matches!(
188        message,
189        AgentMessage::User(_)
190            | AgentMessage::Assistant(_)
191            | AgentMessage::BashExecution(_)
192            | AgentMessage::Custom(_)
193            | AgentMessage::BranchSummary(_)
194    )
195}
196
197fn is_turn_start(message: &AgentMessage) -> bool {
198    matches!(
199        message,
200        AgentMessage::User(_) | AgentMessage::Custom(_) | AgentMessage::BashExecution(_)
201    )
202}
203
204/// Choose the cut point: walk back from the newest message accumulating
205/// estimates until `keep_recent_tokens`, then snap to a turn boundary.
206pub fn plan_compaction(messages: &[AgentMessage], keep_recent_tokens: u64) -> CompactionPlan {
207    let tokens_before = estimate_context_tokens(messages);
208    let mut budget = 0u64;
209    let mut cut = messages.len();
210    for (i, message) in messages.iter().enumerate().rev() {
211        budget += estimate_message_tokens(message);
212        if budget > keep_recent_tokens {
213            break;
214        }
215        if is_valid_cut(message) {
216            cut = i;
217        }
218    }
219    if cut == messages.len() && !messages.is_empty() {
220        // Keep at least the final message when everything is over budget.
221        cut = messages.len() - 1;
222        while cut > 0 && !is_valid_cut(&messages[cut]) {
223            cut -= 1;
224        }
225    }
226
227    // Snap to a turn start when possible. Otherwise this is a split turn.
228    let turn_start = messages[..cut].iter().rposition(is_turn_start);
229    match turn_start {
230        Some(_) if is_turn_start(&messages[cut]) => CompactionPlan {
231            to_summarize: messages[..cut].to_vec(),
232            kept: messages[cut..].to_vec(),
233            turn_prefix: Vec::new(),
234            is_split_turn: false,
235            tokens_before,
236        },
237        Some(start) => CompactionPlan {
238            to_summarize: messages[..start].to_vec(),
239            kept: messages[cut..].to_vec(),
240            turn_prefix: messages[start..cut].to_vec(),
241            is_split_turn: true,
242            tokens_before,
243        },
244        None => CompactionPlan {
245            to_summarize: Vec::new(),
246            kept: messages[cut..].to_vec(),
247            turn_prefix: messages[..cut].to_vec(),
248            is_split_turn: cut > 0,
249            tokens_before,
250        },
251    }
252}
253
254/// Whether auto-compaction should trigger.
255pub fn should_compact(context_tokens: u64, context_window: u64, reserve_tokens: u64) -> bool {
256    context_tokens > context_window.saturating_sub(reserve_tokens)
257}
258
259pub struct SummaryOutcome {
260    pub summary: String,
261    pub usage: Option<kiss_ai::Usage>,
262}
263
264fn summary_failure(message: &kiss_ai::AssistantMessage) -> Option<String> {
265    match message.stop_reason {
266        kiss_ai::StopReason::Error => Some(format!(
267            "summary generation failed: {}",
268            message.error_message.as_deref().unwrap_or("unknown error")
269        )),
270        kiss_ai::StopReason::Length => {
271            Some("summary generation failed: response reached the token limit".into())
272        }
273        _ if message.tool_calls().next().is_some() => {
274            Some("summary generation failed: response contained a tool call".into())
275        }
276        _ => None,
277    }
278}
279
280fn summary_prompt(
281    conversation_text: &str,
282    previous_summary: Option<&str>,
283    custom_instructions: Option<&str>,
284    continuation: bool,
285) -> String {
286    let mut prompt = format!("# Conversation\n{conversation_text}\n\n# Instructions\n");
287    if continuation {
288        prompt.push_str("The conversation above is earlier context from an ongoing task. Later messages are retained separately. Create a concise checkpoint that can be placed before them so work can continue; do not infer or recreate later messages. Use exactly this structure:\n\n");
289    } else {
290        prompt.push_str("Summarize the conversation so a coding agent can pick up where it left off. Use exactly this structure:\n\n");
291    }
292    prompt.push_str(SUMMARY_FORMAT);
293    prompt.push_str("\n\nAlso include, at the end, a <read-files> block listing files that were read and a <modified-files> block listing files that were changed, one path per line, when known.");
294    if let Some(prev) = previous_summary {
295        prompt.push_str("\n\nA previous summary of earlier context exists. Fold it in:\n\n");
296        prompt.push_str(prev);
297    }
298    if let Some(custom) = custom_instructions {
299        prompt.push_str("\n\nAdditional focus requested by the user: ");
300        prompt.push_str(custom);
301    }
302    prompt
303}
304
305/// Generate a structured summary with the LLM. One-off prompt: fresh session
306/// id and no cache writes wanted, so it's a plain request.
307pub async fn generate_summary(
308    model: &Model,
309    credential: Option<kiss_ai::ResolvedCredential>,
310    conversation_text: &str,
311    previous_summary: Option<&str>,
312    custom_instructions: Option<&str>,
313    continuation: bool,
314    cancel: tokio_util::sync::CancellationToken,
315) -> anyhow::Result<SummaryOutcome> {
316    let prompt = summary_prompt(
317        conversation_text,
318        previous_summary,
319        custom_instructions,
320        continuation,
321    );
322    let context = kiss_ai::Context {
323        system_prompt: None,
324        openai_responses_input: None,
325        messages: vec![Message::User(kiss_ai::UserMessage {
326            content: UserContent::Text(prompt),
327            timestamp: kiss_ai::now_ms(),
328        })],
329        tools: vec![],
330    };
331    let options = StreamOptions {
332        credential,
333        reasoning: ThinkingLevel::Off,
334        cancel,
335        ..Default::default()
336    };
337    let message = kiss_ai::stream_simple(model, &context, &options)
338        .result()
339        .await;
340    if let Some(error) = summary_failure(&message) {
341        anyhow::bail!(error);
342    }
343    Ok(SummaryOutcome {
344        summary: message.text(),
345        usage: Some(message.usage),
346    })
347}
348
349/// Details payload stored on compaction/branch-summary entries.
350pub fn file_ops_details(read: &[String], modified: &[String]) -> serde_json::Value {
351    json!({"readFiles": read, "modifiedFiles": modified})
352}
353
354/// Serialize agent messages for summarization (convert then serialize).
355pub fn serialize_agent_messages(messages: &[AgentMessage]) -> String {
356    serialize_conversation(&convert_to_llm(messages))
357}
358
359#[cfg(test)]
360mod tests {
361    use super::*;
362    use kiss_ai::{AssistantMessage, StopReason, ToolCall, ToolResultMessage};
363
364    fn user(text: &str) -> AgentMessage {
365        AgentMessage::user(text)
366    }
367
368    fn assistant(text: &str) -> AgentMessage {
369        let mut m = AssistantMessage::empty("fake", "fake", "fake");
370        m.content.push(ContentBlock::text(text));
371        m.stop_reason = StopReason::Stop;
372        AgentMessage::Assistant(m)
373    }
374
375    fn tool_result(text: &str) -> AgentMessage {
376        AgentMessage::ToolResult(ToolResultMessage {
377            tool_call_id: "c".into(),
378            tool_name: "bash".into(),
379            content: vec![ContentBlock::text(text)],
380            details: None,
381            usage: None,
382            is_error: false,
383            timestamp: 0,
384        })
385    }
386
387    #[test]
388    fn cut_at_turn_boundary_never_tool_result() {
389        // Two turns. Small keep budget keeps only the last turn.
390        let messages = vec![
391            user(&"a".repeat(400)),
392            assistant(&"b".repeat(400)),
393            tool_result(&"c".repeat(400)),
394            user(&"d".repeat(400)),
395            assistant(&"e".repeat(400)),
396        ];
397        let plan = plan_compaction(&messages, 250);
398        assert!(!plan.is_split_turn);
399        assert!(matches!(plan.kept.first().unwrap(), AgentMessage::User(_)));
400        assert_eq!(plan.to_summarize.len(), 3);
401    }
402
403    #[test]
404    fn split_turn_detected() {
405        // One huge turn: user + many assistant/tool pairs.
406        let messages = vec![
407            user("start"),
408            assistant(&"x".repeat(2000)),
409            tool_result(&"y".repeat(2000)),
410            assistant(&"z".repeat(2000)),
411        ];
412        let plan = plan_compaction(&messages, 600);
413        assert!(plan.is_split_turn);
414        assert!(plan.to_summarize.is_empty());
415        assert!(!plan.turn_prefix.is_empty());
416        // Kept must not start with a tool result.
417        assert!(is_valid_cut(plan.kept.first().unwrap()));
418    }
419
420    #[test]
421    fn oversized_trailing_tool_result_keeps_its_assistant_call() {
422        let messages = vec![
423            user("start"),
424            assistant("tool call"),
425            tool_result(&"x".repeat(8_000)),
426        ];
427        let plan = plan_compaction(&messages, 10);
428        assert!(matches!(
429            plan.kept.first(),
430            Some(AgentMessage::Assistant(_))
431        ));
432        assert!(matches!(
433            plan.kept.get(1),
434            Some(AgentMessage::ToolResult(_))
435        ));
436    }
437
438    #[test]
439    fn serialization_labels_and_caps() {
440        let big = "L".repeat(3000);
441        let messages = vec![user("do it"), assistant("on it"), tool_result(&big)];
442        let text = serialize_agent_messages(&messages);
443        assert!(text.contains("[User]: do it"));
444        assert!(text.contains("[Assistant]: on it"));
445        assert!(text.contains("characters truncated"));
446    }
447
448    #[test]
449    fn file_ops_extraction() {
450        let mut a = AssistantMessage::empty("f", "f", "f");
451        a.content.push(ContentBlock::ToolCall(ToolCall {
452            id: "1".into(),
453            name: "read".into(),
454            arguments: json!({"path": "src/a.rs"}),
455            thought_signature: None,
456        }));
457        a.content.push(ContentBlock::ToolCall(ToolCall {
458            id: "2".into(),
459            name: "edit".into(),
460            arguments: json!({"path": "src/b.rs"}),
461            thought_signature: None,
462        }));
463        let (read, modified) = extract_file_ops(&[AgentMessage::Assistant(a)]);
464        assert_eq!(read, vec!["src/a.rs"]);
465        assert_eq!(modified, vec!["src/b.rs"]);
466    }
467
468    #[test]
469    fn threshold() {
470        assert!(should_compact(190_000, 200_000, 16_384));
471        assert!(!should_compact(100_000, 200_000, 16_384));
472    }
473
474    #[test]
475    fn split_turn_prompt_separates_conversation_from_continuation_instructions() {
476        let prompt = summary_prompt("user work", None, None, true);
477        assert!(prompt.starts_with("# Conversation\nuser work\n\n# Instructions\n"));
478        assert!(prompt.contains("Later messages are retained separately"));
479        assert!(prompt.contains("do not infer or recreate later messages"));
480    }
481
482    #[test]
483    fn incomplete_summary_is_rejected() {
484        let mut message = AssistantMessage::empty("fake", "fake", "fake");
485        message.stop_reason = StopReason::Length;
486        assert_eq!(
487            summary_failure(&message).as_deref(),
488            Some("summary generation failed: response reached the token limit")
489        );
490    }
491
492    #[test]
493    fn summary_with_tool_call_is_rejected() {
494        let mut message = AssistantMessage::empty("fake", "fake", "fake");
495        message.stop_reason = StopReason::ToolUse;
496        message.content.push(ContentBlock::ToolCall(ToolCall {
497            id: "call-1".into(),
498            name: "read".into(),
499            arguments: json!({"path": "src/lib.rs"}),
500            thought_signature: None,
501        }));
502        assert_eq!(
503            summary_failure(&message).as_deref(),
504            Some("summary generation failed: response contained a tool call")
505        );
506    }
507}
kiss_coding/compaction.rs

kiss_coding/
compaction.rs
