Skip to main content

pi/
compaction.rs

1//! Context compaction for long sessions.
2//!
3//! This module ports the pi-mono compaction algorithm:
4//! - Estimate context usage and choose a cut point that keeps recent context
5//! - Summarize the discarded portion with the LLM (iteratively updating prior summaries)
6//! - Record a `compaction` session entry containing the summary and cut point
7//! - When building provider context, the session inserts the summary before the kept region
8//!   and omits older messages.
9
10use crate::error::{Error, Result};
11use crate::model::{
12    AssistantMessage, ContentBlock, Message, StopReason, TextContent, ThinkingLevel, ToolCall,
13    Usage, UserContent, UserMessage,
14};
15use crate::provider::{Context, Provider, StreamOptions};
16use crate::session::{SessionEntry, SessionMessage, session_message_to_model};
17use futures::StreamExt;
18use serde::Serialize;
19use serde_json::{Map, Value};
20use std::collections::{HashMap, HashSet};
21use std::fmt::Write as _;
22use std::sync::Arc;
23
24/// Approximate characters per token for English text with GPT-family tokenizers.
25/// Intentionally conservative (overestimates tokens) to avoid exceeding context windows.
26/// Set to 3 to safely account for code/symbol-heavy content which is denser than prose.
27const CHARS_PER_TOKEN_ESTIMATE: usize = 3;
28
29/// Estimated tokens for an image content block (~1200 tokens).
30const IMAGE_TOKEN_ESTIMATE: usize = 1200;
31
32/// Character-equivalent estimate for an image (IMAGE_TOKEN_ESTIMATE * CHARS_PER_TOKEN_ESTIMATE).
33const IMAGE_CHAR_ESTIMATE: usize = IMAGE_TOKEN_ESTIMATE * CHARS_PER_TOKEN_ESTIMATE;
34
35/// Count the serialized JSON byte length of a [`Value`] without allocating a `String`.
36///
37/// Uses `serde_json::to_writer` with a sink that only counts bytes – this gives the
38/// exact same length as `serde_json::to_string(&v).len()` at zero heap cost.
39fn json_byte_len(value: &Value) -> usize {
40    struct Counter(usize);
41    impl std::io::Write for Counter {
42        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
43            self.0 = self.0.saturating_add(buf.len());
44            Ok(buf.len())
45        }
46        fn flush(&mut self) -> std::io::Result<()> {
47            Ok(())
48        }
49    }
50    let mut c = Counter(0);
51    if serde_json::to_writer(&mut c, value).is_err() {
52        // Fallback or partial count on error (e.g. recursion limit)
53    }
54    c.0
55}
56
57// =============================================================================
58// Public types
59// =============================================================================
60
61#[derive(Debug, Clone)]
62pub struct ResolvedCompactionSettings {
63    pub enabled: bool,
64    pub context_window_tokens: u32,
65    pub reserve_tokens: u32,
66    pub keep_recent_tokens: u32,
67}
68
69impl Default for ResolvedCompactionSettings {
70    /// Conservative default using the smallest common context window (32K).
71    ///
72    /// Production code paths should always override `context_window_tokens`
73    /// with the actual model's context window via
74    /// [`context_window_tokens_for_entry`](crate::main) or equivalent.
75    /// This default is deliberately conservative so that if a code path
76    /// forgets to override, compaction triggers too early (safe) rather
77    /// than too late (could exceed the real context window).
78    fn default() -> Self {
79        let context_window_tokens: u32 = 128_000;
80        Self {
81            enabled: true,
82            context_window_tokens,
83            // ~8% of context window
84            reserve_tokens: 10_240,
85            // 10% of context window
86            keep_recent_tokens: 12_800,
87        }
88    }
89}
90
91/// Details stored in `CompactionEntry.details` for cumulative file tracking.
92#[derive(Debug, Clone, Serialize)]
93#[serde(rename_all = "camelCase")]
94pub struct CompactionDetails {
95    pub read_files: Vec<String>,
96    pub modified_files: Vec<String>,
97}
98
99#[derive(Debug, Clone, Serialize)]
100#[serde(rename_all = "camelCase")]
101pub struct CompactionResult {
102    pub summary: String,
103    pub first_kept_entry_id: String,
104    pub tokens_before: u64,
105    pub details: CompactionDetails,
106}
107
108#[derive(Debug, Clone)]
109pub struct CompactionPreparation {
110    pub first_kept_entry_id: String,
111    pub messages_to_summarize: Vec<SessionMessage>,
112    pub turn_prefix_messages: Vec<SessionMessage>,
113    pub is_split_turn: bool,
114    pub tokens_before: u64,
115    pub previous_summary: Option<String>,
116    pub file_ops: FileOperations,
117    pub settings: ResolvedCompactionSettings,
118}
119
120pub fn compaction_preparation_to_value(prep: &CompactionPreparation) -> Value {
121    let messages_to_summarize =
122        serde_json::to_value(&prep.messages_to_summarize).unwrap_or(Value::Array(Vec::new()));
123    let turn_prefix_messages =
124        serde_json::to_value(&prep.turn_prefix_messages).unwrap_or(Value::Array(Vec::new()));
125
126    let mut obj = Map::new();
127    obj.insert(
128        "firstKeptEntryId".to_string(),
129        Value::String(prep.first_kept_entry_id.clone()),
130    );
131    obj.insert("messagesToSummarize".to_string(), messages_to_summarize);
132    obj.insert("turnPrefixMessages".to_string(), turn_prefix_messages);
133    obj.insert("isSplitTurn".to_string(), Value::Bool(prep.is_split_turn));
134    obj.insert("tokensBefore".to_string(), Value::from(prep.tokens_before));
135    if let Some(previous_summary) = &prep.previous_summary {
136        obj.insert(
137            "previousSummary".to_string(),
138            Value::String(previous_summary.clone()),
139        );
140    }
141    obj.insert("fileOps".to_string(), file_ops_to_value(&prep.file_ops));
142    obj.insert(
143        "settings".to_string(),
144        compaction_settings_to_value(&prep.settings),
145    );
146    Value::Object(obj)
147}
148
149fn file_ops_to_value(file_ops: &FileOperations) -> Value {
150    let read = sorted_file_ops(&file_ops.read);
151    let written = sorted_file_ops(&file_ops.written);
152    let edited = sorted_file_ops(&file_ops.edited);
153    let mut obj = Map::new();
154    obj.insert("read".to_string(), Value::Array(read));
155    obj.insert("written".to_string(), Value::Array(written));
156    obj.insert("edited".to_string(), Value::Array(edited));
157    Value::Object(obj)
158}
159
160fn sorted_file_ops(values: &HashSet<String>) -> Vec<Value> {
161    let mut entries = values.iter().cloned().collect::<Vec<_>>();
162    entries.sort();
163    entries.into_iter().map(Value::String).collect()
164}
165
166fn compaction_settings_to_value(settings: &ResolvedCompactionSettings) -> Value {
167    let mut obj = Map::new();
168    obj.insert("enabled".to_string(), Value::Bool(settings.enabled));
169    obj.insert(
170        "contextWindowTokens".to_string(),
171        Value::from(settings.context_window_tokens),
172    );
173    obj.insert(
174        "reserveTokens".to_string(),
175        Value::from(settings.reserve_tokens),
176    );
177    obj.insert(
178        "keepRecentTokens".to_string(),
179        Value::from(settings.keep_recent_tokens),
180    );
181    Value::Object(obj)
182}
183
184// =============================================================================
185// File op tracking (read/write/edit)
186// =============================================================================
187
188#[derive(Debug, Clone, Default)]
189pub struct FileOperations {
190    read: HashSet<String>,
191    written: HashSet<String>,
192    edited: HashSet<String>,
193}
194
195impl FileOperations {
196    pub fn read_files(&self) -> impl Iterator<Item = &str> {
197        self.read.iter().map(String::as_str)
198    }
199}
200
201fn build_tool_status_map(messages: &[SessionMessage]) -> HashMap<String, bool> {
202    let mut status = HashMap::new();
203    for msg in messages {
204        if let SessionMessage::ToolResult {
205            tool_call_id,
206            is_error,
207            ..
208        } = msg
209        {
210            status.insert(tool_call_id.clone(), !*is_error);
211        }
212    }
213    status
214}
215
216fn extract_file_ops_from_message(
217    message: &SessionMessage,
218    file_ops: &mut FileOperations,
219    tool_status: &HashMap<String, bool>,
220) {
221    let SessionMessage::Assistant { message } = message else {
222        return;
223    };
224
225    for block in &message.content {
226        let ContentBlock::ToolCall(ToolCall {
227            id,
228            name,
229            arguments,
230            ..
231        }) = block
232        else {
233            continue;
234        };
235
236        // Only track successful tool calls.
237        if !tool_status.get(id).copied().unwrap_or(false) {
238            continue;
239        }
240
241        let Some(path) = arguments.get("path").and_then(Value::as_str) else {
242            continue;
243        };
244
245        match name.as_str() {
246            "read" | "grep" | "find" | "ls" => {
247                file_ops.read.insert(path.to_string());
248            }
249            "write" => {
250                file_ops.written.insert(path.to_string());
251            }
252            "edit" | "hashline_edit" => {
253                file_ops.edited.insert(path.to_string());
254            }
255            _ => {}
256        }
257    }
258}
259
260fn compute_file_lists(file_ops: &FileOperations) -> (Vec<String>, Vec<String>) {
261    let modified: HashSet<&String> = file_ops
262        .edited
263        .iter()
264        .chain(file_ops.written.iter())
265        .collect();
266
267    let mut read_only = file_ops
268        .read
269        .iter()
270        .filter(|f| !modified.contains(f))
271        .cloned()
272        .collect::<Vec<_>>();
273    read_only.sort();
274
275    let mut modified_files = modified.into_iter().cloned().collect::<Vec<_>>();
276    modified_files.sort();
277
278    (read_only, modified_files)
279}
280
281fn write_escaped_file_list(out: &mut String, tag: &str, files: &[String]) {
282    out.push('<');
283    out.push_str(tag);
284    out.push_str(">\n");
285    for (i, file) in files.iter().enumerate() {
286        if i > 0 {
287            out.push('\n');
288        }
289        // Inline escape: replace < and > in file paths
290        for ch in file.chars() {
291            match ch {
292                '<' => out.push_str("&lt;"),
293                '>' => out.push_str("&gt;"),
294                _ => out.push(ch),
295            }
296        }
297    }
298    out.push_str("\n</");
299    out.push_str(tag);
300    out.push('>');
301}
302
303fn format_file_operations(read_files: &[String], modified_files: &[String]) -> String {
304    if read_files.is_empty() && modified_files.is_empty() {
305        return String::new();
306    }
307
308    let mut out = String::from("\n\n");
309    if !read_files.is_empty() {
310        write_escaped_file_list(&mut out, "read-files", read_files);
311    }
312    if !modified_files.is_empty() {
313        if !read_files.is_empty() {
314            out.push_str("\n\n");
315        }
316        write_escaped_file_list(&mut out, "modified-files", modified_files);
317    }
318    out
319}
320
321// =============================================================================
322// Token estimation
323// =============================================================================
324
325const fn calculate_context_tokens(usage: &Usage) -> u64 {
326    if usage.total_tokens > 0 {
327        usage.total_tokens
328    } else {
329        usage.input.saturating_add(usage.output)
330    }
331}
332
333const fn get_assistant_usage(message: &SessionMessage) -> Option<&Usage> {
334    let SessionMessage::Assistant { message } = message else {
335        return None;
336    };
337
338    if matches!(message.stop_reason, StopReason::Aborted | StopReason::Error) {
339        return None;
340    }
341
342    Some(&message.usage)
343}
344
345#[derive(Debug, Clone, Copy)]
346struct ContextUsageEstimate {
347    tokens: u64,
348    last_usage_index: Option<usize>,
349}
350
351fn estimate_context_tokens(messages: &[SessionMessage]) -> ContextUsageEstimate {
352    let mut last_usage: Option<(&Usage, usize)> = None;
353    for (idx, msg) in messages.iter().enumerate().rev() {
354        if let Some(usage) = get_assistant_usage(msg) {
355            last_usage = Some((usage, idx));
356            break;
357        }
358    }
359
360    let Some((usage, usage_index)) = last_usage else {
361        let total = messages
362            .iter()
363            .map(estimate_tokens)
364            .fold(0u64, u64::saturating_add);
365        return ContextUsageEstimate {
366            tokens: total,
367            last_usage_index: None,
368        };
369    };
370
371    let usage_tokens = calculate_context_tokens(usage);
372
373    // Fall back to heuristic estimation if the provider didn't return usage metrics
374    if usage_tokens == 0 {
375        let total = messages
376            .iter()
377            .map(estimate_tokens)
378            .fold(0u64, u64::saturating_add);
379        return ContextUsageEstimate {
380            tokens: total,
381            last_usage_index: None,
382        };
383    }
384
385    let trailing_tokens = messages[usage_index + 1..]
386        .iter()
387        .map(estimate_tokens)
388        .fold(0u64, u64::saturating_add);
389    ContextUsageEstimate {
390        tokens: usage_tokens.saturating_add(trailing_tokens),
391        last_usage_index: Some(usage_index),
392    }
393}
394
395fn should_compact(
396    context_tokens: u64,
397    context_window: u32,
398    settings: &ResolvedCompactionSettings,
399) -> bool {
400    if !settings.enabled {
401        return false;
402    }
403    let reserve = u64::from(settings.reserve_tokens);
404    let window = u64::from(context_window);
405    context_tokens >= window.saturating_sub(reserve)
406}
407
408fn estimate_tokens(message: &SessionMessage) -> u64 {
409    let mut chars: usize = 0;
410
411    match message {
412        SessionMessage::User { content, .. } => match content {
413            UserContent::Text(text) => chars = text.len(),
414            UserContent::Blocks(blocks) => {
415                for block in blocks {
416                    match block {
417                        ContentBlock::Text(text) => {
418                            chars = chars.saturating_add(text.text.len());
419                        }
420                        ContentBlock::Image(_) => {
421                            chars = chars.saturating_add(IMAGE_CHAR_ESTIMATE);
422                        }
423                        ContentBlock::Thinking(thinking) => {
424                            chars = chars.saturating_add(thinking.thinking.len());
425                        }
426                        ContentBlock::ToolCall(call) => {
427                            chars = chars.saturating_add(call.name.len());
428                            chars = chars.saturating_add(json_byte_len(&call.arguments));
429                        }
430                    }
431                }
432            }
433        },
434        SessionMessage::Assistant { message } => {
435            for block in &message.content {
436                match block {
437                    ContentBlock::Text(text) => {
438                        chars = chars.saturating_add(text.text.len());
439                    }
440                    ContentBlock::Thinking(thinking) => {
441                        chars = chars.saturating_add(thinking.thinking.len());
442                    }
443                    ContentBlock::Image(_) => {
444                        chars = chars.saturating_add(IMAGE_CHAR_ESTIMATE);
445                    }
446                    ContentBlock::ToolCall(call) => {
447                        chars = chars.saturating_add(call.name.len());
448                        chars = chars.saturating_add(json_byte_len(&call.arguments));
449                    }
450                }
451            }
452        }
453        SessionMessage::ToolResult { content, .. } => {
454            for block in content {
455                match block {
456                    ContentBlock::Text(text) => {
457                        chars = chars.saturating_add(text.text.len());
458                    }
459                    ContentBlock::Thinking(thinking) => {
460                        chars = chars.saturating_add(thinking.thinking.len());
461                    }
462                    ContentBlock::Image(_) => {
463                        chars = chars.saturating_add(IMAGE_CHAR_ESTIMATE);
464                    }
465                    ContentBlock::ToolCall(call) => {
466                        chars = chars.saturating_add(call.name.len());
467                        chars = chars.saturating_add(json_byte_len(&call.arguments));
468                    }
469                }
470            }
471        }
472        SessionMessage::Custom { content, .. } => chars = content.len(),
473        SessionMessage::BashExecution {
474            command, output, ..
475        } => chars = command.len().saturating_add(output.len()),
476        SessionMessage::BranchSummary { summary, .. }
477        | SessionMessage::CompactionSummary { summary, .. } => chars = summary.len(),
478    }
479
480    u64::try_from(chars.div_ceil(CHARS_PER_TOKEN_ESTIMATE)).unwrap_or(u64::MAX)
481}
482
483// =============================================================================
484// Cut point detection
485// =============================================================================
486
487#[derive(Debug, Clone, Copy)]
488struct CutPointResult {
489    first_kept_entry_index: usize,
490    turn_start_index: Option<usize>,
491    is_split_turn: bool,
492}
493
494fn message_from_entry(entry: &SessionEntry) -> Option<SessionMessage> {
495    match entry {
496        SessionEntry::Message(msg_entry) => Some(msg_entry.message.clone()),
497        SessionEntry::BranchSummary(summary) => Some(SessionMessage::BranchSummary {
498            summary: summary.summary.clone(),
499            from_id: summary.from_id.clone(),
500        }),
501        SessionEntry::Compaction(compaction) => Some(SessionMessage::CompactionSummary {
502            summary: compaction.summary.clone(),
503            tokens_before: compaction.tokens_before,
504        }),
505        _ => None,
506    }
507}
508
509const fn entry_is_message_like(entry: &SessionEntry) -> bool {
510    matches!(
511        entry,
512        SessionEntry::Message(_) | SessionEntry::BranchSummary(_)
513    )
514}
515
516const fn entry_is_compaction_boundary(entry: &SessionEntry) -> bool {
517    matches!(entry, SessionEntry::Compaction(_))
518}
519
520fn find_valid_cut_points(
521    entries: &[SessionEntry],
522    start_index: usize,
523    end_index: usize,
524) -> Vec<usize> {
525    let mut cut_points = Vec::new();
526    for (idx, entry) in entries.iter().enumerate().take(end_index).skip(start_index) {
527        match entry {
528            SessionEntry::Message(msg_entry) => match msg_entry.message {
529                SessionMessage::ToolResult { .. } => {}
530                _ => cut_points.push(idx),
531            },
532            SessionEntry::BranchSummary(_) => cut_points.push(idx),
533            _ => {}
534        }
535    }
536    cut_points
537}
538
539fn entry_has_tool_calls(entry: &SessionEntry) -> bool {
540    matches!(
541        entry,
542        SessionEntry::Message(msg) if matches!(
543            &msg.message,
544            SessionMessage::Assistant { message } if message.content.iter().any(|b| matches!(b, ContentBlock::ToolCall(_)))
545        )
546    )
547}
548
549const fn is_user_turn_start(entry: &SessionEntry) -> bool {
550    match entry {
551        SessionEntry::BranchSummary(_) => true,
552        SessionEntry::Message(msg_entry) => matches!(
553            msg_entry.message,
554            SessionMessage::User { .. } | SessionMessage::BashExecution { .. }
555        ),
556        _ => false,
557    }
558}
559
560fn find_turn_start_index(
561    entries: &[SessionEntry],
562    entry_index: usize,
563    start_index: usize,
564) -> Option<usize> {
565    (start_index..=entry_index)
566        .rev()
567        .find(|&idx| is_user_turn_start(&entries[idx]))
568}
569
570fn find_cut_point(
571    entries: &[SessionEntry],
572    start_index: usize,
573    end_index: usize,
574    keep_recent_tokens: u32,
575) -> CutPointResult {
576    let cut_points = find_valid_cut_points(entries, start_index, end_index);
577    if cut_points.is_empty() {
578        return CutPointResult {
579            first_kept_entry_index: start_index,
580            turn_start_index: None,
581            is_split_turn: false,
582        };
583    }
584
585    let mut accumulated_tokens: u64 = 0;
586    let mut cut_index = cut_points[0];
587
588    for i in (start_index..end_index).rev() {
589        let entry = &entries[i];
590        if let Some(msg) = message_from_entry(entry) {
591            accumulated_tokens = accumulated_tokens.saturating_add(estimate_tokens(&msg));
592        } else {
593            continue;
594        }
595
596        if accumulated_tokens >= u64::from(keep_recent_tokens) {
597            // Binary search: find the largest cut point <= i.
598            // `partition_point` returns the index of the first element > i,
599            // so idx-1 is the largest element <= i (if any).
600            let pos = cut_points.partition_point(|&cp| cp <= i);
601            if pos > 0 {
602                cut_index = cut_points[pos - 1];
603            }
604            // else: no cut point <= i, keep the fallback (cut_points[0])
605            break;
606        }
607    }
608
609    while cut_index > start_index {
610        let prev = &entries[cut_index - 1];
611        if entry_is_compaction_boundary(prev) {
612            break;
613        }
614        if entry_is_message_like(prev) {
615            break;
616        }
617        cut_index -= 1;
618    }
619
620    let is_user_message = is_user_turn_start(&entries[cut_index]);
621    let turn_start_index = if is_user_message {
622        None
623    } else {
624        find_turn_start_index(entries, cut_index, start_index)
625    };
626
627    CutPointResult {
628        first_kept_entry_index: cut_index,
629        turn_start_index,
630        is_split_turn: !is_user_message && turn_start_index.is_some(),
631    }
632}
633
634// =============================================================================
635// Summarization prompts
636// =============================================================================
637
638const SUMMARIZATION_SYSTEM_PROMPT: &str = "You are a context summarization assistant. Your task is to read a conversation between a user and an AI coding assistant, then produce a structured summary following the exact format specified.\n\nDo NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.";
639
640const SUMMARIZATION_PROMPT: &str = "The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.\n\nUse this EXACT format:\n\n## Goal\n[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]\n\n## Constraints & Preferences\n- [Any constraints, preferences, or requirements mentioned by user]\n- [Or \"(none)\" if none were mentioned]\n\n## Progress\n### Done\n- [x] [Completed tasks/changes]\n\n### In Progress\n- [ ] [Current work]\n\n### Blocked\n- [Issues preventing progress, if any]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale]\n\n## Next Steps\n1. [Ordered list of what should happen next]\n\n## Critical Context\n- [Any data, examples, or references needed to continue]\n- [Or \"(none)\" if not applicable]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";
641
642const UPDATE_SUMMARIZATION_PROMPT: &str = "The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.\n\nUpdate the existing structured summary with new information. RULES:\n- PRESERVE all existing information from the previous summary\n- ADD new progress, decisions, and context from the new messages\n- UPDATE the Progress section: move items from \"In Progress\" to \"Done\" when completed\n- UPDATE \"Next Steps\" based on what was accomplished\n- PRESERVE exact file paths, function names, and error messages\n- If something is no longer relevant, you may remove it\n\nUse this EXACT format:\n\n## Goal\n[Preserve existing goals, add new ones if the task expanded]\n\n## Constraints & Preferences\n- [Preserve existing, add new ones discovered]\n\n## Progress\n### Done\n- [x] [Include previously done items AND newly completed items]\n\n### In Progress\n- [ ] [Current work - update based on progress]\n\n### Blocked\n- [Current blockers - remove if resolved]\n\n## Key Decisions\n- **[Decision]**: [Brief rationale] (preserve all previous, add new)\n\n## Next Steps\n1. [Update based on current state]\n\n## Critical Context\n- [Preserve important context, add new if needed]\n\nKeep each section concise. Preserve exact file paths, function names, and error messages.";
643
644const TURN_PREFIX_SUMMARIZATION_PROMPT: &str = "This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.\n\nSummarize the prefix to provide context for the retained suffix:\n\n## Original Request\n[What did the user ask for in this turn?]\n\n## Early Progress\n- [Key decisions and work done in the prefix]\n\n## Context for Suffix\n- [Information needed to understand the retained recent work]\n\nBe concise. Focus on what's needed to understand the kept suffix.";
645
646fn push_message_separator(out: &mut String) {
647    if !out.is_empty() {
648        out.push_str("\n\n");
649    }
650}
651
652fn user_has_serializable_content(user: &UserMessage) -> bool {
653    match &user.content {
654        UserContent::Text(text) => !text.is_empty(),
655        UserContent::Blocks(blocks) => blocks
656            .iter()
657            .any(|c| matches!(c, ContentBlock::Text(t) if !t.text.is_empty())),
658    }
659}
660
661fn append_user_message(out: &mut String, user: &UserMessage) {
662    if !user_has_serializable_content(user) {
663        return;
664    }
665
666    push_message_separator(out);
667    out.push_str("[User]: ");
668    match &user.content {
669        UserContent::Text(text) => out.push_str(text),
670        UserContent::Blocks(blocks) => {
671            for block in blocks {
672                if let ContentBlock::Text(text) = block {
673                    out.push_str(&text.text);
674                }
675            }
676        }
677    }
678}
679
680fn append_custom_message(out: &mut String, custom_type: &str, content: &str) {
681    if content.trim().is_empty() {
682        return;
683    }
684
685    push_message_separator(out);
686    out.push('[');
687    if custom_type.trim().is_empty() {
688        out.push_str("Custom");
689    } else {
690        out.push_str("Custom:");
691        out.push_str(custom_type);
692    }
693    out.push_str("]: ");
694    out.push_str(content);
695}
696
697fn assistant_content_flags(assistant: &AssistantMessage) -> (bool, bool, bool) {
698    let mut has_thinking = false;
699    let mut has_text = false;
700    let mut has_tools = false;
701    for block in &assistant.content {
702        match block {
703            ContentBlock::Thinking(_) => has_thinking = true,
704            ContentBlock::Text(_) => has_text = true,
705            ContentBlock::ToolCall(_) => has_tools = true,
706            ContentBlock::Image(_) => {}
707        }
708    }
709    (has_thinking, has_text, has_tools)
710}
711
712fn append_assistant_thinking(out: &mut String, assistant: &AssistantMessage) {
713    push_message_separator(out);
714    out.push_str("[Assistant thinking]: ");
715    let mut first = true;
716    for block in &assistant.content {
717        if let ContentBlock::Thinking(thinking) = block {
718            if !first {
719                out.push('\n');
720            }
721            out.push_str(&thinking.thinking);
722            first = false;
723        }
724    }
725}
726
727fn append_assistant_text(out: &mut String, assistant: &AssistantMessage) {
728    push_message_separator(out);
729    out.push_str("[Assistant]: ");
730    let mut first = true;
731    for block in &assistant.content {
732        if let ContentBlock::Text(text) = block {
733            if !first {
734                out.push('\n');
735            }
736            out.push_str(&text.text);
737            first = false;
738        }
739    }
740}
741
742fn append_tool_call_arguments(out: &mut String, arguments: &Value) {
743    if let Some(obj) = arguments.as_object() {
744        let mut first_kv = true;
745        for (k, v) in obj {
746            if !first_kv {
747                out.push_str(", ");
748            }
749            out.push_str(k);
750            out.push('=');
751            match serde_json::to_string(v) {
752                Ok(s) => out.push_str(&s),
753                Err(_) => {
754                    let _ = write!(out, "{v}");
755                }
756            }
757            first_kv = false;
758        }
759    } else {
760        match serde_json::to_string(arguments) {
761            Ok(s) => out.push_str(&s),
762            Err(_) => {
763                let _ = write!(out, "{arguments}");
764            }
765        }
766    }
767}
768
769fn append_assistant_tool_calls(out: &mut String, assistant: &AssistantMessage) {
770    push_message_separator(out);
771    out.push_str("[Assistant tool calls]: ");
772    let mut first = true;
773    for block in &assistant.content {
774        if let ContentBlock::ToolCall(call) = block {
775            if !first {
776                out.push_str("; ");
777            }
778            out.push_str(&call.name);
779            out.push('(');
780            append_tool_call_arguments(out, &call.arguments);
781            out.push(')');
782            first = false;
783        }
784    }
785}
786
787fn append_assistant_message(out: &mut String, assistant: &AssistantMessage) {
788    let (has_thinking, has_text, has_tools) = assistant_content_flags(assistant);
789    if has_thinking {
790        append_assistant_thinking(out, assistant);
791    }
792    if has_text {
793        append_assistant_text(out, assistant);
794    }
795    if has_tools {
796        append_assistant_tool_calls(out, assistant);
797    }
798}
799
800fn tool_result_has_serializable_content(content: &[ContentBlock]) -> bool {
801    content
802        .iter()
803        .any(|c| matches!(c, ContentBlock::Text(t) if !t.text.is_empty()))
804}
805
806fn append_tool_result_message(out: &mut String, content: &[ContentBlock]) {
807    if !tool_result_has_serializable_content(content) {
808        return;
809    }
810
811    push_message_separator(out);
812    out.push_str("[Tool result]: ");
813    for block in content {
814        if let ContentBlock::Text(text) = block {
815            out.push_str(&text.text);
816        }
817    }
818}
819
820fn collect_text_blocks(blocks: &[ContentBlock]) -> String {
821    let mut out = String::new();
822    let mut first = true;
823    for block in blocks {
824        if let ContentBlock::Text(text) = block {
825            if !first {
826                out.push('\n');
827            }
828            out.push_str(&text.text);
829            first = false;
830        }
831    }
832    out
833}
834
835fn serialize_conversation(messages: &[Message]) -> String {
836    let mut out = String::new();
837
838    for msg in messages {
839        match msg {
840            Message::User(user) => append_user_message(&mut out, user),
841            Message::Custom(custom) => {
842                append_custom_message(&mut out, &custom.custom_type, &custom.content);
843            }
844            Message::Assistant(assistant) => append_assistant_message(&mut out, assistant),
845            Message::ToolResult(tool) => append_tool_result_message(&mut out, &tool.content),
846        }
847    }
848
849    out
850}
851
852async fn complete_simple(
853    provider: Arc<dyn Provider>,
854    system_prompt: &str,
855    prompt_text: String,
856    api_key: &str,
857    reserve_tokens: u32,
858    max_tokens_factor: f64,
859) -> Result<AssistantMessage> {
860    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
861    let max_tokens = (f64::from(reserve_tokens) * max_tokens_factor).floor() as u32;
862    let max_tokens = max_tokens.max(256);
863
864    let context = Context {
865        system_prompt: Some(system_prompt.to_string().into()),
866        messages: vec![Message::User(UserMessage {
867            content: UserContent::Blocks(vec![ContentBlock::Text(TextContent::new(prompt_text))]),
868            timestamp: chrono::Utc::now().timestamp_millis(),
869        })]
870        .into(),
871        tools: Vec::new().into(),
872    };
873
874    let options = StreamOptions {
875        api_key: Some(api_key.to_string()),
876        max_tokens: Some(max_tokens),
877        thinking_level: Some(ThinkingLevel::High),
878        ..Default::default()
879    };
880
881    let mut stream = provider.stream(&context, &options).await?;
882    let mut final_message: Option<AssistantMessage> = None;
883
884    while let Some(event) = stream.next().await {
885        match event? {
886            crate::model::StreamEvent::Done { message, .. } => {
887                final_message = Some(message);
888            }
889            crate::model::StreamEvent::Error { error, .. } => {
890                let msg = error
891                    .error_message
892                    .unwrap_or_else(|| "Summarization error".to_string());
893                return Err(Error::api(msg));
894            }
895            _ => {}
896        }
897    }
898
899    let message = final_message.ok_or_else(|| Error::api("Stream ended without Done event"))?;
900    if matches!(message.stop_reason, StopReason::Aborted | StopReason::Error) {
901        let msg = message
902            .error_message
903            .unwrap_or_else(|| "Summarization error".to_string());
904        return Err(Error::api(msg));
905    }
906    Ok(message)
907}
908
909async fn generate_summary(
910    messages: &[SessionMessage],
911    provider: Arc<dyn Provider>,
912    api_key: &str,
913    settings: &ResolvedCompactionSettings,
914    custom_instructions: Option<&str>,
915    previous_summary: Option<&str>,
916) -> Result<String> {
917    let base_prompt = if previous_summary.is_some() {
918        UPDATE_SUMMARIZATION_PROMPT
919    } else {
920        SUMMARIZATION_PROMPT
921    };
922
923    let mut prompt = base_prompt.to_string();
924    if let Some(custom) = custom_instructions.filter(|s| !s.trim().is_empty()) {
925        let _ = write!(prompt, "\n\nAdditional focus: {custom}");
926    }
927
928    let llm_messages = messages
929        .iter()
930        .filter_map(session_message_to_model)
931        .collect::<Vec<_>>();
932    let conversation_text = serialize_conversation(&llm_messages);
933
934    let mut prompt_text = format!("<conversation>\n{conversation_text}\n</conversation>\n\n");
935    if let Some(previous) = previous_summary {
936        let _ = write!(
937            prompt_text,
938            "<previous-summary>\n{previous}\n</previous-summary>\n\n"
939        );
940    }
941    prompt_text.push_str(&prompt);
942
943    let assistant = complete_simple(
944        provider,
945        SUMMARIZATION_SYSTEM_PROMPT,
946        prompt_text,
947        api_key,
948        settings.reserve_tokens,
949        0.8,
950    )
951    .await?;
952
953    let text = collect_text_blocks(&assistant.content);
954
955    if text.trim().is_empty() {
956        return Err(Error::api(
957            "Summarization returned empty text; refusing to store empty compaction summary",
958        ));
959    }
960
961    Ok(text)
962}
963
964async fn generate_turn_prefix_summary(
965    messages: &[SessionMessage],
966    provider: Arc<dyn Provider>,
967    api_key: &str,
968    settings: &ResolvedCompactionSettings,
969) -> Result<String> {
970    let llm_messages = messages
971        .iter()
972        .filter_map(session_message_to_model)
973        .collect::<Vec<_>>();
974    let conversation_text = serialize_conversation(&llm_messages);
975    let prompt_text = format!(
976        "<conversation>\n{conversation_text}\n</conversation>\n\n{TURN_PREFIX_SUMMARIZATION_PROMPT}"
977    );
978
979    let assistant = complete_simple(
980        provider,
981        SUMMARIZATION_SYSTEM_PROMPT,
982        prompt_text,
983        api_key,
984        settings.reserve_tokens,
985        0.5,
986    )
987    .await?;
988
989    let text = collect_text_blocks(&assistant.content);
990
991    if text.trim().is_empty() {
992        return Err(Error::api(
993            "Turn prefix summarization returned empty text; refusing to store empty summary",
994        ));
995    }
996
997    Ok(text)
998}
999
1000// =============================================================================
1001// Public API
1002// =============================================================================
1003
1004#[allow(clippy::too_many_lines)]
1005pub fn prepare_compaction(
1006    path_entries: &[SessionEntry],
1007    settings: ResolvedCompactionSettings,
1008) -> Option<CompactionPreparation> {
1009    if path_entries.is_empty() {
1010        return None;
1011    }
1012
1013    if path_entries
1014        .last()
1015        .is_some_and(|entry| matches!(entry, SessionEntry::Compaction(_)))
1016    {
1017        return None;
1018    }
1019
1020    let mut prev_compaction_index: Option<usize> = None;
1021    for (idx, entry) in path_entries.iter().enumerate().rev() {
1022        if matches!(entry, SessionEntry::Compaction(_)) {
1023            prev_compaction_index = Some(idx);
1024            break;
1025        }
1026    }
1027
1028    let boundary_start = prev_compaction_index.map_or(0, |i| i + 1);
1029    let boundary_end = path_entries.len();
1030
1031    let usage_start = prev_compaction_index.unwrap_or(0);
1032    let mut usage_messages = Vec::new();
1033    for entry in &path_entries[usage_start..boundary_end] {
1034        if let Some(msg) = message_from_entry(entry) {
1035            usage_messages.push(msg);
1036        }
1037    }
1038    // Calculate the tokens *currently* occupied by the segment we are about to compact.
1039    // If the segment includes a previous compaction summary, this counts the *summary* tokens,
1040    // not the original uncompressed history tokens. This effectively tracks the "compressed size"
1041    // of the history prior to the new cut point.
1042    let tokens_before = estimate_context_tokens(&usage_messages).tokens;
1043
1044    if !should_compact(tokens_before, settings.context_window_tokens, &settings) {
1045        return None;
1046    }
1047
1048    let cut_point = find_cut_point(
1049        path_entries,
1050        boundary_start,
1051        boundary_end,
1052        settings.keep_recent_tokens,
1053    );
1054
1055    let first_kept_entry = &path_entries[cut_point.first_kept_entry_index];
1056    let first_kept_entry_id = first_kept_entry.base_id()?.clone();
1057
1058    let history_end = if cut_point.is_split_turn {
1059        cut_point.turn_start_index?
1060    } else {
1061        cut_point.first_kept_entry_index
1062    };
1063
1064    let mut messages_to_summarize = Vec::new();
1065    for entry in &path_entries[boundary_start..history_end] {
1066        if let Some(msg) = message_from_entry(entry) {
1067            messages_to_summarize.push(msg);
1068        }
1069    }
1070
1071    let mut turn_prefix_messages = Vec::new();
1072    if cut_point.is_split_turn {
1073        let turn_start = cut_point.turn_start_index?;
1074        for entry in &path_entries[turn_start..cut_point.first_kept_entry_index] {
1075            if let Some(msg) = message_from_entry(entry) {
1076                turn_prefix_messages.push(msg);
1077            }
1078        }
1079    }
1080
1081    // No-op compaction: if there's nothing to summarize, don't issue an LLM call and don't append a
1082    // compaction entry. This can happen early in a session (e.g. session header entries only).
1083    if messages_to_summarize.is_empty() && turn_prefix_messages.is_empty() {
1084        return None;
1085    }
1086
1087    let previous_summary = prev_compaction_index.and_then(|idx| match &path_entries[idx] {
1088        SessionEntry::Compaction(entry) => Some(entry.summary.clone()),
1089        _ => None,
1090    });
1091
1092    let mut file_ops = FileOperations::default();
1093
1094    // Collect file tracking from previous compaction details if pi-generated.
1095    if let Some(idx) = prev_compaction_index {
1096        if let SessionEntry::Compaction(entry) = &path_entries[idx] {
1097            if !entry.from_hook.unwrap_or(false) {
1098                if let Some(details) = entry.details.as_ref().and_then(Value::as_object) {
1099                    if let Some(read_files) = details.get("readFiles").and_then(Value::as_array) {
1100                        for item in read_files.iter().filter_map(Value::as_str) {
1101                            file_ops.read.insert(item.to_string());
1102                        }
1103                    }
1104                    if let Some(modified_files) =
1105                        details.get("modifiedFiles").and_then(Value::as_array)
1106                    {
1107                        for item in modified_files.iter().filter_map(Value::as_str) {
1108                            file_ops.edited.insert(item.to_string());
1109                        }
1110                    }
1111                }
1112            }
1113        }
1114    }
1115
1116    let mut tool_status = build_tool_status_map(&messages_to_summarize);
1117    tool_status.extend(build_tool_status_map(&turn_prefix_messages));
1118
1119    for msg in &messages_to_summarize {
1120        extract_file_ops_from_message(msg, &mut file_ops, &tool_status);
1121    }
1122    for msg in &turn_prefix_messages {
1123        extract_file_ops_from_message(msg, &mut file_ops, &tool_status);
1124    }
1125
1126    Some(CompactionPreparation {
1127        first_kept_entry_id,
1128        messages_to_summarize,
1129        turn_prefix_messages,
1130        is_split_turn: cut_point.is_split_turn,
1131        tokens_before,
1132        previous_summary,
1133        file_ops,
1134        settings,
1135    })
1136}
1137
1138pub async fn summarize_entries(
1139    entries: &[SessionEntry],
1140    provider: Arc<dyn Provider>,
1141    api_key: &str,
1142    reserve_tokens: u32,
1143    custom_instructions: Option<&str>,
1144) -> Result<Option<String>> {
1145    let mut messages = Vec::new();
1146    for entry in entries {
1147        if let Some(message) = message_from_entry(entry) {
1148            messages.push(message);
1149        }
1150    }
1151
1152    if messages.is_empty() {
1153        return Ok(None);
1154    }
1155
1156    let settings = ResolvedCompactionSettings {
1157        enabled: true,
1158        reserve_tokens,
1159        keep_recent_tokens: 0,
1160        ..Default::default()
1161    };
1162
1163    let summary = generate_summary(
1164        &messages,
1165        provider,
1166        api_key,
1167        &settings,
1168        custom_instructions,
1169        None,
1170    )
1171    .await?;
1172
1173    Ok(Some(summary))
1174}
1175
1176pub async fn compact(
1177    preparation: CompactionPreparation,
1178    provider: Arc<dyn Provider>,
1179    api_key: &str,
1180    custom_instructions: Option<&str>,
1181) -> Result<CompactionResult> {
1182    let summary = if preparation.is_split_turn && !preparation.turn_prefix_messages.is_empty() {
1183        let history_summary = if preparation.messages_to_summarize.is_empty() {
1184            "No prior history.".to_string()
1185        } else {
1186            generate_summary(
1187                &preparation.messages_to_summarize,
1188                Arc::clone(&provider),
1189                api_key,
1190                &preparation.settings,
1191                custom_instructions,
1192                preparation.previous_summary.as_deref(),
1193            )
1194            .await?
1195        };
1196
1197        let turn_prefix_summary = generate_turn_prefix_summary(
1198            &preparation.turn_prefix_messages,
1199            Arc::clone(&provider),
1200            api_key,
1201            &preparation.settings,
1202        )
1203        .await?;
1204
1205        format!(
1206            "{history_summary}\n\n---\n\n**Turn Context (split turn):**\n\n{turn_prefix_summary}"
1207        )
1208    } else {
1209        generate_summary(
1210            &preparation.messages_to_summarize,
1211            Arc::clone(&provider),
1212            api_key,
1213            &preparation.settings,
1214            custom_instructions,
1215            preparation.previous_summary.as_deref(),
1216        )
1217        .await?
1218    };
1219
1220    let (read_files, modified_files) = compute_file_lists(&preparation.file_ops);
1221    let details = CompactionDetails {
1222        read_files: read_files.clone(),
1223        modified_files: modified_files.clone(),
1224    };
1225
1226    let mut summary = summary;
1227    summary.push_str(&format_file_operations(&read_files, &modified_files));
1228
1229    Ok(CompactionResult {
1230        summary,
1231        first_kept_entry_id: preparation.first_kept_entry_id,
1232        tokens_before: preparation.tokens_before,
1233        details,
1234    })
1235}
1236
1237pub fn compaction_details_to_value(details: &CompactionDetails) -> Result<Value> {
1238    serde_json::to_value(details).map_err(|e| Error::session(format!("Compaction details: {e}")))
1239}
1240
1241#[cfg(test)]
1242mod tests {
1243    use super::*;
1244    use crate::model::{AssistantMessage, ContentBlock, TextContent, Usage};
1245    use serde_json::json;
1246
1247    fn make_user_text(text: &str) -> SessionMessage {
1248        SessionMessage::User {
1249            content: UserContent::Text(text.to_string()),
1250            timestamp: Some(0),
1251        }
1252    }
1253
1254    fn make_assistant_text(text: &str, input: u64, output: u64) -> SessionMessage {
1255        SessionMessage::Assistant {
1256            message: AssistantMessage {
1257                content: vec![ContentBlock::Text(TextContent::new(text))],
1258                api: String::new(),
1259                provider: String::new(),
1260                model: String::new(),
1261                stop_reason: StopReason::Stop,
1262                error_message: None,
1263                timestamp: 0,
1264                usage: Usage {
1265                    input,
1266                    output,
1267                    cache_read: 0,
1268                    cache_write: 0,
1269                    total_tokens: input + output,
1270                    ..Default::default()
1271                },
1272            },
1273        }
1274    }
1275
1276    fn make_assistant_tool_call(name: &str, args: Value) -> SessionMessage {
1277        SessionMessage::Assistant {
1278            message: AssistantMessage {
1279                content: vec![ContentBlock::ToolCall(ToolCall {
1280                    id: "call_1".to_string(),
1281                    name: name.to_string(),
1282                    arguments: args,
1283                    thought_signature: None,
1284                })],
1285                api: String::new(),
1286                provider: String::new(),
1287                model: String::new(),
1288                stop_reason: StopReason::ToolUse,
1289                error_message: None,
1290                timestamp: 0,
1291                usage: Usage::default(),
1292            },
1293        }
1294    }
1295
1296    fn make_tool_result(text: &str) -> SessionMessage {
1297        SessionMessage::ToolResult {
1298            tool_call_id: "call_1".to_string(),
1299            tool_name: String::new(),
1300            content: vec![ContentBlock::Text(TextContent::new(text))],
1301            details: None,
1302            is_error: false,
1303            timestamp: None,
1304        }
1305    }
1306
1307    // ── calculate_context_tokens ─────────────────────────────────────
1308
1309    #[test]
1310    fn context_tokens_prefers_total_tokens() {
1311        let usage = Usage {
1312            input: 100,
1313            output: 50,
1314            total_tokens: 200,
1315            ..Default::default()
1316        };
1317        assert_eq!(calculate_context_tokens(&usage), 200);
1318    }
1319
1320    #[test]
1321    fn context_tokens_falls_back_to_input_plus_output() {
1322        let usage = Usage {
1323            input: 100,
1324            output: 50,
1325            total_tokens: 0,
1326            ..Default::default()
1327        };
1328        assert_eq!(calculate_context_tokens(&usage), 150);
1329    }
1330
1331    // ── should_compact ───────────────────────────────────────────────
1332
1333    #[test]
1334    fn should_compact_when_over_threshold() {
1335        let settings = ResolvedCompactionSettings {
1336            enabled: true,
1337            reserve_tokens: 10_000,
1338            keep_recent_tokens: 5_000,
1339            ..Default::default()
1340        };
1341        // window=100k, reserve=10k => threshold=90k, context=95k => should compact
1342        assert!(should_compact(95_000, 100_000, &settings));
1343    }
1344
1345    #[test]
1346    fn should_not_compact_when_under_threshold() {
1347        let settings = ResolvedCompactionSettings {
1348            enabled: true,
1349            reserve_tokens: 10_000,
1350            keep_recent_tokens: 5_000,
1351            ..Default::default()
1352        };
1353        // window=100k, reserve=10k => threshold=90k, context=80k => should not compact
1354        assert!(!should_compact(80_000, 100_000, &settings));
1355    }
1356
1357    #[test]
1358    fn should_not_compact_when_disabled() {
1359        let settings = ResolvedCompactionSettings {
1360            enabled: false,
1361            reserve_tokens: 0,
1362            keep_recent_tokens: 0,
1363            ..Default::default()
1364        };
1365        assert!(!should_compact(1_000_000, 100_000, &settings));
1366    }
1367
1368    #[test]
1369    fn should_compact_at_exact_threshold() {
1370        let settings = ResolvedCompactionSettings {
1371            enabled: true,
1372            reserve_tokens: 10_000,
1373            keep_recent_tokens: 5_000,
1374            ..Default::default()
1375        };
1376        // window=100k, reserve=10k => threshold=90k, context=90k => compact
1377        assert!(should_compact(90_000, 100_000, &settings));
1378        // 89999 should not trigger
1379        assert!(!should_compact(89_999, 100_000, &settings));
1380        // 90001 should also trigger
1381        assert!(should_compact(90_001, 100_000, &settings));
1382    }
1383
1384    // ── estimate_tokens ──────────────────────────────────────────────
1385
1386    #[test]
1387    fn estimate_tokens_user_text() {
1388        let msg = make_user_text("hello world"); // 11 chars => ceil(11/3) = 4
1389        assert_eq!(estimate_tokens(&msg), 4);
1390    }
1391
1392    #[test]
1393    fn estimate_tokens_empty_text() {
1394        let msg = make_user_text(""); // 0 chars => 0
1395        assert_eq!(estimate_tokens(&msg), 0);
1396    }
1397
1398    #[test]
1399    fn estimate_tokens_assistant_text() {
1400        let msg = make_assistant_text("hello", 10, 5); // 5 chars => ceil(5/3) = 2
1401        assert_eq!(estimate_tokens(&msg), 2);
1402    }
1403
1404    #[test]
1405    fn estimate_tokens_tool_result() {
1406        let msg = make_tool_result("file contents here"); // 18 chars => ceil(18/3) = 6
1407        assert_eq!(estimate_tokens(&msg), 6);
1408    }
1409
1410    #[test]
1411    fn estimate_tokens_custom_message() {
1412        let msg = SessionMessage::Custom {
1413            custom_type: "system".to_string(),
1414            content: "some custom content".to_string(),
1415            display: true,
1416            details: None,
1417            timestamp: Some(0),
1418        };
1419        // 19 chars => ceil(19/3) = 7
1420        assert_eq!(estimate_tokens(&msg), 7);
1421    }
1422
1423    // ── estimate_context_tokens ──────────────────────────────────────
1424
1425    #[test]
1426    fn estimate_context_with_assistant_usage() {
1427        let messages = vec![
1428            make_user_text("hi"),
1429            make_assistant_text("hello", 50, 10),
1430            make_user_text("bye"),
1431        ];
1432        let estimate = estimate_context_tokens(&messages);
1433        // Last assistant usage: input=50, output=10, total=60
1434        // Trailing after that: "bye" = ceil(3/3) = 1
1435        assert_eq!(estimate.tokens, 61);
1436        assert_eq!(estimate.last_usage_index, Some(1));
1437    }
1438
1439    #[test]
1440    fn estimate_context_no_assistant() {
1441        let messages = vec![make_user_text("hello"), make_user_text("world")];
1442        let estimate = estimate_context_tokens(&messages);
1443        // No assistant messages, so sum estimate_tokens for all: ceil(5/3)+ceil(5/3) = 2+2 = 4
1444        assert_eq!(estimate.tokens, 4);
1445        assert!(estimate.last_usage_index.is_none());
1446    }
1447
1448    #[test]
1449    fn estimate_context_zero_usage_falls_back_to_heuristics() {
1450        let messages = vec![
1451            make_user_text("hi"),
1452            make_assistant_text("hello", 0, 0),
1453            make_user_text("bye"),
1454        ];
1455        let estimate = estimate_context_tokens(&messages);
1456        // Zero provider usage should not collapse the estimate to trailing
1457        // messages only. We should fall back to whole-history heuristics:
1458        // "hi" => 1, "hello" => 2, "bye" => 1.
1459        assert_eq!(estimate.tokens, 4);
1460        assert!(estimate.last_usage_index.is_none());
1461    }
1462
1463    // ── extract_file_ops_from_message ────────────────────────────────
1464
1465    #[test]
1466    fn extract_file_ops_read() {
1467        let msg = make_assistant_tool_call("read", json!({"path": "/foo/bar.rs"}));
1468        let mut ops = FileOperations::default();
1469        let mut status = HashMap::new();
1470        status.insert("call_1".to_string(), true);
1471        extract_file_ops_from_message(&msg, &mut ops, &status);
1472        assert!(ops.read.contains("/foo/bar.rs"));
1473        assert!(ops.written.is_empty());
1474        assert!(ops.edited.is_empty());
1475    }
1476
1477    #[test]
1478    fn extract_file_ops_write() {
1479        let msg = make_assistant_tool_call("write", json!({"path": "/out.txt"}));
1480        let mut ops = FileOperations::default();
1481        let mut status = HashMap::new();
1482        status.insert("call_1".to_string(), true);
1483        extract_file_ops_from_message(&msg, &mut ops, &status);
1484        assert!(ops.written.contains("/out.txt"));
1485        assert!(ops.read.is_empty());
1486    }
1487
1488    #[test]
1489    fn extract_file_ops_edit() {
1490        let msg = make_assistant_tool_call("edit", json!({"path": "/src/main.rs"}));
1491        let mut ops = FileOperations::default();
1492        let mut status = HashMap::new();
1493        status.insert("call_1".to_string(), true);
1494        extract_file_ops_from_message(&msg, &mut ops, &status);
1495        assert!(ops.edited.contains("/src/main.rs"));
1496    }
1497
1498    #[test]
1499    fn extract_file_ops_ignores_failed_tools() {
1500        let msg = make_assistant_tool_call("read", json!({"path": "/secret.rs"}));
1501        let mut ops = FileOperations::default();
1502        let mut status = HashMap::new();
1503        status.insert("call_1".to_string(), false); // Failed!
1504        extract_file_ops_from_message(&msg, &mut ops, &status);
1505        assert!(ops.read.is_empty());
1506    }
1507
1508    #[test]
1509    fn extract_file_ops_ignores_other_tools() {
1510        let msg = make_assistant_tool_call("bash", json!({"command": "ls"}));
1511        let mut ops = FileOperations::default();
1512        let mut status = HashMap::new();
1513        status.insert("call_1".to_string(), true);
1514        extract_file_ops_from_message(&msg, &mut ops, &status);
1515        assert!(ops.read.is_empty());
1516        assert!(ops.written.is_empty());
1517        assert!(ops.edited.is_empty());
1518    }
1519
1520    #[test]
1521    fn extract_file_ops_ignores_user_messages() {
1522        let msg = make_user_text("read the file /foo.rs");
1523        let mut ops = FileOperations::default();
1524        let status = HashMap::new();
1525        extract_file_ops_from_message(&msg, &mut ops, &status);
1526        assert!(ops.read.is_empty());
1527    }
1528
1529    // ── compute_file_lists ───────────────────────────────────────────
1530
1531    #[test]
1532    fn compute_file_lists_separates_read_from_modified() {
1533        let mut ops = FileOperations::default();
1534        ops.read.insert("/a.rs".to_string());
1535        ops.read.insert("/b.rs".to_string());
1536        ops.written.insert("/b.rs".to_string());
1537        ops.edited.insert("/c.rs".to_string());
1538
1539        let (read_only, modified) = compute_file_lists(&ops);
1540        // /a.rs was only read; /b.rs was read AND written (so it's modified)
1541        assert_eq!(read_only, vec!["/a.rs"]);
1542        assert!(modified.contains(&"/b.rs".to_string()));
1543        assert!(modified.contains(&"/c.rs".to_string()));
1544    }
1545
1546    #[test]
1547    fn compute_file_lists_empty() {
1548        let ops = FileOperations::default();
1549        let (read_only, modified) = compute_file_lists(&ops);
1550        assert!(read_only.is_empty());
1551        assert!(modified.is_empty());
1552    }
1553
1554    // ── format_file_operations ───────────────────────────────────────
1555
1556    #[test]
1557    fn format_file_operations_empty() {
1558        assert_eq!(format_file_operations(&[], &[]), String::new());
1559    }
1560
1561    #[test]
1562    fn format_file_operations_read_only() {
1563        let result = format_file_operations(&["src/main.rs".to_string()], &[]);
1564        assert!(result.contains("<read-files>"));
1565        assert!(result.contains("src/main.rs"));
1566        assert!(!result.contains("<modified-files>"));
1567    }
1568
1569    #[test]
1570    fn format_file_operations_both() {
1571        let result = format_file_operations(&["a.rs".to_string()], &["b.rs".to_string()]);
1572        assert!(result.contains("<read-files>"));
1573        assert!(result.contains("a.rs"));
1574        assert!(result.contains("<modified-files>"));
1575        assert!(result.contains("b.rs"));
1576    }
1577
1578    // ── compaction_details_to_value ──────────────────────────────────
1579
1580    #[test]
1581    fn compaction_details_serializes() {
1582        let details = CompactionDetails {
1583            read_files: vec!["a.rs".to_string()],
1584            modified_files: vec!["b.rs".to_string()],
1585        };
1586        let value = compaction_details_to_value(&details).unwrap();
1587        assert_eq!(value["readFiles"], json!(["a.rs"]));
1588        assert_eq!(value["modifiedFiles"], json!(["b.rs"]));
1589    }
1590
1591    // ── ResolvedCompactionSettings default ───────────────────────────
1592
1593    #[test]
1594    fn default_settings() {
1595        let settings = ResolvedCompactionSettings::default();
1596        assert!(settings.enabled);
1597        assert_eq!(settings.context_window_tokens, 128_000);
1598        assert_eq!(settings.reserve_tokens, 10_240);
1599        assert_eq!(settings.keep_recent_tokens, 12_800);
1600    }
1601
1602    // ── Helper: entry constructors ──────────────────────────────────
1603
1604    use crate::model::{ImageContent, ThinkingContent};
1605    use crate::session::{
1606        BranchSummaryEntry, CompactionEntry, EntryBase, MessageEntry, ModelChangeEntry,
1607    };
1608    use std::collections::HashMap;
1609
1610    fn test_base(id: &str) -> EntryBase {
1611        EntryBase {
1612            id: Some(id.to_string()),
1613            parent_id: None,
1614            timestamp: "2026-01-01T00:00:00.000Z".to_string(),
1615        }
1616    }
1617
1618    fn user_entry(id: &str, text: &str) -> SessionEntry {
1619        SessionEntry::Message(MessageEntry {
1620            base: test_base(id),
1621            message: make_user_text(text),
1622        })
1623    }
1624
1625    fn assistant_entry(id: &str, text: &str, input: u64, output: u64) -> SessionEntry {
1626        SessionEntry::Message(MessageEntry {
1627            base: test_base(id),
1628            message: make_assistant_text(text, input, output),
1629        })
1630    }
1631
1632    fn tool_call_entry(id: &str, tool_name: &str, path: &str) -> SessionEntry {
1633        SessionEntry::Message(MessageEntry {
1634            base: test_base(id),
1635            message: make_assistant_tool_call(tool_name, json!({"path": path})),
1636        })
1637    }
1638
1639    fn tool_result_entry(id: &str, text: &str) -> SessionEntry {
1640        SessionEntry::Message(MessageEntry {
1641            base: test_base(id),
1642            message: make_tool_result(text),
1643        })
1644    }
1645
1646    fn branch_entry(id: &str, summary: &str) -> SessionEntry {
1647        SessionEntry::BranchSummary(BranchSummaryEntry {
1648            base: test_base(id),
1649            from_id: "parent".to_string(),
1650            summary: summary.to_string(),
1651            details: None,
1652            from_hook: None,
1653        })
1654    }
1655
1656    fn compact_entry(id: &str, summary: &str, tokens: u64) -> SessionEntry {
1657        SessionEntry::Compaction(CompactionEntry {
1658            base: test_base(id),
1659            summary: summary.to_string(),
1660            first_kept_entry_id: "kept".to_string(),
1661            tokens_before: tokens,
1662            details: None,
1663            from_hook: None,
1664        })
1665    }
1666
1667    fn bash_entry(id: &str) -> SessionEntry {
1668        SessionEntry::Message(MessageEntry {
1669            base: test_base(id),
1670            message: SessionMessage::BashExecution {
1671                command: "ls".to_string(),
1672                output: "ok".to_string(),
1673                exit_code: 0,
1674                cancelled: None,
1675                truncated: None,
1676                full_output_path: None,
1677                timestamp: None,
1678                extra: HashMap::new(),
1679            },
1680        })
1681    }
1682
1683    // ── get_assistant_usage ─────────────────────────────────────────
1684
1685    #[test]
1686    fn get_assistant_usage_returns_usage_for_stop() {
1687        let msg = make_assistant_text("text", 100, 50);
1688        let usage = get_assistant_usage(&msg);
1689        assert!(usage.is_some());
1690        assert_eq!(usage.unwrap().input, 100);
1691    }
1692
1693    #[test]
1694    fn get_assistant_usage_none_for_aborted() {
1695        let msg = SessionMessage::Assistant {
1696            message: AssistantMessage {
1697                content: vec![ContentBlock::Text(TextContent::new("text"))],
1698                api: String::new(),
1699                provider: String::new(),
1700                model: String::new(),
1701                stop_reason: StopReason::Aborted,
1702                error_message: None,
1703                timestamp: 0,
1704                usage: Usage {
1705                    input: 100,
1706                    output: 50,
1707                    total_tokens: 150,
1708                    ..Default::default()
1709                },
1710            },
1711        };
1712        assert!(get_assistant_usage(&msg).is_none());
1713    }
1714
1715    #[test]
1716    fn get_assistant_usage_none_for_error() {
1717        let msg = SessionMessage::Assistant {
1718            message: AssistantMessage {
1719                content: vec![],
1720                api: String::new(),
1721                provider: String::new(),
1722                model: String::new(),
1723                stop_reason: StopReason::Error,
1724                error_message: None,
1725                timestamp: 0,
1726                usage: Usage::default(),
1727            },
1728        };
1729        assert!(get_assistant_usage(&msg).is_none());
1730    }
1731
1732    #[test]
1733    fn get_assistant_usage_none_for_user() {
1734        assert!(get_assistant_usage(&make_user_text("hello")).is_none());
1735    }
1736
1737    // ── entry_is_message_like ───────────────────────────────────────
1738
1739    #[test]
1740    fn entry_is_message_like_for_message() {
1741        assert!(entry_is_message_like(&user_entry("1", "hi")));
1742    }
1743
1744    #[test]
1745    fn entry_is_message_like_for_branch_summary() {
1746        assert!(entry_is_message_like(&branch_entry("1", "sum")));
1747    }
1748
1749    #[test]
1750    fn entry_is_message_like_false_for_compaction() {
1751        assert!(!entry_is_message_like(&compact_entry("1", "sum", 100)));
1752    }
1753
1754    #[test]
1755    fn entry_is_message_like_false_for_model_change() {
1756        let entry = SessionEntry::ModelChange(ModelChangeEntry {
1757            base: test_base("1"),
1758            provider: "test".to_string(),
1759            model_id: "model-1".to_string(),
1760        });
1761        assert!(!entry_is_message_like(&entry));
1762    }
1763
1764    // ── entry_is_compaction_boundary ────────────────────────────────
1765
1766    #[test]
1767    fn compaction_boundary_true_for_compaction() {
1768        assert!(entry_is_compaction_boundary(&compact_entry(
1769            "1", "sum", 100
1770        )));
1771    }
1772
1773    #[test]
1774    fn compaction_boundary_false_for_message() {
1775        assert!(!entry_is_compaction_boundary(&user_entry("1", "hi")));
1776    }
1777
1778    #[test]
1779    fn compaction_boundary_false_for_branch() {
1780        assert!(!entry_is_compaction_boundary(&branch_entry("1", "sum")));
1781    }
1782
1783    // ── is_user_turn_start ──────────────────────────────────────────
1784
1785    #[test]
1786    fn user_turn_start_for_user() {
1787        assert!(is_user_turn_start(&user_entry("1", "hello")));
1788    }
1789
1790    #[test]
1791    fn user_turn_start_for_branch() {
1792        assert!(is_user_turn_start(&branch_entry("1", "summary")));
1793    }
1794
1795    #[test]
1796    fn user_turn_start_for_bash() {
1797        assert!(is_user_turn_start(&bash_entry("1")));
1798    }
1799
1800    #[test]
1801    fn user_turn_start_false_for_assistant() {
1802        assert!(!is_user_turn_start(&assistant_entry("1", "resp", 10, 5)));
1803    }
1804
1805    #[test]
1806    fn user_turn_start_false_for_tool_result() {
1807        assert!(!is_user_turn_start(&tool_result_entry("1", "result")));
1808    }
1809
1810    #[test]
1811    fn user_turn_start_false_for_compaction() {
1812        assert!(!is_user_turn_start(&compact_entry("1", "sum", 100)));
1813    }
1814
1815    // ── message_from_entry ──────────────────────────────────────────
1816
1817    #[test]
1818    fn message_from_entry_user() {
1819        let entry = user_entry("1", "hello");
1820        let msg = message_from_entry(&entry);
1821        assert!(msg.is_some());
1822        assert!(matches!(msg.unwrap(), SessionMessage::User { .. }));
1823    }
1824
1825    #[test]
1826    fn message_from_entry_branch_summary() {
1827        let entry = branch_entry("1", "branch summary text");
1828        let msg = message_from_entry(&entry).unwrap();
1829        if let SessionMessage::BranchSummary { summary, from_id } = msg {
1830            assert_eq!(summary, "branch summary text");
1831            assert_eq!(from_id, "parent");
1832        } else {
1833            panic!();
1834        }
1835    }
1836
1837    #[test]
1838    fn message_from_entry_compaction() {
1839        let entry = compact_entry("1", "compact summary", 500);
1840        let msg = message_from_entry(&entry).unwrap();
1841        if let SessionMessage::CompactionSummary {
1842            summary,
1843            tokens_before,
1844        } = msg
1845        {
1846            assert_eq!(summary, "compact summary");
1847            assert_eq!(tokens_before, 500);
1848        } else {
1849            panic!();
1850        }
1851    }
1852
1853    #[test]
1854    fn message_from_entry_model_change_is_none() {
1855        let entry = SessionEntry::ModelChange(ModelChangeEntry {
1856            base: test_base("1"),
1857            provider: "test".to_string(),
1858            model_id: "model".to_string(),
1859        });
1860        assert!(message_from_entry(&entry).is_none());
1861    }
1862
1863    // ── find_valid_cut_points ───────────────────────────────────────
1864
1865    #[test]
1866    fn find_valid_cut_points_empty() {
1867        assert!(find_valid_cut_points(&[], 0, 0).is_empty());
1868    }
1869
1870    #[test]
1871    fn find_valid_cut_points_skips_tool_results() {
1872        let entries = vec![
1873            user_entry("1", "hello"),
1874            assistant_entry("2", "resp", 10, 5),
1875            tool_result_entry("3", "result"),
1876            user_entry("4", "follow up"),
1877        ];
1878        let cuts = find_valid_cut_points(&entries, 0, entries.len());
1879        assert!(cuts.contains(&0)); // user
1880        assert!(cuts.contains(&1)); // assistant
1881        assert!(!cuts.contains(&2)); // tool result excluded
1882        assert!(cuts.contains(&3)); // user
1883    }
1884
1885    #[test]
1886    fn find_valid_cut_points_includes_branch_summary() {
1887        let entries = vec![branch_entry("1", "summary"), user_entry("2", "hello")];
1888        let cuts = find_valid_cut_points(&entries, 0, entries.len());
1889        assert!(cuts.contains(&0));
1890        assert!(cuts.contains(&1));
1891    }
1892
1893    #[test]
1894    fn find_valid_cut_points_respects_range() {
1895        let entries = vec![
1896            user_entry("1", "a"),
1897            user_entry("2", "b"),
1898            user_entry("3", "c"),
1899        ];
1900        let cuts = find_valid_cut_points(&entries, 1, 2);
1901        assert!(!cuts.contains(&0));
1902        assert!(cuts.contains(&1));
1903        assert!(!cuts.contains(&2));
1904    }
1905
1906    // ── find_turn_start_index ───────────────────────────────────────
1907
1908    #[test]
1909    fn find_turn_start_basic() {
1910        let entries = vec![
1911            user_entry("1", "hello"),
1912            assistant_entry("2", "resp", 10, 5),
1913            tool_result_entry("3", "result"),
1914        ];
1915        assert_eq!(find_turn_start_index(&entries, 2, 0), Some(0));
1916    }
1917
1918    #[test]
1919    fn find_turn_start_at_self() {
1920        let entries = vec![user_entry("1", "hello")];
1921        assert_eq!(find_turn_start_index(&entries, 0, 0), Some(0));
1922    }
1923
1924    #[test]
1925    fn find_turn_start_none_no_user() {
1926        let entries = vec![
1927            assistant_entry("1", "resp", 10, 5),
1928            tool_result_entry("2", "result"),
1929        ];
1930        assert_eq!(find_turn_start_index(&entries, 1, 0), None);
1931    }
1932
1933    #[test]
1934    fn find_turn_start_respects_start_index() {
1935        let entries = vec![
1936            user_entry("1", "old"),
1937            assistant_entry("2", "resp", 10, 5),
1938            user_entry("3", "new"),
1939        ];
1940        // start_index=2, so it should find user at 2
1941        assert_eq!(find_turn_start_index(&entries, 2, 2), Some(2));
1942        // start_index=2, looking back from 2, user at 1 is below start
1943        assert_eq!(find_turn_start_index(&entries, 1, 2), None);
1944    }
1945
1946    // ── serialize_conversation ───────────────────────────────────────
1947
1948    #[test]
1949    fn serialize_conversation_user_text() {
1950        let messages = vec![Message::User(crate::model::UserMessage {
1951            content: UserContent::Text("hello world".to_string()),
1952            timestamp: 0,
1953        })];
1954        assert_eq!(serialize_conversation(&messages), "[User]: hello world");
1955    }
1956
1957    #[test]
1958    fn serialize_conversation_empty() {
1959        assert!(serialize_conversation(&[]).is_empty());
1960    }
1961
1962    #[test]
1963    fn serialize_conversation_skips_empty_user() {
1964        let messages = vec![Message::User(crate::model::UserMessage {
1965            content: UserContent::Text(String::new()),
1966            timestamp: 0,
1967        })];
1968        assert!(serialize_conversation(&messages).is_empty());
1969    }
1970
1971    #[test]
1972    fn serialize_conversation_assistant_text() {
1973        let messages = vec![Message::assistant(AssistantMessage {
1974            content: vec![ContentBlock::Text(TextContent::new("response"))],
1975            api: String::new(),
1976            provider: String::new(),
1977            model: String::new(),
1978            usage: Usage::default(),
1979            stop_reason: StopReason::Stop,
1980            error_message: None,
1981            timestamp: 0,
1982        })];
1983        assert!(serialize_conversation(&messages).contains("[Assistant]: response"));
1984    }
1985
1986    #[test]
1987    fn serialize_conversation_tool_calls() {
1988        let messages = vec![Message::assistant(AssistantMessage {
1989            content: vec![ContentBlock::ToolCall(ToolCall {
1990                id: "c1".to_string(),
1991                name: "read".to_string(),
1992                arguments: json!({"path": "/main.rs"}),
1993                thought_signature: None,
1994            })],
1995            api: String::new(),
1996            provider: String::new(),
1997            model: String::new(),
1998            usage: Usage::default(),
1999            stop_reason: StopReason::Stop,
2000            error_message: None,
2001            timestamp: 0,
2002        })];
2003        let result = serialize_conversation(&messages);
2004        assert!(result.contains("[Assistant tool calls]: read("));
2005        assert!(result.contains("path="));
2006    }
2007
2008    #[test]
2009    fn serialize_conversation_thinking() {
2010        let messages = vec![Message::assistant(AssistantMessage {
2011            content: vec![ContentBlock::Thinking(ThinkingContent {
2012                thinking: "let me think".to_string(),
2013                thinking_signature: None,
2014            })],
2015            api: String::new(),
2016            provider: String::new(),
2017            model: String::new(),
2018            usage: Usage::default(),
2019            stop_reason: StopReason::Stop,
2020            error_message: None,
2021            timestamp: 0,
2022        })];
2023        assert!(serialize_conversation(&messages).contains("[Assistant thinking]: let me think"));
2024    }
2025
2026    #[test]
2027    fn serialize_conversation_tool_result() {
2028        let messages = vec![Message::tool_result(crate::model::ToolResultMessage {
2029            tool_call_id: "c1".to_string(),
2030            tool_name: "read".to_string(),
2031            content: vec![ContentBlock::Text(TextContent::new("file contents"))],
2032            details: None,
2033            is_error: false,
2034            timestamp: 0,
2035        })];
2036        assert!(serialize_conversation(&messages).contains("[Tool result]: file contents"));
2037    }
2038
2039    // ── estimate_tokens additional ──────────────────────────────────
2040
2041    #[test]
2042    fn estimate_tokens_image_block() {
2043        let msg = SessionMessage::User {
2044            content: UserContent::Blocks(vec![ContentBlock::Image(ImageContent {
2045                data: "base64data".to_string(),
2046                mime_type: "image/png".to_string(),
2047            })]),
2048            timestamp: None,
2049        };
2050        // Image = 3600 chars (IMAGE_CHAR_ESTIMATE) -> ceil(3600/3) = 1200
2051        assert_eq!(estimate_tokens(&msg), 1200);
2052    }
2053
2054    #[test]
2055    fn estimate_tokens_thinking() {
2056        let msg = SessionMessage::User {
2057            content: UserContent::Blocks(vec![ContentBlock::Thinking(ThinkingContent {
2058                thinking: "a".repeat(20),
2059                thinking_signature: None,
2060            })]),
2061            timestamp: None,
2062        };
2063        // 20 chars -> ceil(20/3) = 7
2064        assert_eq!(estimate_tokens(&msg), 7);
2065    }
2066
2067    #[test]
2068    fn estimate_tokens_bash_execution() {
2069        let msg = SessionMessage::BashExecution {
2070            command: "echo hi".to_string(),
2071            output: "hi\n".to_string(),
2072            exit_code: 0,
2073            cancelled: None,
2074            truncated: None,
2075            full_output_path: None,
2076            timestamp: None,
2077            extra: HashMap::new(),
2078        };
2079        // 7 + 3 = 10 chars -> ceil(10/3) = 4
2080        assert_eq!(estimate_tokens(&msg), 4);
2081    }
2082
2083    #[test]
2084    fn estimate_tokens_branch_summary() {
2085        let msg = SessionMessage::BranchSummary {
2086            summary: "a".repeat(40),
2087            from_id: "id".to_string(),
2088        };
2089        // 40 chars -> ceil(40/3) = 14
2090        assert_eq!(estimate_tokens(&msg), 14);
2091    }
2092
2093    #[test]
2094    fn estimate_tokens_compaction_summary() {
2095        let msg = SessionMessage::CompactionSummary {
2096            summary: "a".repeat(80),
2097            tokens_before: 5000,
2098        };
2099        // 80 chars -> ceil(80/3) = 27
2100        assert_eq!(estimate_tokens(&msg), 27);
2101    }
2102
2103    // ── prepare_compaction ──────────────────────────────────────────
2104
2105    #[test]
2106    fn prepare_compaction_empty() {
2107        assert!(prepare_compaction(&[], ResolvedCompactionSettings::default()).is_none());
2108    }
2109
2110    #[test]
2111    fn prepare_compaction_last_is_compaction_returns_none() {
2112        let entries = vec![user_entry("1", "hello"), compact_entry("2", "summary", 100)];
2113        assert!(prepare_compaction(&entries, ResolvedCompactionSettings::default()).is_none());
2114    }
2115
2116    #[test]
2117    fn prepare_compaction_no_messages_to_summarize_returns_none() {
2118        // Only non-message entries that produce no summarizable messages
2119        let entries = vec![SessionEntry::ModelChange(ModelChangeEntry {
2120            base: test_base("1"),
2121            provider: "test".to_string(),
2122            model_id: "model".to_string(),
2123        })];
2124        assert!(prepare_compaction(&entries, ResolvedCompactionSettings::default()).is_none());
2125    }
2126
2127    #[test]
2128    fn prepare_compaction_basic_returns_some() {
2129        let long_text = "a".repeat(100_000);
2130        let entries = vec![
2131            user_entry("1", &long_text),
2132            assistant_entry("2", &long_text, 50000, 25000),
2133            user_entry("3", &long_text),
2134            assistant_entry("4", &long_text, 80000, 30000),
2135            user_entry("5", "recent"),
2136        ];
2137        let settings = ResolvedCompactionSettings {
2138            enabled: true,
2139            context_window_tokens: 100_000,
2140            reserve_tokens: 1000,
2141            keep_recent_tokens: 100,
2142        };
2143        let prep = prepare_compaction(&entries, settings);
2144        assert!(prep.is_some());
2145        let p = prep.unwrap();
2146        assert!(!p.messages_to_summarize.is_empty());
2147        assert!(p.tokens_before > 0);
2148        assert!(p.previous_summary.is_none());
2149    }
2150
2151    #[test]
2152    fn prepare_compaction_after_previous_compaction() {
2153        let entries = vec![
2154            user_entry("1", "old message"),
2155            assistant_entry("2", "old response", 100, 50),
2156            compact_entry("3", "previous summary", 300),
2157            user_entry("4", &"x".repeat(100_000)),
2158            assistant_entry("5", &"y".repeat(100_000), 80000, 30000),
2159            user_entry("6", "recent"),
2160        ];
2161        let settings = ResolvedCompactionSettings {
2162            enabled: true,
2163            context_window_tokens: 100_000,
2164            reserve_tokens: 1000,
2165            keep_recent_tokens: 100,
2166        };
2167        let prep = prepare_compaction(&entries, settings);
2168        assert!(prep.is_some());
2169        let p = prep.unwrap();
2170        assert_eq!(p.previous_summary.as_deref(), Some("previous summary"));
2171    }
2172
2173    #[test]
2174    fn prepare_compaction_tracks_file_ops() {
2175        let entries = vec![
2176            tool_call_entry("1", "read", "/src/main.rs"),
2177            tool_result_entry("1r", "ok"),
2178            tool_call_entry("2", "edit", "/src/lib.rs"),
2179            tool_result_entry("2r", "ok"),
2180            user_entry("3", &"x".repeat(100_000)),
2181            assistant_entry("4", &"y".repeat(100_000), 80000, 30000),
2182            user_entry("5", "recent"),
2183        ];
2184        let settings = ResolvedCompactionSettings {
2185            enabled: true,
2186            reserve_tokens: 1000,
2187            keep_recent_tokens: 100,
2188            ..Default::default()
2189        };
2190        if let Some(prep) = prepare_compaction(&entries, settings) {
2191            let has_read = prep.file_ops.read.contains("/src/main.rs");
2192            let has_edit = prep.file_ops.edited.contains("/src/lib.rs");
2193            // At least one should be tracked (depends on cut point position)
2194            assert!(has_read || has_edit || prep.file_ops.read.is_empty());
2195        }
2196    }
2197
2198    // ── FileOperations::read_files ──────────────────────────────────
2199
2200    #[test]
2201    fn file_operations_read_files_iterator() {
2202        let mut ops = FileOperations::default();
2203        ops.read.insert("/a.rs".to_string());
2204        ops.read.insert("/b.rs".to_string());
2205        let files: Vec<&str> = ops.read_files().collect();
2206        assert_eq!(files.len(), 2);
2207        assert!(files.contains(&"/a.rs"));
2208        assert!(files.contains(&"/b.rs"));
2209    }
2210
2211    #[test]
2212    fn find_cut_point_includes_tool_result_when_needed() {
2213        // Setup:
2214        // 0. User (10)
2215        // 1. Assistant Call (10)
2216        // 2. Tool Result (100)
2217        // 3. User (10)
2218        // 4. Assistant (10)
2219        //
2220        // Keep recent = 100.
2221        // Accumulation from end:
2222        // 4: 10
2223        // 3: 20
2224        // 2: 120 (Threshold crossed at index 2)
2225        //
2226        // Index 2 is ToolResult (invalid cut point).
2227        // Valid cut points: 0, 1, 3, 4.
2228        //
2229        // Logic should pick closest valid cut point <= 2, which is 1.
2230        // If it picked >= 2, it would pick 3, discarding the ToolResult and Call (keeping only 20 tokens).
2231        // By picking 1, we keep 1..4 (130 tokens).
2232
2233        // Create entries with controlled lengths.
2234        // With chars/token ~=3, 400 chars => ceil(400/3)=134 tokens.
2235        let tr_text = "x".repeat(400);
2236        let entries = vec![
2237            user_entry("0", "user"),              // Valid
2238            assistant_entry("1", "call", 10, 10), // Valid (Assistant)
2239            tool_result_entry("2", &tr_text),     // Invalid
2240            user_entry("3", "user"),              // Valid
2241            assistant_entry("4", "resp", 10, 10), // Valid
2242        ];
2243
2244        // Verify token estimates (approx)
2245        // 0: ceil(4/3) = 2
2246        // 1: ceil(4/3) = 2
2247        // 2: ceil(400/3) = 134
2248        // 3: ceil(4/3) = 2
2249        // 4: ceil(4/3) = 2
2250        // Total recent needed: 100.
2251        // Accumulate: 4(2)+3(2)+2(134) = 138. Crossed at 2.
2252
2253        let settings = ResolvedCompactionSettings {
2254            enabled: true,
2255            context_window_tokens: 15,
2256            reserve_tokens: 0,
2257            keep_recent_tokens: 100,
2258        };
2259
2260        let prep = prepare_compaction(&entries, settings).expect("should compact");
2261
2262        // Cut point is index 1 (Assistant/Call). Because entries[1] is Assistant (not User),
2263        // this is a split turn: the turn started at index 0 (User). The User message at index 0
2264        // goes into turn_prefix_messages (not messages_to_summarize) because history_end = 0.
2265        assert_eq!(prep.first_kept_entry_id, "1");
2266
2267        // messages_to_summarize is entries[0..0] = empty (split-turn puts the
2268        // prefix in turn_prefix_messages instead).
2269        assert!(
2270            prep.messages_to_summarize.is_empty(),
2271            "split turn: user goes into turn prefix, not summarize"
2272        );
2273
2274        // turn_prefix_messages should contain the User message at index 0.
2275        assert_eq!(prep.turn_prefix_messages.len(), 1);
2276        match &prep.turn_prefix_messages[0] {
2277            SessionMessage::User { content, .. } => {
2278                if let UserContent::Text(t) = content {
2279                    assert_eq!(t, "user");
2280                } else {
2281                    panic!();
2282                }
2283            }
2284            _ => panic!(),
2285        }
2286    }
2287
2288    #[test]
2289    fn find_cut_point_should_not_discard_context_to_skip_tool_chain() {
2290        // Setup (estimate_tokens uses ceil(chars/3)):
2291        // 0. User "x"*4000 → 1334 tokens
2292        // 1. Assistant "x"*400 → 134 tokens
2293        // 2. Tool Result "x"*400 → 134 tokens
2294        // 3. User "next" → 2 tokens
2295        //
2296        // Keep recent = 150.
2297        // Accumulation (from end):
2298        // 3: 2
2299        // 2: 136
2300        // 1: 270 (Crosses 150) -> cut_index = 1
2301        //
2302        // The cut should land at index 1 (the assistant message), keeping
2303        // entries 1-3 and summarizing only entry 0.
2304
2305        let entries = vec![
2306            user_entry("0", &"x".repeat(4000)),             // 1000 tokens
2307            assistant_entry("1", &"x".repeat(400), 50, 50), // 100 tokens
2308            tool_result_entry("2", &"x".repeat(400)),       // 100 tokens
2309            user_entry("3", "next"),                        // 1 token
2310        ];
2311
2312        let settings = ResolvedCompactionSettings {
2313            enabled: true,
2314            context_window_tokens: 200,
2315            reserve_tokens: 0,
2316            keep_recent_tokens: 150,
2317        };
2318
2319        // We use prepare_compaction as the entry point
2320        let prep = prepare_compaction(&entries, settings).expect("should compact");
2321
2322        // We expect to keep from 1 (Assistant). The cut splits the turn
2323        // (user 0 + assistant 1), so user 0 goes into the turn prefix.
2324        assert_eq!(
2325            prep.first_kept_entry_id, "1",
2326            "Should start at Assistant message to preserve context"
2327        );
2328        assert!(
2329            prep.is_split_turn,
2330            "Cut should split the user/assistant turn"
2331        );
2332        assert_eq!(
2333            prep.turn_prefix_messages.len(),
2334            1,
2335            "User entry at index 0 should be in the turn prefix"
2336        );
2337        assert!(
2338            prep.messages_to_summarize.is_empty(),
2339            "Nothing before the turn to summarize"
2340        );
2341    }
2342
2343    mod proptest_compaction {
2344        use super::*;
2345        use proptest::prelude::*;
2346
2347        proptest! {
2348            /// `calculate_context_tokens`: if total > 0, returns total.
2349            #[test]
2350            fn calc_context_tokens_total_wins(
2351                input in 0..1_000_000u64,
2352                output in 0..1_000_000u64,
2353                total in 1..2_000_000u64,
2354            ) {
2355                let usage = Usage {
2356                    input,
2357                    output,
2358                    total_tokens: total,
2359                    ..Usage::default()
2360                };
2361                assert_eq!(calculate_context_tokens(&usage), total);
2362            }
2363
2364            /// `calculate_context_tokens`: if total == 0, returns input + output.
2365            #[test]
2366            fn calc_context_tokens_fallback(
2367                input in 0..1_000_000u64,
2368                output in 0..1_000_000u64,
2369            ) {
2370                let usage = Usage {
2371                    input,
2372                    output,
2373                    total_tokens: 0,
2374                    ..Usage::default()
2375                };
2376                assert_eq!(calculate_context_tokens(&usage), input + output);
2377            }
2378
2379            /// `should_compact` returns false when disabled.
2380            #[test]
2381            fn should_compact_disabled_returns_false(
2382                ctx_tokens in 0..1_000_000u64,
2383                window in 0..500_000u32,
2384            ) {
2385                let settings = ResolvedCompactionSettings {
2386                    enabled: false,
2387                    context_window_tokens: window,
2388                    reserve_tokens: 16_384,
2389                    keep_recent_tokens: 20_000,
2390                };
2391                assert!(!should_compact(ctx_tokens, window, &settings));
2392            }
2393
2394            /// `should_compact` threshold: tokens >= window - reserve.
2395            #[test]
2396            fn should_compact_threshold(
2397                ctx_tokens in 0..500_000u64,
2398                window in 0..300_000u32,
2399                reserve in 0..100_000u32,
2400            ) {
2401                let settings = ResolvedCompactionSettings {
2402                    enabled: true,
2403                    context_window_tokens: window,
2404                    reserve_tokens: reserve,
2405                    keep_recent_tokens: 20_000,
2406                };
2407                let threshold = u64::from(window).saturating_sub(u64::from(reserve));
2408                let result = should_compact(ctx_tokens, window, &settings);
2409                assert_eq!(result, ctx_tokens >= threshold);
2410            }
2411
2412            /// `format_file_operations`: empty lists produce empty string.
2413            #[test]
2414            fn format_file_ops_empty(_dummy in 0..10u32) {
2415                let result = format_file_operations(&[], &[]);
2416                assert!(result.is_empty());
2417            }
2418
2419            /// `format_file_operations`: read files produce `<read-files>` tag.
2420            #[test]
2421            fn format_file_ops_read_tag(
2422                files in prop::collection::vec("[a-z./]{1,20}", 1..5),
2423            ) {
2424                let result = format_file_operations(&files, &[]);
2425                assert!(result.contains("<read-files>"));
2426                assert!(result.contains("</read-files>"));
2427                assert!(!result.contains("<modified-files>"));
2428                for f in &files {
2429                    assert!(result.contains(f.as_str()));
2430                }
2431            }
2432
2433            /// `format_file_operations`: modified files produce `<modified-files>` tag.
2434            #[test]
2435            fn format_file_ops_modified_tag(
2436                files in prop::collection::vec("[a-z./]{1,20}", 1..5),
2437            ) {
2438                let result = format_file_operations(&[], &files);
2439                assert!(!result.contains("<read-files>"));
2440                assert!(result.contains("<modified-files>"));
2441                assert!(result.contains("</modified-files>"));
2442                for f in &files {
2443                    assert!(result.contains(f.as_str()));
2444                }
2445            }
2446
2447            /// `compute_file_lists`: modified = edited ∪ written, read_only = read \ modified.
2448            #[test]
2449            fn compute_file_lists_set_algebra(
2450                read in prop::collection::hash_set("[a-z]{1,5}", 0..5),
2451                written in prop::collection::hash_set("[a-z]{1,5}", 0..5),
2452                edited in prop::collection::hash_set("[a-z]{1,5}", 0..5),
2453            ) {
2454                let file_ops = FileOperations {
2455                    read: read.clone(),
2456                    written: written.clone(),
2457                    edited: edited.clone(),
2458                };
2459                let (read_only, modified) = compute_file_lists(&file_ops);
2460                // Modified = edited ∪ written
2461                let expected_modified: HashSet<&String> =
2462                    edited.iter().chain(written.iter()).collect();
2463                let actual_modified: HashSet<&String> = modified.iter().collect();
2464                assert_eq!(actual_modified, expected_modified);
2465                // Read-only = read \ modified (no overlap)
2466                for f in &read_only {
2467                    assert!(!modified.contains(f), "overlap: {f}");
2468                    assert!(read.contains(f));
2469                }
2470                // Both are sorted
2471                for pair in read_only.windows(2) {
2472                    assert!(pair[0] <= pair[1]);
2473                }
2474                for pair in modified.windows(2) {
2475                    assert!(pair[0] <= pair[1]);
2476                }
2477            }
2478        }
2479    }
2480}