Skip to main content

bamboo_compression/
compression_tooling.rs

1use crate::counter::{TiktokenTokenCounter, TokenCounter};
2use crate::limits::{create_budget_for_model, ModelLimitsRegistry};
3use crate::{BudgetStrategy, TokenBudget};
4use bamboo_domain::MessagePhase;
5use bamboo_domain::{
6    CompressionEvent, CompressionTriggerType, ConversationSummary, Message, Session,
7};
8
9/// Checks if a message is part of a skill tool chain (load_skill / read_skill_resource).
10fn is_skill_tool_chain_message(message: &Message) -> bool {
11    message.tool_calls.as_ref().is_some_and(|calls| {
12        calls.iter().any(|call| {
13            matches!(
14                call.function.name.as_str(),
15                "load_skill" | "read_skill_resource"
16            )
17        })
18    })
19}
20
21/// The `tool_call_id`s a message participates in: the ids of the tool calls an
22/// assistant message initiates, plus the id of the call a `tool` result answers.
23/// Two messages belong to the same tool chain iff these overlap.
24fn tool_chain_call_ids(message: &Message) -> impl Iterator<Item = String> + '_ {
25    message
26        .tool_calls
27        .iter()
28        .flatten()
29        .map(|call| call.id.clone())
30        .chain(message.tool_call_id.clone())
31}
32
33/// Close the compressed (`messages_to_summarize`) set over tool chains: repeatedly
34/// move any message still in `messages_to_keep` that shares a `tool_call_id` with
35/// an already-compressed message into the summarize set. This keeps an assistant
36/// `tool_calls` message and its matching `tool` result(s) on the same side — a
37/// split leaves an orphan `tool_result` (or a `tool_use` with no result) in the
38/// active set, which providers reject with a 400 that then poisons every
39/// subsequent request in the session (#340). Protected messages are left in place
40/// (skill chains are already fully protected upstream, so are never partially
41/// compressed; protected user messages carry no `tool_call_id`).
42fn close_compressed_set_over_tool_chains(
43    messages_to_keep: &mut Vec<Message>,
44    messages_to_summarize: &mut Vec<Message>,
45    protected_user_ids: &HashSet<String>,
46    never_compress_ids: &[String],
47) {
48    loop {
49        let compressed_call_ids: HashSet<String> = messages_to_summarize
50            .iter()
51            .flat_map(tool_chain_call_ids)
52            .collect();
53        let split_index = messages_to_keep.iter().position(|message| {
54            !protected_user_ids.contains(message.id.as_str())
55                && !never_compress_ids.contains(&message.id)
56                && tool_chain_call_ids(message).any(|id| compressed_call_ids.contains(&id))
57        });
58        match split_index {
59            Some(index) => messages_to_summarize.push(messages_to_keep.remove(index)),
60            None => break,
61        }
62    }
63}
64use chrono::Utc;
65use std::collections::HashSet;
66
67/// Structured reason why a compression plan could not be built.
68#[derive(Debug, Clone)]
69pub enum CompressionPlanError {
70    /// The exposure gate (threshold not reached) prevented building.
71    ExposureGateNotMet {
72        usage_percent: f64,
73        trigger_percent: u8,
74    },
75    /// No active messages in the session.
76    NoActiveMessages,
77    /// Not enough non-system messages to compress (need >=3).
78    NotEnoughMessages { non_system_count: usize },
79    /// Nothing to compress after anchor/keep splitting.
80    NothingToCompress {
81        anchor_index: usize,
82        non_system_count: usize,
83    },
84}
85
86impl std::fmt::Display for CompressionPlanError {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        match self {
89            Self::ExposureGateNotMet {
90                usage_percent,
91                trigger_percent,
92            } => write!(
93                f,
94                "compression threshold not reached (usage={:.1}%, trigger={}%)",
95                usage_percent, trigger_percent
96            ),
97            Self::NoActiveMessages => write!(f, "no active messages to compress"),
98            Self::NotEnoughMessages { non_system_count } => write!(
99                f,
100                "not enough non-system messages to compress ({}, need >=3)",
101                non_system_count
102            ),
103            Self::NothingToCompress {
104                anchor_index,
105                non_system_count,
106            } => write!(
107                f,
108                "nothing to compress after anchor/keep splitting (anchor_index={}, non_system={})",
109                anchor_index, non_system_count
110            ),
111        }
112    }
113}
114
115/// Metadata about current context pressure, used to decide when compression
116/// should be requested by host-side control flow.
117#[derive(Debug, Clone)]
118pub struct ContextCompressionExposure {
119    pub budget: TokenBudget,
120    pub active_tokens: u32,
121    pub active_usage_percent: f64,
122    pub active_usage_percent_rounded: u8,
123    pub should_expose_tool: bool,
124}
125
126/// A compression plan describing which active historical messages should be
127/// archived and summarized.
128#[derive(Debug, Clone)]
129pub struct CompressionPlan {
130    pub compressed_message_ids: Vec<String>,
131    pub messages_to_summarize: Vec<Message>,
132    pub summary_tokens: u32,
133    pub summary_content: String,
134    pub active_usage_before_percent: f64,
135    pub active_usage_after_percent: f64,
136    pub trigger_percent: u8,
137    pub target_percent: u8,
138    pub segments_removed: usize,
139    pub trigger_type: CompressionTriggerType,
140    pub compression_ratio: f64,
141    pub model_used: Option<String>,
142    pub latency_ms: u64,
143}
144
145pub fn context_window_usage_percent(total_tokens: u32, context_window_tokens: u32) -> f64 {
146    if context_window_tokens == 0 {
147        return 0.0;
148    }
149    (total_tokens as f64 / context_window_tokens as f64) * 100.0
150}
151
152pub fn normalized_trigger_percent(trigger_percent: u8) -> f64 {
153    match trigger_percent {
154        0 => 100.0,
155        1..=100 => trigger_percent as f64,
156        _ => 100.0,
157    }
158}
159
160/// Estimate whether context pressure has crossed the configured threshold for
161/// compression eligibility.
162pub fn estimate_context_compression_exposure(
163    session: &Session,
164    model_name: &str,
165    configured_budget: Option<&TokenBudget>,
166) -> ContextCompressionExposure {
167    // When a budget was already resolved upstream (the production path — see
168    // `resolve_token_budget`, which caches it in `session.resolved_token_budget` (#180),
169    // issue #20 bug 1), use it directly. Only when none is available do we fall
170    // back to a model-derived budget. No `model_limits.json` registry is in
171    // scope synchronously here, so this fallback resolves to the global default
172    // rather than silently fabricating an empty override registry (#20 bug 2).
173    let budget = configured_budget.cloned().unwrap_or_else(|| {
174        create_budget_for_model(
175            model_name,
176            BudgetStrategy::default(),
177            &ModelLimitsRegistry::new(),
178        )
179    });
180    let counter = TiktokenTokenCounter::default();
181    let active_messages = active_messages_for_budget(session);
182    let active_message_tokens = counter.count_messages(&active_messages);
183    let summary_tokens = session
184        .conversation_summary
185        .as_ref()
186        .map(|summary| counter.count_messages(&[compression_summary_message(&summary.content)]))
187        .unwrap_or(0);
188    let active_tokens = active_message_tokens.saturating_add(summary_tokens);
189    // Use context window as the denominator for a single, provider-aligned
190    // pressure scale across backend and frontend.
191    let context_window = budget.max_context_tokens;
192    let estimated_usage = context_window_usage_percent(active_tokens, context_window);
193    let usage = session
194        .token_usage
195        .as_ref()
196        .and_then(|token_usage| {
197            let denominator = if token_usage.max_context_tokens > 0 {
198                token_usage.max_context_tokens
199            } else if token_usage.budget_limit > 0 {
200                // Legacy payload compatibility.
201                token_usage.budget_limit
202            } else {
203                context_window
204            };
205            (denominator > 0).then_some(context_window_usage_percent(
206                token_usage.total_tokens,
207                denominator,
208            ))
209        })
210        .map(|persisted_usage| persisted_usage.max(estimated_usage))
211        .unwrap_or(estimated_usage);
212
213    let rounded = usage.clamp(0.0, 100.0).round() as u8;
214    let trigger_tokens = budget.compression_trigger_context_tokens();
215    let trigger_percent = if budget.max_context_tokens > 0 {
216        (trigger_tokens as f64 / budget.max_context_tokens as f64) * 100.0
217    } else {
218        0.0
219    };
220    let threshold_reached = usage >= trigger_percent;
221
222    // Check non-system message count to stay consistent with the plan
223    // building requirement of >=3 non-system messages.  Using
224    // active_messages.len() would include system messages and expose the
225    // tool even when plan building would immediately fail.
226    let non_system_count = active_messages
227        .iter()
228        .filter(|m| !matches!(m.role, bamboo_domain::Role::System))
229        .count();
230
231    let should_expose_tool = threshold_reached && non_system_count >= 3;
232
233    ContextCompressionExposure {
234        budget,
235        active_tokens,
236        active_usage_percent: usage,
237        active_usage_percent_rounded: rounded,
238        should_expose_tool,
239    }
240}
241
242/// Build a compression plan that archives older active messages and replaces
243/// them with a caller-provided summary.
244pub fn build_compression_plan_with_summary(
245    session: &Session,
246    model_name: &str,
247    configured_budget: Option<&TokenBudget>,
248    summary_content: String,
249) -> Result<CompressionPlan, CompressionPlanError> {
250    build_compression_plan_with_summary_internal(
251        session,
252        model_name,
253        configured_budget,
254        summary_content,
255        true,
256        CompressionTriggerType::Auto,
257    )
258}
259
260/// Build a compression plan while bypassing "tool exposure" gating.
261///
262/// This is intended for host-enforced fallback paths when context pressure is
263/// critically high and compression must be attempted regardless of the normal
264/// trigger gate.
265pub fn build_forced_compression_plan_with_summary(
266    session: &Session,
267    model_name: &str,
268    configured_budget: Option<&TokenBudget>,
269    summary_content: String,
270    trigger_type: CompressionTriggerType,
271) -> Result<CompressionPlan, CompressionPlanError> {
272    build_compression_plan_with_summary_internal(
273        session,
274        model_name,
275        configured_budget,
276        summary_content,
277        false,
278        trigger_type,
279    )
280}
281
282fn build_compression_plan_with_summary_internal(
283    session: &Session,
284    model_name: &str,
285    configured_budget: Option<&TokenBudget>,
286    summary_content: String,
287    require_exposure_gate: bool,
288    trigger_type: CompressionTriggerType,
289) -> Result<CompressionPlan, CompressionPlanError> {
290    let exposure = estimate_context_compression_exposure(session, model_name, configured_budget);
291    if require_exposure_gate && !exposure.should_expose_tool {
292        return Err(CompressionPlanError::ExposureGateNotMet {
293            usage_percent: exposure.active_usage_percent,
294            trigger_percent: exposure.budget.compression_trigger_percent,
295        });
296    }
297
298    let budget = &exposure.budget;
299    let counter = TiktokenTokenCounter::default();
300    let summary_message = compression_summary_message(&summary_content);
301    let summary_tokens = counter.count_messages(&[summary_message]);
302
303    let context_window = budget.max_context_tokens;
304    let target_limit = budget.compression_target_context_tokens();
305
306    let mut active_messages = active_messages_for_budget(session);
307    if active_messages.is_empty() {
308        tracing::debug!("compression plan: no active messages, cannot build plan");
309        return Err(CompressionPlanError::NoActiveMessages);
310    }
311
312    let system_messages: Vec<Message> = active_messages
313        .iter()
314        .filter(|m| matches!(m.role, bamboo_domain::Role::System))
315        .cloned()
316        .collect();
317    let system_tokens = counter.count_messages(&system_messages);
318    let reserved_non_window_tokens = system_tokens.saturating_add(summary_tokens);
319    let window_limit = target_limit.saturating_sub(reserved_non_window_tokens);
320
321    let non_system: Vec<Message> = active_messages
322        .drain(..)
323        .filter(|m| !matches!(m.role, bamboo_domain::Role::System))
324        .collect();
325
326    if non_system.len() < 3 {
327        tracing::debug!(
328            "compression plan: not enough non-system messages ({}), need at least 3",
329            non_system.len()
330        );
331        return Err(CompressionPlanError::NotEnoughMessages {
332            non_system_count: non_system.len(),
333        });
334    }
335
336    let user_indexes = non_system
337        .iter()
338        .enumerate()
339        .filter_map(|(index, message)| {
340            matches!(message.role, bamboo_domain::Role::User).then_some(index)
341        })
342        .collect::<Vec<_>>();
343    let keep_user_count = user_indexes.len().min(3);
344    let anchor_index = if keep_user_count > 0 {
345        user_indexes[user_indexes.len() - keep_user_count]
346    } else {
347        non_system
348            .iter()
349            .rposition(|m| matches!(m.role, bamboo_domain::Role::User))
350            .unwrap_or(non_system.len().saturating_sub(1))
351    };
352    let protected_user_ids: HashSet<String> = if keep_user_count > 0 {
353        user_indexes[user_indexes.len() - keep_user_count..]
354            .iter()
355            .filter_map(|idx| non_system.get(*idx))
356            .map(|message| message.id.clone())
357            .collect()
358    } else {
359        HashSet::new()
360    };
361
362    tracing::debug!(
363        "compression plan: context_window={}, target_limit={}, system_tokens={}, summary_tokens={}, window_limit={}, non_system_messages={}, keep_user_count={}, keep_from_index={}",
364        context_window, target_limit, system_tokens, summary_tokens, window_limit, non_system.len(), keep_user_count, anchor_index
365    );
366
367    // Keep the newest 3 user turns (or fewer if there are not enough user
368    // turns) as active context and summarize older history before that
369    // boundary. If budget is still too high, continue moving the oldest
370    // non-protected messages into the summarize set.
371    let mut messages_to_summarize = non_system[..anchor_index].to_vec();
372
373    // Protected messages must never be summarized — move them to the keep set.
374    let mut never_compress_ids: Vec<String> = messages_to_summarize
375        .iter()
376        .filter(|m| m.never_compress || is_skill_tool_chain_message(m))
377        .map(|m| m.id.clone())
378        .collect();
379
380    // Also protect tool result messages that correspond to skill tool calls.
381    let skill_call_ids: Vec<String> = messages_to_summarize
382        .iter()
383        .filter(|m| is_skill_tool_chain_message(m))
384        .flat_map(|m| m.tool_calls.iter().flatten().map(|c| c.id.clone()))
385        .collect();
386    if !skill_call_ids.is_empty() {
387        for m in &*messages_to_summarize {
388            if let Some(ref call_id) = m.tool_call_id {
389                if skill_call_ids.contains(call_id) && !never_compress_ids.contains(&m.id) {
390                    never_compress_ids.push(m.id.clone());
391                }
392            }
393        }
394    }
395
396    if !never_compress_ids.is_empty() {
397        messages_to_summarize.retain(|m| !never_compress_ids.contains(&m.id));
398    }
399
400    let non_system_count = non_system.len();
401    let mut messages_to_keep = non_system[anchor_index..].to_vec();
402    // Add never_compress / skill messages to the keep set.
403    for id in &never_compress_ids {
404        if let Some(msg) = non_system.iter().find(|m| &m.id == id) {
405            if !messages_to_keep.iter().any(|m| m.id == *id) {
406                messages_to_keep.push(msg.clone());
407            }
408        }
409    }
410
411    while !messages_to_keep.is_empty() {
412        let keep_tokens = counter.count_messages(&messages_to_keep);
413        if keep_tokens <= window_limit {
414            break;
415        }
416
417        let Some(remove_index) = messages_to_keep.iter().position(|message| {
418            !protected_user_ids.contains(message.id.as_str())
419                && !never_compress_ids.contains(&message.id)
420        }) else {
421            // Remaining messages are all protected; stop shrinking.
422            break;
423        };
424        let moved = messages_to_keep.remove(remove_index);
425        messages_to_summarize.push(moved);
426    }
427
428    // Tool-chain atomicity. The one-at-a-time eviction above can split a generic
429    // (non-skill) tool chain — compressing an assistant `tool_calls` message
430    // while keeping its `tool` result active, or vice versa. That leaves an
431    // orphan `tool_result` (or a `tool_use` with no result) in the active set,
432    // which providers reject with a 400 that then poisons EVERY subsequent
433    // request in the session. Close the compressed set over tool chains: keep
434    // moving any kept message that shares a `tool_call_id` with an
435    // already-compressed message into the summarize set until none remain.
436    // (Skill chains are fully protected above, so are never partially
437    // compressed; protected user messages carry no `tool_call_id`.)
438    close_compressed_set_over_tool_chains(
439        &mut messages_to_keep,
440        &mut messages_to_summarize,
441        &protected_user_ids,
442        &never_compress_ids,
443    );
444
445    if messages_to_summarize.is_empty() {
446        tracing::debug!(
447            "compression plan: messages_to_summarize is empty after anchor/keep splitting"
448        );
449        return Err(CompressionPlanError::NothingToCompress {
450            anchor_index,
451            non_system_count,
452        });
453    }
454
455    let compressed_message_ids = messages_to_summarize
456        .iter()
457        .map(|message| message.id.clone())
458        .collect::<Vec<_>>();
459
460    let keep_tokens = counter.count_messages(&messages_to_keep);
461    let active_before = exposure.active_usage_percent;
462    // Use context_window as denominator, consistent with
463    // estimate_context_compression_exposure().
464    let active_after = if context_window == 0 {
465        0.0
466    } else {
467        let after_total = reserved_non_window_tokens.saturating_add(keep_tokens);
468        (after_total as f64 / context_window as f64) * 100.0
469    };
470
471    // Count actual segments being compressed using the same segmenter that
472    // prepare_hybrid_context uses, so the segment count is accurate.
473    let segmenter = crate::segmenter::MessageSegmenter::new();
474    let segments_removed = segmenter.segment(messages_to_summarize.clone()).len();
475
476    Ok(CompressionPlan {
477        compressed_message_ids,
478        messages_to_summarize,
479        summary_tokens,
480        summary_content,
481        active_usage_before_percent: active_before,
482        active_usage_after_percent: active_after,
483        trigger_percent: budget.compression_trigger_percent,
484        target_percent: budget.compression_target_percent,
485        segments_removed,
486        trigger_type,
487        compression_ratio: 0.0,
488        model_used: None,
489        latency_ms: 0,
490    })
491}
492
493/// Apply a previously computed compression plan to the session.
494/// Extract recently modified files from tool calls in the given messages.
495pub(super) fn extract_recently_modified_files(messages: &[Message]) -> Vec<(String, String)> {
496    let mut files = Vec::new();
497    for message in messages {
498        if let Some(ref tool_calls) = message.tool_calls {
499            for call in tool_calls {
500                let tool_name = call.function.name.as_str();
501                if !matches!(tool_name, "Write" | "Edit" | "Bash") {
502                    continue;
503                }
504                let args = &call.function.arguments;
505                if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(args) {
506                    if let Some(path) = parsed.get("file_path").and_then(|v| v.as_str()) {
507                        files.push((path.to_string(), tool_name.to_string()));
508                    } else if let Some(cmd) = parsed.get("command").and_then(|v| v.as_str()) {
509                        // Extract file paths from shell commands heuristically
510                        for part in cmd.split_whitespace() {
511                            if part.contains('/')
512                                && (part.ends_with(".rs")
513                                    || part.ends_with(".ts")
514                                    || part.ends_with(".js")
515                                    || part.ends_with(".toml")
516                                    || part.ends_with(".json")
517                                    || part.ends_with(".md"))
518                            {
519                                files.push((part.to_string(), "Bash".to_string()));
520                            }
521                        }
522                    }
523                }
524            }
525        }
526    }
527    files.truncate(10);
528    files
529}
530
531/// Extract key decision snippets from assistant messages.
532pub(super) fn extract_key_decisions(messages: &[Message], limit: usize) -> Vec<String> {
533    let decision_keywords = [
534        "decided to",
535        "approach is",
536        "use ",
537        "using ",
538        "we'll go with",
539        "the plan is",
540        "strategy:",
541        "solution:",
542        "chose to",
543        "switched to",
544        "refactored to",
545        "migrated to",
546        "replaced with",
547    ];
548    let mut decisions = Vec::new();
549    for message in messages {
550        if !matches!(message.role, bamboo_domain::Role::Assistant) {
551            continue;
552        }
553        let content = &message.content;
554        for line in content.lines() {
555            let line_lower = line.to_lowercase();
556            if decision_keywords.iter().any(|kw| line_lower.contains(kw)) {
557                let truncated: String = line.chars().take(200).collect();
558                decisions.push(truncated);
559                if decisions.len() >= limit {
560                    return decisions;
561                }
562            }
563        }
564    }
565    decisions
566}
567
568/// Build a post-compaction recovery message that preserves critical context
569/// from the compressed messages so the LLM can continue work without losing
570/// track of active files, tasks, and decisions.
571fn build_post_compaction_recovery_message(
572    compressed_messages: &[Message],
573    session: &Session,
574) -> Option<Message> {
575    if compressed_messages.is_empty() {
576        return None;
577    }
578
579    let mut sections = Vec::new();
580
581    // 1. Recently modified files
582    let files = extract_recently_modified_files(compressed_messages);
583    if !files.is_empty() {
584        let mut section = String::from("## Recently Modified Files\n");
585        for (path, tool) in &files {
586            section.push_str(&format!("- {} ({})\n", path, tool));
587        }
588        sections.push(section);
589    }
590
591    // 2. Active tasks from task list
592    if let Some(ref task_list) = session.task_list {
593        let active_items: Vec<_> = task_list
594            .items
595            .iter()
596            .filter(|item| !matches!(item.status, bamboo_domain::TaskItemStatus::Completed))
597            .collect();
598        if !active_items.is_empty() {
599            let mut section = String::from("## Active Tasks\n");
600            for item in active_items.iter().take(10) {
601                section.push_str(&format!("- [{:?}] {}\n", item.status, item.description));
602            }
603            sections.push(section);
604        }
605    }
606
607    // 3. Key decisions
608    let decisions = extract_key_decisions(compressed_messages, 5);
609    if !decisions.is_empty() {
610        let mut section = String::from("## Key Decisions\n");
611        for decision in &decisions {
612            section.push_str(&format!("- {}\n", decision));
613        }
614        sections.push(section);
615    }
616
617    if sections.is_empty() {
618        return None;
619    }
620
621    let mut content = String::from("[post-compaction-recovery]\nContext extracted from compressed messages for continued work.\n\n");
622    content.push_str(&sections.join("\n"));
623
624    let mut message = Message::assistant(content, None);
625    message.never_compress = true;
626    Some(message)
627}
628
629struct SummaryQualityMetrics {
630    file_coverage: f64,
631    decision_coverage: f64,
632}
633
634fn validate_summary_quality(summary: &str, messages: &[Message]) -> SummaryQualityMetrics {
635    let files = extract_recently_modified_files(messages);
636    let decisions = extract_key_decisions(messages, 10);
637
638    let files_mentioned = files
639        .iter()
640        .filter(|(path, _)| summary.contains(path.as_str()))
641        .count();
642    let file_coverage = if files.is_empty() {
643        1.0
644    } else {
645        files_mentioned as f64 / files.len() as f64
646    };
647
648    let decisions_mentioned = decisions
649        .iter()
650        .filter(|d| {
651            let check_str: String = d.chars().take(50).collect();
652            summary.contains(&check_str)
653        })
654        .count();
655    let decision_coverage = if decisions.is_empty() {
656        1.0
657    } else {
658        decisions_mentioned as f64 / decisions.len() as f64
659    };
660
661    SummaryQualityMetrics {
662        file_coverage,
663        decision_coverage,
664    }
665}
666
667pub fn apply_compression_plan(session: &mut Session, plan: CompressionPlan) -> usize {
668    let compressed_ids: HashSet<&str> = plan
669        .compressed_message_ids
670        .iter()
671        .map(String::as_str)
672        .collect();
673
674    let mut changed_indexes = Vec::new();
675    for (index, message) in session.messages.iter_mut().enumerate() {
676        if message.compressed || !compressed_ids.contains(message.id.as_str()) {
677            continue;
678        }
679        message.compressed = true;
680        changed_indexes.push(index);
681    }
682
683    if changed_indexes.is_empty() {
684        return 0;
685    }
686
687    let event = CompressionEvent::new(
688        changed_indexes.len(),
689        plan.segments_removed,
690        plan.active_usage_before_percent,
691        plan.active_usage_after_percent,
692        plan.summary_tokens,
693        plan.trigger_type,
694        plan.compression_ratio,
695        plan.model_used.clone(),
696        plan.latency_ms,
697    );
698    let event_id = event.id.clone();
699    for index in changed_indexes {
700        session.messages[index].compressed_by_event_id = Some(event_id.clone());
701    }
702    session.compression_events.push(event);
703    session.conversation_summary = Some(ConversationSummary::new(
704        &plan.summary_content,
705        plan.compressed_message_ids.len(),
706        plan.summary_tokens,
707    ));
708
709    // Inject a post-compaction recovery message to preserve critical context
710    // from the compressed messages (files, tasks, decisions).
711    let compressed_messages: Vec<Message> = session
712        .messages
713        .iter()
714        .filter(|m| compressed_ids.contains(m.id.as_str()))
715        .cloned()
716        .collect();
717    if let Some(recovery) = build_post_compaction_recovery_message(&compressed_messages, session) {
718        // Insert just before the last user message, or at the end
719        let insert_pos = session
720            .messages
721            .iter()
722            .rposition(|m| matches!(m.role, bamboo_domain::Role::User) && !m.compressed)
723            .map(|pos| pos + 1)
724            .unwrap_or(session.messages.len());
725        session.messages.insert(insert_pos, recovery);
726    }
727
728    let quality = validate_summary_quality(&plan.summary_content, &compressed_messages);
729    if quality.file_coverage < 0.5 || quality.decision_coverage < 0.3 {
730        tracing::warn!(
731            "[{}] Summary quality: file_coverage={:.0}%, decision_coverage={:.0}%",
732            session.id,
733            quality.file_coverage * 100.0,
734            quality.decision_coverage * 100.0
735        );
736    }
737
738    // Instead of clearing token_usage entirely (which forces the next round
739    // to rely on heuristic estimates that don't account for tool schema
740    // tokens), recompute an approximate post-compression snapshot.  We
741    // preserve the context-window denominator from the previous usage snapshot
742    // so percentages stay consistent across rounds.
743    let counter = TiktokenTokenCounter::default();
744    let remaining_active: Vec<_> = session
745        .messages
746        .iter()
747        .filter(|m| !m.compressed)
748        .cloned()
749        .collect();
750    let system_msgs: Vec<_> = remaining_active
751        .iter()
752        .filter(|m| matches!(m.role, bamboo_domain::Role::System))
753        .cloned()
754        .collect();
755    let window_msgs: Vec<_> = remaining_active
756        .iter()
757        .filter(|m| !matches!(m.role, bamboo_domain::Role::System))
758        .cloned()
759        .collect();
760    let system_tokens = counter.count_messages(&system_msgs);
761    let new_summary_tokens = plan.summary_tokens;
762    let window_tokens = counter.count_messages(&window_msgs);
763    let total_tokens = system_tokens
764        .saturating_add(new_summary_tokens)
765        .saturating_add(window_tokens);
766    let previous_usage = session.token_usage.take();
767    let budget_limit = previous_usage
768        .as_ref()
769        .map(|u| {
770            if u.max_context_tokens > 0 {
771                u.max_context_tokens
772            } else {
773                u.budget_limit
774            }
775        })
776        .unwrap_or(0);
777    let max_context_tokens = previous_usage
778        .as_ref()
779        .map(|u| u.max_context_tokens)
780        .unwrap_or(0);
781    session.token_usage = Some(bamboo_domain::TokenBudgetUsage {
782        system_tokens,
783        summary_tokens: new_summary_tokens,
784        window_tokens,
785        total_tokens,
786        max_context_tokens,
787        budget_limit,
788        truncation_occurred: false,
789        segments_removed: 0,
790        prompt_cached_tool_outputs: 0,
791        prompt_cached_tool_tokens_saved: 0,
792        thinking_tokens: 0,
793        cache_read_input_tokens: 0,
794    });
795
796    session.updated_at = Utc::now();
797    plan.compressed_message_ids.len()
798}
799
800pub fn compression_summary_message(summary_content: &str) -> Message {
801    Message::system(format!(
802        "<!-- CONVERSATION_SUMMARY_START -->\n\
803         ## Previous Conversation Summary\n\
804         The following is compressed historical context for continuity only.\n\
805         It is background memory, not a new user request. Follow the current task list and recent messages over this summary when they conflict.\n\n\
806         {}\n\
807         <!-- CONVERSATION_SUMMARY_END -->",
808        summary_content
809    ))
810}
811
812pub fn active_messages_for_budget(session: &Session) -> Vec<Message> {
813    session
814        .messages
815        .iter()
816        .filter(|message| !message.compressed)
817        .cloned()
818        .collect()
819}
820
821pub fn summary_source_messages(session: &Session) -> Vec<Message> {
822    session
823        .messages
824        .iter()
825        .filter(|message| !message.compressed)
826        .filter(|message| !matches!(message.role, bamboo_domain::Role::System))
827        .cloned()
828        .collect()
829}
830
831pub fn build_summary_prompt(
832    session: &Session,
833    messages: &[Message],
834    existing_summary: Option<&str>,
835) -> String {
836    let mut content = String::new();
837    content.push_str(
838        "You are compressing conversation history for continued work. Produce a compact but reliable working-memory summary.\n\n",
839    );
840    content.push_str(
841        "Critical requirements:\n- First capture the in-flight work right before compression (what was being done, where, and with which tool/file)\n- Distinguish clearly between ACTIVE work, COMPLETED work, and OBSOLETE or superseded work\n- Do not restate old tasks as active unless they are still unresolved\n- The current task list is the source of truth for what is actively being worked on\n- Preserve constraints, decisions, file paths, code changes, errors, tool findings, blockers, and the next step\n- If earlier plans conflict with the current task list or newer messages, treat the earlier plans as obsolete or completed\n- Explicitly evaluate each clear user requirement (e.g. requirement 1, requirement 2) with a status and evidence\n- Return only summary text in the same language as the conversation\n\n",
842    );
843
844    if let Some(existing) = existing_summary.map(str::trim).filter(|s| !s.is_empty()) {
845        content.push_str("## Existing Summary\n");
846        content.push_str(existing);
847        content.push_str("\n\n");
848    }
849
850    let task_list_prompt = session.format_task_list_for_prompt();
851    if !task_list_prompt.trim().is_empty() {
852        content.push_str("## Current Task List\n");
853        content.push_str(task_list_prompt.trim());
854        content.push_str("\n\n");
855    }
856
857    content.push_str(
858        "## Required Output Sections\n1. Pre-compression in-flight work (what was being done immediately before compression)\n2. Current active objective\n3. Requirement checklist (Requirement | Status: completed/in_progress/pending/blocked/obsolete | Evidence)\n4. Active tasks\n5. Completed tasks\n6. Obsolete or superseded tasks\n7. Important context and constraints\n8. Files, code, and tool findings\n9. Open issues and next step\n\n",
859    );
860
861    content.push_str("## Messages To Compress\n\n");
862    for message in messages {
863        let role = match message.role {
864            bamboo_domain::Role::System => continue,
865            bamboo_domain::Role::User => "User",
866            bamboo_domain::Role::Assistant => match message.phase {
867                Some(MessagePhase::Commentary) => "Assistant Commentary",
868                Some(MessagePhase::FinalAnswer) => "Assistant Final",
869                None => "Assistant",
870            },
871            bamboo_domain::Role::Tool => "Tool Result",
872        };
873
874        content.push_str("### ");
875        content.push_str(role);
876        content.push('\n');
877        if let Some(tool_calls) = &message.tool_calls {
878            if !tool_calls.is_empty() {
879                let names = tool_calls
880                    .iter()
881                    .map(|call| call.function.name.as_str())
882                    .collect::<Vec<_>>()
883                    .join(", ");
884                content.push_str("Called tools: ");
885                content.push_str(&names);
886                content.push('\n');
887            }
888        }
889        if let Some(tool_call_id) = &message.tool_call_id {
890            content.push_str("Tool call id: ");
891            content.push_str(tool_call_id);
892            content.push('\n');
893        }
894        let snippet = truncate_chars(&message.content, 2000);
895        content.push_str(&snippet);
896        content.push_str("\n\n");
897    }
898
899    content.push_str(
900        "Return only the summary text. Be explicit about what is active now versus what is already done or no longer relevant.",
901    );
902    content
903}
904
905fn truncate_chars(value: &str, max_chars: usize) -> String {
906    if value.chars().count() <= max_chars {
907        return value.to_string();
908    }
909    value.chars().take(max_chars).collect::<String>() + "..."
910}
911
912#[cfg(test)]
913mod tests {
914    use super::*;
915    use bamboo_domain::TokenBudgetUsage;
916    use bamboo_domain::{FunctionCall, TaskItem, TaskItemStatus, TaskList, ToolCall};
917    use chrono::Utc;
918
919    fn make_budget() -> TokenBudget {
920        TokenBudget {
921            max_context_tokens: 1000,
922            max_output_tokens: 100,
923            strategy: BudgetStrategy::Hybrid {
924                window_size: 20,
925                enable_summarization: true,
926            },
927            safety_margin: 0,
928            compression_trigger_percent: 50,
929            compression_target_percent: 20,
930            working_reserve_tokens: 0,
931            fallback_trigger_percent: 75,
932            prompt_cache_min_tool_output_chars: 1_200,
933            prompt_cache_head_chars: 280,
934            prompt_cache_tail_chars: 180,
935            prompt_cache_recent_user_turns: 2,
936            prompt_cache_recent_tool_chains: 2,
937            max_tool_output_tokens: 0,
938        }
939    }
940
941    fn make_session_with_pressure() -> Session {
942        let mut session = Session::new("compression-hysteresis", "gpt-4o-mini");
943        session.token_budget = Some(make_budget());
944        session.add_message(Message::system("system"));
945        for i in 0..3 {
946            session.add_message(Message::user(format!(
947                "User message {i}: {}",
948                "alpha beta gamma delta epsilon ".repeat(2)
949            )));
950            session.add_message(Message::assistant(
951                format!(
952                    "Assistant message {i}: {}",
953                    "work log decisions next steps ".repeat(2)
954                ),
955                None,
956            ));
957        }
958        session
959    }
960
961    #[test]
962    fn context_window_usage_percent_uses_context_window_denominator() {
963        assert_eq!(context_window_usage_percent(0, 0), 0.0);
964        assert_eq!(context_window_usage_percent(500, 1000), 50.0);
965    }
966
967    #[test]
968    fn estimate_context_compression_exposure_crosses_trigger_when_usage_is_high_enough() {
969        let mut session = make_session_with_pressure();
970        if let Some(budget) = session.token_budget.as_mut() {
971            budget.compression_trigger_percent = 10;
972        }
973        let exposure = estimate_context_compression_exposure(
974            &session,
975            "gpt-4o-mini",
976            session.token_budget.as_ref(),
977        );
978        assert!(exposure.active_usage_percent >= 10.0);
979        assert!(exposure.should_expose_tool);
980    }
981
982    #[test]
983    fn estimate_context_compression_exposure_stays_below_trigger_when_usage_is_low() {
984        let mut session = make_session_with_pressure();
985        if let Some(budget) = session.token_budget.as_mut() {
986            budget.compression_trigger_percent = 99;
987        }
988
989        let exposure = estimate_context_compression_exposure(
990            &session,
991            "gpt-4o-mini",
992            session.token_budget.as_ref(),
993        );
994
995        assert!(exposure.active_usage_percent < 99.0);
996        assert!(!exposure.should_expose_tool);
997    }
998
999    #[test]
1000    fn build_summary_prompt_includes_task_list_and_state_sections() {
1001        let mut session = Session::new("summary-prompt", "gpt-4o-mini");
1002        session.set_task_list(TaskList {
1003            session_id: session.id.clone(),
1004            title: "Task List".to_string(),
1005            items: vec![
1006                TaskItem {
1007                    id: "task_1".to_string(),
1008                    description: "检查 51% 又回落到 50% 的触发逻辑".to_string(),
1009                    status: TaskItemStatus::InProgress,
1010                    depends_on: Vec::new(),
1011                    notes: "避免刚压缩完又立刻再次压缩".to_string(),
1012                    ..TaskItem::default()
1013                },
1014                TaskItem {
1015                    id: "task_2".to_string(),
1016                    description: "重写 summarizer prompt 并纳入 task list".to_string(),
1017                    status: TaskItemStatus::Pending,
1018                    depends_on: Vec::new(),
1019                    notes: String::new(),
1020                    ..TaskItem::default()
1021                },
1022            ],
1023            created_at: Utc::now(),
1024            updated_at: Utc::now(),
1025        });
1026        let prompt = build_summary_prompt(
1027            &session,
1028            &[
1029                Message::user("继续修复 context compression"),
1030                Message::assistant("先分析 trigger / target / summary", None),
1031            ],
1032            Some("old summary"),
1033        );
1034
1035        assert!(prompt.contains("## Current Task List"));
1036        assert!(prompt.contains("Current active objective"));
1037        assert!(prompt.contains("Requirement checklist"));
1038        assert!(prompt.contains("Active tasks"));
1039        assert!(prompt.contains("Completed tasks"));
1040        assert!(prompt.contains("Obsolete or superseded tasks"));
1041        assert!(prompt.contains("检查 51% 又回落到 50% 的触发逻辑"));
1042        assert!(prompt.contains("old summary"));
1043    }
1044
1045    #[test]
1046    fn forced_plan_keeps_last_three_user_messages_active() {
1047        let budget = TokenBudget {
1048            max_context_tokens: 1200,
1049            max_output_tokens: 100,
1050            strategy: BudgetStrategy::Hybrid {
1051                window_size: 20,
1052                enable_summarization: true,
1053            },
1054            safety_margin: 0,
1055            compression_trigger_percent: 80,
1056            compression_target_percent: 20,
1057            working_reserve_tokens: 0,
1058            fallback_trigger_percent: 75,
1059            prompt_cache_min_tool_output_chars: 1_200,
1060            prompt_cache_head_chars: 280,
1061            prompt_cache_tail_chars: 180,
1062            prompt_cache_recent_user_turns: 2,
1063            prompt_cache_recent_tool_chains: 2,
1064            max_tool_output_tokens: 0,
1065        };
1066        let mut session = Session::new("keep-last-three-user-turns", "gpt-4o-mini");
1067        session.token_budget = Some(budget.clone());
1068        session.add_message(Message::system("system"));
1069        for i in 0..6 {
1070            session.add_message(Message::user(format!(
1071                "U{i}: {}",
1072                "alpha beta gamma ".repeat(8)
1073            )));
1074            session.add_message(Message::assistant(
1075                format!("A{i}: {}", "analysis plan steps ".repeat(8)),
1076                None,
1077            ));
1078        }
1079
1080        let plan = build_forced_compression_plan_with_summary(
1081            &session,
1082            "gpt-4o-mini",
1083            Some(&budget),
1084            "summary".to_string(),
1085            CompressionTriggerType::CriticalOverflow,
1086        )
1087        .expect("forced plan should build");
1088
1089        let compressed_ids = plan
1090            .compressed_message_ids
1091            .iter()
1092            .map(String::as_str)
1093            .collect::<HashSet<_>>();
1094        let kept_user_contents = session
1095            .messages
1096            .iter()
1097            .filter(|message| !matches!(message.role, bamboo_domain::Role::System))
1098            .filter(|message| !compressed_ids.contains(message.id.as_str()))
1099            .filter(|message| matches!(message.role, bamboo_domain::Role::User))
1100            .map(|message| message.content.clone())
1101            .collect::<Vec<_>>();
1102
1103        assert!(
1104            kept_user_contents.len() >= 3,
1105            "expected to keep at least 3 user messages, got {}",
1106            kept_user_contents.len()
1107        );
1108        assert!(kept_user_contents
1109            .iter()
1110            .any(|content| content.starts_with("U3:")));
1111        assert!(kept_user_contents
1112            .iter()
1113            .any(|content| content.starts_with("U4:")));
1114        assert!(kept_user_contents
1115            .iter()
1116            .any(|content| content.starts_with("U5:")));
1117    }
1118
1119    #[test]
1120    fn estimate_exposure_prefers_persisted_budget_usage_when_higher() {
1121        let mut session = Session::new("persisted-usage", "gpt-4o-mini");
1122        session.token_budget = Some(TokenBudget {
1123            max_context_tokens: 100_000,
1124            max_output_tokens: 1_000,
1125            strategy: BudgetStrategy::Hybrid {
1126                window_size: 20,
1127                enable_summarization: true,
1128            },
1129            safety_margin: 0,
1130            compression_trigger_percent: 80,
1131            compression_target_percent: 50,
1132            working_reserve_tokens: 0,
1133            fallback_trigger_percent: 75,
1134            prompt_cache_min_tool_output_chars: 1_200,
1135            prompt_cache_head_chars: 280,
1136            prompt_cache_tail_chars: 180,
1137            prompt_cache_recent_user_turns: 2,
1138            prompt_cache_recent_tool_chains: 2,
1139            max_tool_output_tokens: 0,
1140        });
1141        session.add_message(Message::system("system"));
1142        session.add_message(Message::user("short"));
1143        session.add_message(Message::assistant("short", None));
1144        session.add_message(Message::user("follow-up"));
1145        session.add_message(Message::assistant("reply", None));
1146        session.token_usage = Some(TokenBudgetUsage {
1147            system_tokens: 100,
1148            summary_tokens: 0,
1149            window_tokens: 95_900,
1150            total_tokens: 96_000,
1151            max_context_tokens: 100_000,
1152            budget_limit: 10_000,
1153            truncation_occurred: true,
1154            segments_removed: 12,
1155            prompt_cached_tool_outputs: 0,
1156            prompt_cached_tool_tokens_saved: 0,
1157            thinking_tokens: 0,
1158            cache_read_input_tokens: 0,
1159        });
1160
1161        let exposure = estimate_context_compression_exposure(
1162            &session,
1163            "gpt-4o-mini",
1164            session.token_budget.as_ref(),
1165        );
1166
1167        assert!(
1168            exposure.active_usage_percent >= 96.0,
1169            "expected persisted context-window usage to drive exposure, got {}",
1170            exposure.active_usage_percent
1171        );
1172        assert!(exposure.should_expose_tool);
1173    }
1174
1175    #[test]
1176    fn never_compress_messages_are_excluded_from_summarize_set() {
1177        let budget = TokenBudget {
1178            max_context_tokens: 1200,
1179            max_output_tokens: 100,
1180            strategy: BudgetStrategy::Hybrid {
1181                window_size: 20,
1182                enable_summarization: true,
1183            },
1184            safety_margin: 0,
1185            compression_trigger_percent: 80,
1186            compression_target_percent: 20,
1187            working_reserve_tokens: 0,
1188            fallback_trigger_percent: 75,
1189            prompt_cache_min_tool_output_chars: 1_200,
1190            prompt_cache_head_chars: 280,
1191            prompt_cache_tail_chars: 180,
1192            prompt_cache_recent_user_turns: 2,
1193            prompt_cache_recent_tool_chains: 2,
1194            max_tool_output_tokens: 0,
1195        };
1196        let mut session = Session::new("never-compress-test", "gpt-4o-mini");
1197        session.token_budget = Some(budget.clone());
1198        session.add_message(Message::system("system"));
1199
1200        // Old user message that should be summarized
1201        session.add_message(Message::user("Old question about X"));
1202        session.add_message(Message::assistant("Old answer about X", None));
1203
1204        // Protected user message (never_compress = true)
1205        let mut protected = Message::user("Critical context that must survive");
1206        protected.never_compress = true;
1207        session.add_message(protected);
1208        session.add_message(Message::assistant("Response to critical", None));
1209
1210        // Recent user messages that anchor the keep window
1211        for i in 0..4 {
1212            session.add_message(Message::user(format!(
1213                "Recent U{i}: {}",
1214                "padding text to fill budget ".repeat(6)
1215            )));
1216            session.add_message(Message::assistant(
1217                format!("Recent A{i}: {}", "reply padding text ".repeat(6)),
1218                None,
1219            ));
1220        }
1221
1222        let plan = build_forced_compression_plan_with_summary(
1223            &session,
1224            "gpt-4o-mini",
1225            Some(&budget),
1226            "summary".to_string(),
1227            CompressionTriggerType::Auto,
1228        )
1229        .expect("plan should build");
1230
1231        let compressed_ids: HashSet<&str> = plan
1232            .compressed_message_ids
1233            .iter()
1234            .map(String::as_str)
1235            .collect();
1236
1237        // Find the never_compress message
1238        let protected_msg = session
1239            .messages
1240            .iter()
1241            .find(|m| m.never_compress)
1242            .expect("should find the protected message");
1243
1244        assert!(
1245            !compressed_ids.contains(protected_msg.id.as_str()),
1246            "never_compress message should NOT be in the compressed set"
1247        );
1248    }
1249
1250    #[test]
1251    fn skill_tool_chain_messages_are_protected_from_compression() {
1252        let budget = TokenBudget {
1253            max_context_tokens: 1200,
1254            max_output_tokens: 100,
1255            strategy: BudgetStrategy::Hybrid {
1256                window_size: 20,
1257                enable_summarization: true,
1258            },
1259            safety_margin: 0,
1260            compression_trigger_percent: 80,
1261            compression_target_percent: 20,
1262            working_reserve_tokens: 0,
1263            fallback_trigger_percent: 75,
1264            prompt_cache_min_tool_output_chars: 1_200,
1265            prompt_cache_head_chars: 280,
1266            prompt_cache_tail_chars: 180,
1267            prompt_cache_recent_user_turns: 2,
1268            prompt_cache_recent_tool_chains: 2,
1269            max_tool_output_tokens: 0,
1270        };
1271        let mut session = Session::new("skill-chain-test", "gpt-4o-mini");
1272        session.token_budget = Some(budget.clone());
1273        session.add_message(Message::system("system"));
1274
1275        // Skill tool chain (load_skill + read_skill_resource)
1276        let mut skill_call = Message::assistant(String::new(), None);
1277        skill_call.tool_calls = Some(vec![ToolCall {
1278            id: "tc-skill".to_string(),
1279            tool_type: "function".to_string(),
1280            function: FunctionCall {
1281                name: "load_skill".to_string(),
1282                arguments: r#"{"skill_id":"my-skill"}"#.to_string(),
1283            },
1284        }]);
1285        session.add_message(skill_call);
1286
1287        let mut skill_result = Message::tool_result("tc-skill", "skill loaded");
1288        skill_result.tool_success = Some(true);
1289        session.add_message(skill_result);
1290
1291        // Regular messages to fill budget
1292        for i in 0..6 {
1293            session.add_message(Message::user(format!(
1294                "U{i}: {}",
1295                "alpha beta gamma delta ".repeat(8)
1296            )));
1297            session.add_message(Message::assistant(
1298                format!("A{i}: {}", "analysis steps plan ".repeat(8)),
1299                None,
1300            ));
1301        }
1302
1303        let plan = build_forced_compression_plan_with_summary(
1304            &session,
1305            "gpt-4o-mini",
1306            Some(&budget),
1307            "summary".to_string(),
1308            CompressionTriggerType::Auto,
1309        )
1310        .expect("plan should build");
1311
1312        let compressed_ids: HashSet<&str> = plan
1313            .compressed_message_ids
1314            .iter()
1315            .map(String::as_str)
1316            .collect();
1317
1318        // Skill tool chain messages should not be compressed
1319        let skill_messages: Vec<&Message> = session
1320            .messages
1321            .iter()
1322            .filter(|m| {
1323                m.tool_calls
1324                    .as_ref()
1325                    .is_some_and(|calls| calls.iter().any(|c| c.function.name == "load_skill"))
1326                    || m.tool_call_id.as_deref() == Some("tc-skill")
1327            })
1328            .collect();
1329
1330        for msg in &skill_messages {
1331            assert!(
1332                !compressed_ids.contains(msg.id.as_str()),
1333                "skill tool chain message {} should NOT be compressed",
1334                msg.id
1335            );
1336        }
1337    }
1338
1339    #[test]
1340    fn generic_tool_chain_is_never_split_by_forced_compression() {
1341        // A generic (non-skill) tool_use and its tool_result must never be split
1342        // across the compression boundary — a split orphans one of them in the
1343        // active set and the provider 400s, poisoning the session. #340.
1344        let budget = TokenBudget {
1345            max_context_tokens: 1200,
1346            max_output_tokens: 100,
1347            strategy: BudgetStrategy::Hybrid {
1348                window_size: 20,
1349                enable_summarization: true,
1350            },
1351            safety_margin: 0,
1352            compression_trigger_percent: 80,
1353            compression_target_percent: 20,
1354            working_reserve_tokens: 0,
1355            fallback_trigger_percent: 75,
1356            prompt_cache_min_tool_output_chars: 1_200,
1357            prompt_cache_head_chars: 280,
1358            prompt_cache_tail_chars: 180,
1359            prompt_cache_recent_user_turns: 2,
1360            prompt_cache_recent_tool_chains: 2,
1361            max_tool_output_tokens: 0,
1362        };
1363        let mut session = Session::new("generic-chain-test", "gpt-4o-mini");
1364        session.token_budget = Some(budget.clone());
1365        session.add_message(Message::system("system"));
1366
1367        // A generic tool chain (search + its result) placed early so it is an
1368        // eviction candidate; the large result makes it a prime compression target.
1369        let mut call = Message::assistant(String::new(), None);
1370        call.tool_calls = Some(vec![ToolCall {
1371            id: "tc-gen".to_string(),
1372            tool_type: "function".to_string(),
1373            function: FunctionCall {
1374                name: "search".to_string(),
1375                arguments: r#"{"q":"rust"}"#.to_string(),
1376            },
1377        }]);
1378        session.add_message(call);
1379        let mut result = Message::tool_result("tc-gen", &"search result payload ".repeat(20));
1380        result.tool_success = Some(true);
1381        session.add_message(result);
1382
1383        // Filler to push usage over budget and force eviction.
1384        for i in 0..8 {
1385            session.add_message(Message::user(format!(
1386                "U{i}: {}",
1387                "alpha beta gamma delta ".repeat(8)
1388            )));
1389            session.add_message(Message::assistant(
1390                format!("A{i}: {}", "analysis steps plan ".repeat(8)),
1391                None,
1392            ));
1393        }
1394
1395        let plan = build_forced_compression_plan_with_summary(
1396            &session,
1397            "gpt-4o-mini",
1398            Some(&budget),
1399            "summary".to_string(),
1400            CompressionTriggerType::Auto,
1401        )
1402        .expect("plan should build");
1403
1404        let compressed: HashSet<&str> = plan
1405            .compressed_message_ids
1406            .iter()
1407            .map(String::as_str)
1408            .collect();
1409
1410        let assistant_id = session
1411            .messages
1412            .iter()
1413            .find(|m| {
1414                m.tool_calls
1415                    .as_ref()
1416                    .is_some_and(|c| c.iter().any(|tc| tc.id == "tc-gen"))
1417            })
1418            .map(|m| m.id.as_str())
1419            .expect("assistant tool_use message present");
1420        let result_id = session
1421            .messages
1422            .iter()
1423            .find(|m| m.tool_call_id.as_deref() == Some("tc-gen"))
1424            .map(|m| m.id.as_str())
1425            .expect("tool_result message present");
1426
1427        // The tool_use and its tool_result must land on the SAME side.
1428        assert_eq!(
1429            compressed.contains(assistant_id),
1430            compressed.contains(result_id),
1431            "generic tool_use ({assistant_id}) and its tool_result ({result_id}) must not be split"
1432        );
1433    }
1434
1435    fn tool_use_message(call_id: &str) -> Message {
1436        let mut assistant = Message::assistant(String::new(), None);
1437        assistant.tool_calls = Some(vec![ToolCall {
1438            id: call_id.to_string(),
1439            tool_type: "function".to_string(),
1440            function: FunctionCall {
1441                name: "search".to_string(),
1442                arguments: "{}".to_string(),
1443            },
1444        }]);
1445        assistant
1446    }
1447
1448    #[test]
1449    fn close_compressed_set_over_tool_chains_reunites_a_split_chain() {
1450        // Pre-split state: the assistant tool_use is compressed while its
1451        // tool_result was left in the keep (active) set — an orphan. #340.
1452        let assistant = tool_use_message("tc-1");
1453        let result = Message::tool_result("tc-1", "result payload");
1454        let assistant_id = assistant.id.clone();
1455        let result_id = result.id.clone();
1456
1457        let mut messages_to_keep = vec![result];
1458        let mut messages_to_summarize = vec![assistant];
1459
1460        close_compressed_set_over_tool_chains(
1461            &mut messages_to_keep,
1462            &mut messages_to_summarize,
1463            &HashSet::new(),
1464            &[],
1465        );
1466
1467        // The orphan tool_result must be pulled into the compressed set.
1468        assert!(
1469            messages_to_keep.is_empty(),
1470            "orphaned tool_result must be moved into the compressed set"
1471        );
1472        let summarized: HashSet<&str> = messages_to_summarize
1473            .iter()
1474            .map(|m| m.id.as_str())
1475            .collect();
1476        assert!(summarized.contains(assistant_id.as_str()));
1477        assert!(summarized.contains(result_id.as_str()));
1478    }
1479
1480    #[test]
1481    fn close_compressed_set_over_tool_chains_respects_protected_messages() {
1482        // A protected (never-compress) chain member must NOT be force-compressed.
1483        let assistant = tool_use_message("tc-2");
1484        let result = Message::tool_result("tc-2", "result");
1485        let result_id = result.id.clone();
1486
1487        let mut messages_to_keep = vec![result];
1488        let mut messages_to_summarize = vec![assistant];
1489        let never_compress_ids = vec![result_id.clone()];
1490
1491        close_compressed_set_over_tool_chains(
1492            &mut messages_to_keep,
1493            &mut messages_to_summarize,
1494            &HashSet::new(),
1495            &never_compress_ids,
1496        );
1497
1498        assert_eq!(
1499            messages_to_keep.len(),
1500            1,
1501            "protected result must stay in keep"
1502        );
1503        assert_eq!(messages_to_keep[0].id, result_id);
1504    }
1505
1506    #[test]
1507    fn recovery_message_returns_none_for_empty_messages() {
1508        let session = Session::new("recovery-empty", "model");
1509        let result = build_post_compaction_recovery_message(&[], &session);
1510        assert!(result.is_none());
1511    }
1512
1513    #[test]
1514    fn recovery_message_has_never_compress_flag() {
1515        let mut session = Session::new("recovery-flag", "model");
1516        let messages = vec![Message::assistant("no decisions here", None)];
1517        session.set_task_list(TaskList {
1518            session_id: session.id.clone(),
1519            title: "Tasks".to_string(),
1520            items: vec![TaskItem {
1521                id: "t1".to_string(),
1522                description: "Active task".to_string(),
1523                status: TaskItemStatus::InProgress,
1524                ..TaskItem::default()
1525            }],
1526            created_at: Utc::now(),
1527            updated_at: Utc::now(),
1528        });
1529        let recovery = build_post_compaction_recovery_message(&messages, &session)
1530            .expect("should return recovery message");
1531        assert!(recovery.never_compress);
1532        assert!(recovery.content.contains("[post-compaction-recovery]"));
1533    }
1534
1535    #[test]
1536    fn recovery_message_extracts_file_paths_from_tool_calls() {
1537        let session = Session::new("recovery-files", "model");
1538        let mut write_call = Message::assistant("writing file", None);
1539        write_call.tool_calls = Some(vec![ToolCall {
1540            id: "tc1".to_string(),
1541            tool_type: "function".to_string(),
1542            function: FunctionCall {
1543                name: "Write".to_string(),
1544                arguments: r#"{"file_path":"/src/main.rs","content":"fn main() {}"}"#.to_string(),
1545            },
1546        }]);
1547        let mut edit_call = Message::assistant("editing file", None);
1548        edit_call.tool_calls = Some(vec![ToolCall {
1549            id: "tc2".to_string(),
1550            tool_type: "function".to_string(),
1551            function: FunctionCall {
1552                name: "Edit".to_string(),
1553                arguments: r#"{"file_path":"/lib/utils.rs","old":"x","new":"y"}"#.to_string(),
1554            },
1555        }]);
1556        let messages = vec![write_call, edit_call];
1557
1558        let recovery = build_post_compaction_recovery_message(&messages, &session)
1559            .expect("should return recovery");
1560        assert!(recovery.content.contains("/src/main.rs"));
1561        assert!(recovery.content.contains("/lib/utils.rs"));
1562        assert!(recovery.content.contains("Recently Modified Files"));
1563    }
1564
1565    #[test]
1566    fn recovery_message_includes_active_tasks() {
1567        let mut session = Session::new("recovery-tasks", "model");
1568        session.set_task_list(TaskList {
1569            session_id: session.id.clone(),
1570            title: "Tasks".to_string(),
1571            items: vec![
1572                TaskItem {
1573                    id: "t1".to_string(),
1574                    description: "Fix auth middleware".to_string(),
1575                    status: TaskItemStatus::InProgress,
1576                    ..TaskItem::default()
1577                },
1578                TaskItem {
1579                    id: "t2".to_string(),
1580                    description: "Add tests".to_string(),
1581                    status: TaskItemStatus::Pending,
1582                    ..TaskItem::default()
1583                },
1584                TaskItem {
1585                    id: "t3".to_string(),
1586                    description: "Done task".to_string(),
1587                    status: TaskItemStatus::Completed,
1588                    ..TaskItem::default()
1589                },
1590            ],
1591            created_at: Utc::now(),
1592            updated_at: Utc::now(),
1593        });
1594        let messages = vec![Message::assistant("some work", None)];
1595
1596        let recovery = build_post_compaction_recovery_message(&messages, &session)
1597            .expect("should return recovery");
1598        assert!(recovery.content.contains("Active Tasks"));
1599        assert!(recovery.content.contains("Fix auth middleware"));
1600        assert!(recovery.content.contains("Add tests"));
1601        // Completed tasks should NOT appear in active tasks
1602        assert!(!recovery.content.contains("Done task"));
1603    }
1604
1605    #[test]
1606    fn apply_compression_plan_injects_recovery_message() {
1607        let budget = TokenBudget {
1608            max_context_tokens: 1200,
1609            max_output_tokens: 100,
1610            strategy: BudgetStrategy::Hybrid {
1611                window_size: 20,
1612                enable_summarization: true,
1613            },
1614            safety_margin: 0,
1615            compression_trigger_percent: 80,
1616            compression_target_percent: 20,
1617            working_reserve_tokens: 0,
1618            fallback_trigger_percent: 75,
1619            prompt_cache_min_tool_output_chars: 1_200,
1620            prompt_cache_head_chars: 280,
1621            prompt_cache_tail_chars: 180,
1622            prompt_cache_recent_user_turns: 2,
1623            prompt_cache_recent_tool_chains: 2,
1624            max_tool_output_tokens: 0,
1625        };
1626        let mut session = Session::new("recovery-inject", "gpt-4o-mini");
1627        session.token_budget = Some(budget.clone());
1628        session.add_message(Message::system("system"));
1629
1630        // Old messages with tool calls containing file paths
1631        let mut write_msg = Message::assistant("writing", None);
1632        write_msg.tool_calls = Some(vec![ToolCall {
1633            id: "tc-w".to_string(),
1634            tool_type: "function".to_string(),
1635            function: FunctionCall {
1636                name: "Write".to_string(),
1637                arguments: r#"{"file_path":"/src/lib.rs","content":"pub fn hello() {}"}"#
1638                    .to_string(),
1639            },
1640        }]);
1641        session.add_message(Message::user("Write the file"));
1642        session.add_message(write_msg);
1643
1644        // Fill with enough messages to force compression
1645        for i in 0..6 {
1646            session.add_message(Message::user(format!(
1647                "U{i}: {}",
1648                "alpha beta gamma delta ".repeat(8)
1649            )));
1650            session.add_message(Message::assistant(
1651                format!("A{i}: {}", "analysis plan ".repeat(8)),
1652                None,
1653            ));
1654        }
1655
1656        let plan = build_forced_compression_plan_with_summary(
1657            &session,
1658            "gpt-4o-mini",
1659            Some(&budget),
1660            "summary text".to_string(),
1661            CompressionTriggerType::Auto,
1662        )
1663        .expect("plan should build");
1664
1665        assert!(!plan.compressed_message_ids.is_empty());
1666
1667        let compressed_count = apply_compression_plan(&mut session, plan);
1668        assert!(compressed_count > 0);
1669
1670        // Verify recovery message was injected
1671        let has_recovery = session.messages.iter().any(|m| {
1672            m.never_compress
1673                && m.content.contains("[post-compaction-recovery]")
1674                && m.content.contains("/src/lib.rs")
1675        });
1676        assert!(
1677            has_recovery,
1678            "session should contain a post-compaction recovery message with the file path"
1679        );
1680    }
1681
1682    #[test]
1683    fn summary_quality_full_coverage_when_all_files_mentioned() {
1684        let messages = vec![{
1685            let mut m = Message::assistant("writing", None);
1686            m.tool_calls = Some(vec![ToolCall {
1687                id: "tc1".to_string(),
1688                tool_type: "function".to_string(),
1689                function: FunctionCall {
1690                    name: "Write".to_string(),
1691                    arguments: r#"{"file_path":"/src/main.rs","content":"fn main() {}"}"#
1692                        .to_string(),
1693                },
1694            }]);
1695            m
1696        }];
1697        let summary = "Modified /src/main.rs to add main function";
1698        let quality = validate_summary_quality(summary, &messages);
1699        assert!(
1700            quality.file_coverage >= 0.99,
1701            "file_coverage should be ~1.0, got {:.2}",
1702            quality.file_coverage
1703        );
1704    }
1705
1706    #[test]
1707    fn summary_quality_zero_coverage_when_no_files_mentioned() {
1708        let messages = vec![{
1709            let mut m = Message::assistant("writing", None);
1710            m.tool_calls = Some(vec![ToolCall {
1711                id: "tc1".to_string(),
1712                tool_type: "function".to_string(),
1713                function: FunctionCall {
1714                    name: "Write".to_string(),
1715                    arguments: r#"{"file_path":"/src/main.rs","content":"fn main() {}"}"#
1716                        .to_string(),
1717                },
1718            }]);
1719            m
1720        }];
1721        let summary = "Summary that mentions nothing about files";
1722        let quality = validate_summary_quality(summary, &messages);
1723        assert!(
1724            quality.file_coverage < 0.01,
1725            "file_coverage should be ~0.0, got {:.2}",
1726            quality.file_coverage
1727        );
1728    }
1729
1730    #[test]
1731    fn summary_quality_handles_empty_messages() {
1732        let quality = validate_summary_quality("some summary", &[]);
1733        assert_eq!(quality.file_coverage, 1.0);
1734        assert_eq!(quality.decision_coverage, 1.0);
1735    }
1736}