1use chrono::{DateTime, Local};
10use serde::{Deserialize, Serialize};
11
12use crate::constants::{
13 COMPACTION_AUTO_THRESHOLD_PERCENT, COMPACTION_MAX_RESPONSE_RESERVE_TOKENS,
14 COMPACTION_MIN_RESPONSE_RESERVE_TOKENS, COMPACTION_SUMMARIZER_INPUT_TOKEN_BUDGET,
15 COMPACTION_SUMMARY_MAX_TOKENS, COMPACTION_TAIL_TOKEN_BUDGET, COMPACTION_TAIL_TURNS,
16 COMPACTION_TOOL_OUTPUT_MAX_CHARS,
17};
18use crate::models::{ChatMessage, ChatMessageKind, MessageRole, ReasoningLevel, TokenUsage};
19
20use super::cmd::ChatRequest;
21use super::state::ContextUsageSnapshot;
22
23const CHECKPOINT_MARKER: &str = "MERMAID CONTEXT CHECKPOINT";
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(rename_all = "snake_case")]
27pub enum CompactionTrigger {
28 Manual,
29 AutoThreshold,
30 ContextLimitRetry,
31 TruncationRecovery,
34}
35
36impl CompactionTrigger {
37 pub fn as_str(self) -> &'static str {
38 match self {
39 Self::Manual => "manual",
40 Self::AutoThreshold => "auto_threshold",
41 Self::ContextLimitRetry => "context_limit_retry",
42 Self::TruncationRecovery => "truncation_recovery",
43 }
44 }
45
46 pub fn label(self) -> &'static str {
47 match self {
48 Self::Manual => "manual",
49 Self::AutoThreshold => "automatic",
50 Self::ContextLimitRetry => "context-limit retry",
51 Self::TruncationRecovery => "truncation recovery",
52 }
53 }
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
57pub struct CompactionPolicy {
58 pub auto_enabled: bool,
59 pub auto_threshold_percent: u8,
60 pub tail_turns: usize,
61 pub tail_token_budget: usize,
62 pub tool_output_max_chars: usize,
63 pub summary_max_tokens: usize,
64 pub summarizer_input_token_budget: usize,
65 pub min_response_reserve_tokens: usize,
66 pub max_response_reserve_tokens: usize,
67}
68
69impl Default for CompactionPolicy {
70 fn default() -> Self {
71 Self {
72 auto_enabled: true,
73 auto_threshold_percent: COMPACTION_AUTO_THRESHOLD_PERCENT,
74 tail_turns: COMPACTION_TAIL_TURNS,
75 tail_token_budget: COMPACTION_TAIL_TOKEN_BUDGET,
76 tool_output_max_chars: COMPACTION_TOOL_OUTPUT_MAX_CHARS,
77 summary_max_tokens: COMPACTION_SUMMARY_MAX_TOKENS,
78 summarizer_input_token_budget: COMPACTION_SUMMARIZER_INPUT_TOKEN_BUDGET,
79 min_response_reserve_tokens: COMPACTION_MIN_RESPONSE_RESERVE_TOKENS,
80 max_response_reserve_tokens: COMPACTION_MAX_RESPONSE_RESERVE_TOKENS,
81 }
82 }
83}
84
85impl CompactionPolicy {
86 pub fn response_reserve(self, request_max_tokens: usize) -> usize {
87 request_max_tokens
88 .max(self.min_response_reserve_tokens)
89 .min(self.max_response_reserve_tokens)
90 }
91}
92
93#[derive(Debug, Clone)]
94pub struct CompactionRequest {
95 pub chat: ChatRequest,
96 pub trigger: CompactionTrigger,
97 pub instructions: Option<String>,
98 pub force: bool,
99 pub policy: CompactionPolicy,
100}
101
102impl CompactionRequest {
103 pub fn manual(chat: ChatRequest, instructions: Option<String>) -> Self {
104 Self {
105 chat,
106 trigger: CompactionTrigger::Manual,
107 instructions,
108 force: true,
109 policy: CompactionPolicy::default(),
110 }
111 }
112
113 pub fn auto(chat: ChatRequest, trigger: CompactionTrigger) -> Self {
114 Self {
115 chat,
116 trigger,
117 instructions: None,
118 force: false,
119 policy: CompactionPolicy::default(),
120 }
121 }
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct CompactionRecord {
126 pub id: String,
127 pub trigger: CompactionTrigger,
128 pub created_at: DateTime<Local>,
129 pub before_tokens: usize,
130 pub after_tokens: usize,
131 pub archived_message_count: usize,
132 pub preserved_message_count: usize,
133 pub summary_tokens: usize,
134 pub duration_secs: f64,
135 #[serde(default = "default_verified")]
136 pub verified: bool,
137 #[serde(default)]
138 pub verification_error: Option<String>,
139 #[serde(default)]
140 pub focus: Option<String>,
141 #[serde(default)]
142 pub archive_path: Option<String>,
143}
144
145fn default_verified() -> bool {
146 true
147}
148
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct CompactionArchive {
151 pub id: String,
152 pub conversation_id: String,
153 pub created_at: DateTime<Local>,
154 pub messages: Vec<ChatMessage>,
155}
156
157#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
158pub struct CompactionResult {
159 pub record: CompactionRecord,
160 pub replacement_messages: Vec<ChatMessage>,
161 pub archived_messages: Vec<ChatMessage>,
162 pub before_snapshot: ContextUsageSnapshot,
163 pub after_snapshot: ContextUsageSnapshot,
164 pub usage: Option<TokenUsage>,
165}
166
167#[derive(Debug, Clone)]
168pub struct PreparedCompaction {
169 pub archived_messages: Vec<ChatMessage>,
170 pub preserved_messages: Vec<ChatMessage>,
171 pub previous_summary: Option<String>,
172 pub history_excerpt: String,
173}
174
175#[derive(Debug, Clone, PartialEq, Eq)]
176pub enum CompactionSkip {
177 NoKnownContextLimit,
178 AutoDisabled,
179 BelowThreshold,
180 NothingToCompact,
181}
182
183impl std::fmt::Display for CompactionSkip {
184 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
185 match self {
186 Self::NoKnownContextLimit => write!(f, "model context limit is unknown"),
187 Self::AutoDisabled => write!(f, "automatic compaction is disabled"),
188 Self::BelowThreshold => write!(f, "context is below compaction threshold"),
189 Self::NothingToCompact => write!(f, "not enough conversation history to summarize"),
190 }
191 }
192}
193
194pub fn should_auto_compact(
195 snapshot: &ContextUsageSnapshot,
196 request: &ChatRequest,
197 policy: CompactionPolicy,
198) -> Result<(), CompactionSkip> {
199 if !policy.auto_enabled {
200 return Err(CompactionSkip::AutoDisabled);
201 }
202 let Some(max_tokens) = snapshot.max_tokens else {
203 return Err(CompactionSkip::NoKnownContextLimit);
204 };
205 if max_tokens == 0 {
206 return Err(CompactionSkip::NoKnownContextLimit);
207 }
208
209 let reserve = policy.response_reserve(request.max_tokens);
210 let over_percent = snapshot
211 .used_percent
212 .is_some_and(|p| p >= policy.auto_threshold_percent);
213 let low_remaining = snapshot
214 .remaining_tokens
215 .is_some_and(|remaining| remaining <= reserve);
216 if over_percent || low_remaining {
217 Ok(())
218 } else {
219 Err(CompactionSkip::BelowThreshold)
220 }
221}
222
223pub fn context_exceeds_hard_limit(
224 snapshot: &ContextUsageSnapshot,
225 request: &ChatRequest,
226 policy: CompactionPolicy,
227) -> bool {
228 let Some(max_tokens) = snapshot.max_tokens else {
229 return false;
230 };
231 let reserve = policy.response_reserve(request.max_tokens);
232 snapshot.used_tokens.saturating_add(reserve) >= max_tokens
233}
234
235pub fn prepare_compaction(
236 request: &CompactionRequest,
237 max_context_tokens: Option<usize>,
238) -> Result<PreparedCompaction, CompactionSkip> {
239 let messages = &request.chat.messages;
240 if messages.len() < 3 {
241 return Err(CompactionSkip::NothingToCompact);
242 }
243
244 let split =
245 tail_start_index(messages, request.policy).ok_or(CompactionSkip::NothingToCompact)?;
246 if split == 0 {
247 return Err(CompactionSkip::NothingToCompact);
248 }
249
250 let archived_messages = messages[..split].to_vec();
251 let mut preserved_messages = messages[split..].to_vec();
252 if archived_messages.is_empty() || preserved_messages.is_empty() {
253 return Err(CompactionSkip::NothingToCompact);
254 }
255 let preserve_pending_tail = matches!(
264 request.trigger,
265 CompactionTrigger::ContextLimitRetry | CompactionTrigger::TruncationRecovery
266 );
267 drop_orphan_tool_calls(&mut preserved_messages, preserve_pending_tail);
268
269 let previous_summary = archived_messages
270 .iter()
271 .rev()
272 .find(|m| {
273 m.kind == ChatMessageKind::ContextCheckpoint || m.content.contains(CHECKPOINT_MARKER)
274 })
275 .map(|m| m.content.clone());
276
277 let max_input_tokens = max_context_tokens
278 .map(|max| max.saturating_sub(request.policy.response_reserve(request.chat.max_tokens)))
279 .filter(|max| *max > 0)
280 .unwrap_or(request.policy.summarizer_input_token_budget)
281 .min(request.policy.summarizer_input_token_budget);
282 let max_chars = max_input_tokens.saturating_mul(4).max(4_000);
283 let history_excerpt = truncate_middle(
284 &format_history_excerpt(&archived_messages, request.policy),
285 max_chars,
286 );
287
288 Ok(PreparedCompaction {
289 archived_messages,
290 preserved_messages,
291 previous_summary,
292 history_excerpt,
293 })
294}
295
296pub fn build_summary_request(
297 base: &ChatRequest,
298 prepared: &PreparedCompaction,
299 focus: Option<&str>,
300 policy: CompactionPolicy,
301) -> ChatRequest {
302 ChatRequest {
303 model_id: base.model_id.clone(),
304 messages: vec![ChatMessage::user(summary_prompt(prepared, focus))],
305 system_prompt: compaction_system_prompt().to_string(),
306 instructions: None,
307 reasoning: compaction_reasoning(base.reasoning),
308 temperature: 0.0,
309 max_tokens: policy.summary_max_tokens,
310 tools: Vec::new(),
311 ollama_num_ctx: base.ollama_num_ctx,
312 ollama_allow_ram_offload: base.ollama_allow_ram_offload,
313 }
314}
315
316pub fn build_verification_request(
317 base: &ChatRequest,
318 prepared: &PreparedCompaction,
319 draft_summary: &str,
320 focus: Option<&str>,
321 policy: CompactionPolicy,
322) -> ChatRequest {
323 let prompt = format!(
324 "{}\n\n# Draft Summary\n{}\n\n# Verification Task\nCritically check the draft against the conversation excerpt. If it omitted specific file paths, commands, test results, tool results, user constraints, current state, or next steps, return an improved complete checkpoint. Otherwise return the draft unchanged. Return only the final checkpoint markdown.",
325 summary_prompt(prepared, focus),
326 draft_summary.trim()
327 );
328 ChatRequest {
329 model_id: base.model_id.clone(),
330 messages: vec![ChatMessage::user(prompt)],
331 system_prompt: compaction_system_prompt().to_string(),
332 instructions: None,
333 reasoning: compaction_reasoning(base.reasoning),
334 temperature: 0.0,
335 max_tokens: policy.summary_max_tokens,
336 tools: Vec::new(),
337 ollama_num_ctx: base.ollama_num_ctx,
338 ollama_allow_ram_offload: base.ollama_allow_ram_offload,
339 }
340}
341
342pub fn build_replacement_messages(
343 summary: &str,
344 prepared: &PreparedCompaction,
345 record: &CompactionRecord,
346) -> Vec<ChatMessage> {
347 let summary = crate::utils::redact_secrets(summary);
351 let summary = summary.as_str();
352 let checkpoint = format!(
353 "# {}\n\nCompaction id: {}\nTrigger: {}\nCreated: {}\nArchived messages: {}\nPreserved messages: {}\n\n{}",
354 CHECKPOINT_MARKER,
355 record.id,
356 record.trigger.as_str(),
357 record.created_at.to_rfc3339(),
358 record.archived_message_count,
359 record.preserved_message_count,
360 summary.trim()
361 );
362 let mut user = ChatMessage::user(checkpoint);
363 user.kind = ChatMessageKind::ContextCheckpoint;
364 user.metadata = Some(serde_json::json!({
365 "compaction_id": record.id,
366 "trigger": record.trigger.as_str(),
367 "before_tokens": record.before_tokens,
368 "after_tokens": record.after_tokens,
369 "archived_message_count": record.archived_message_count,
370 "preserved_message_count": record.preserved_message_count,
371 "duration_secs": record.duration_secs,
372 "verified": record.verified,
373 "verification_error": record.verification_error,
374 }));
375
376 let mut assistant = ChatMessage::assistant(compaction_receipt(record));
377 assistant.kind = ChatMessageKind::ContextCheckpoint;
378 assistant.metadata = user.metadata.clone();
379
380 let mut messages = Vec::with_capacity(2 + prepared.preserved_messages.len());
381 messages.push(user);
382 messages.push(assistant);
383 messages.extend(prepared.preserved_messages.clone());
384 messages
385}
386
387pub fn compaction_receipt(record: &CompactionRecord) -> String {
388 let verification = if record.verified {
389 "Verified.".to_string()
390 } else if let Some(error) = &record.verification_error {
391 format!("Used draft summary because verification failed: {error}.")
392 } else {
393 "Used draft summary without verification.".to_string()
394 };
395 format!(
396 "Context compacted: {} -> {} tokens, archived {} messages, preserved {} messages, took {:.1}s. {} I will continue from this checkpoint.",
397 format_compact_count(record.before_tokens),
398 format_compact_count(record.after_tokens),
399 record.archived_message_count,
400 record.preserved_message_count,
401 record.duration_secs,
402 verification
403 )
404}
405
406pub fn normalize_summary(text: &str) -> String {
407 let trimmed = text.trim();
408 if let Some(summary) = extract_tagged_summary(trimmed) {
409 return summary.trim().to_string();
410 }
411 trimmed.to_string()
412}
413
414pub fn combine_usage(a: Option<TokenUsage>, b: Option<TokenUsage>) -> Option<TokenUsage> {
415 match (a, b) {
416 (None, None) => None,
417 (Some(u), None) | (None, Some(u)) => Some(u),
418 (Some(mut left), Some(right)) => {
419 left.prompt_tokens = left.prompt_tokens.saturating_add(right.prompt_tokens);
420 left.completion_tokens = left
421 .completion_tokens
422 .saturating_add(right.completion_tokens);
423 left.total_tokens = left.total_tokens.saturating_add(right.total_tokens);
424 left.cached_input_tokens = left
425 .cached_input_tokens
426 .saturating_add(right.cached_input_tokens);
427 left.cache_creation_input_tokens = left
428 .cache_creation_input_tokens
429 .saturating_add(right.cache_creation_input_tokens);
430 left.reasoning_output_tokens = left
431 .reasoning_output_tokens
432 .saturating_add(right.reasoning_output_tokens);
433 Some(left)
434 },
435 }
436}
437
438pub fn estimate_messages_tokens(messages: &[ChatMessage]) -> usize {
439 messages.iter().map(estimate_message_tokens).sum()
440}
441
442pub fn format_compact_count(value: usize) -> String {
449 if value >= 1_000_000 {
450 format_scaled(value, 1_000_000, "M")
451 } else if value >= 1_000 {
452 format_scaled(value, 1_000, "k")
453 } else {
454 value.to_string()
455 }
456}
457
458fn format_scaled(value: usize, divisor: usize, suffix: &str) -> String {
459 let whole = value / divisor;
460 let decimal = ((value % divisor) * 10) / divisor;
461 if decimal == 0 {
462 format!("{}{}", whole, suffix)
463 } else {
464 format!("{}.{}{}", whole, decimal, suffix)
465 }
466}
467
468fn compaction_system_prompt() -> &'static str {
469 "You are performing context checkpoint compaction for Mermaid, a model-agnostic agentic coding CLI. Produce a faithful handoff summary for the next model call. Preserve exact file paths, commands, errors, tool results, user preferences, decisions, current state, and next steps. Do not invent facts. Be concise but complete."
470}
471
472fn compaction_reasoning(current: ReasoningLevel) -> ReasoningLevel {
473 match current {
474 ReasoningLevel::None | ReasoningLevel::Minimal => current,
475 _ => ReasoningLevel::Low,
476 }
477}
478
479fn summary_prompt(prepared: &PreparedCompaction, focus: Option<&str>) -> String {
480 let anchor = prepared
481 .previous_summary
482 .as_deref()
483 .map(|summary| {
484 format!(
485 "A previous checkpoint exists. Update it with the newer history, preserve still-true details, and remove stale details.\n\n<previous_checkpoint>\n{}\n</previous_checkpoint>",
486 summary.trim()
487 )
488 })
489 .unwrap_or_else(|| "Create a new checkpoint from the conversation history below.".to_string());
490
491 let focus = focus
492 .filter(|s| !s.trim().is_empty())
493 .map(|s| format!("\n# User Focus Instructions\n{}\n", s.trim()))
494 .unwrap_or_default();
495
496 format!(
497 "{anchor}{focus}\n# Required Output\nReturn exactly this Markdown structure and keep section order:\n\n## Goal\n- [single-sentence task summary]\n\n## User Preferences And Constraints\n- [preferences, constraints, mode, or \"(none)\"]\n\n## Project State\n- [repo/product state and important architecture facts]\n\n## Completed Work\n- [what has already been done]\n\n## Current Work\n- [what is actively in progress]\n\n## Key Decisions\n- [decision and rationale]\n\n## Critical Files And Symbols\n- [file path or symbol: why it matters]\n\n## Commands Tests And Results\n- [command/test/result/error]\n\n## Open Questions Or Risks\n- [risk/question/blocker]\n\n## Next Steps\n- [ordered next action]\n\nRules:\n- Preserve exact paths, commands, error strings, identifiers, and numeric facts when known.\n- Mention important omitted or truncated data explicitly.\n- Do not mention that you are an AI or explain the compaction process.\n\n# Conversation History To Compact\n{}",
498 prepared.history_excerpt
499 )
500}
501
502fn drop_orphan_tool_calls(messages: &mut Vec<ChatMessage>, preserve_pending_tail: bool) {
525 let pending_tail = if preserve_pending_tail
526 && messages.last().is_some_and(|m| {
527 m.role == MessageRole::Assistant && m.tool_calls.as_ref().is_some_and(|c| !c.is_empty())
528 }) {
529 Some(messages.len() - 1)
530 } else {
531 None
532 };
533
534 let answered: std::collections::HashSet<String> = messages
535 .iter()
536 .filter(|m| m.role == MessageRole::Tool)
537 .filter_map(|m| m.tool_call_id.clone())
538 .collect();
539
540 for (idx, m) in messages.iter_mut().enumerate() {
542 if Some(idx) == pending_tail {
543 continue;
544 }
545 let Some(calls) = m.tool_calls.as_mut() else {
546 continue;
547 };
548 calls.retain(|c| c.id.as_deref().is_some_and(|id| answered.contains(id)));
549 if calls.is_empty() {
550 m.tool_calls = None;
551 }
552 }
553
554 let emitted: std::collections::HashSet<String> = messages
557 .iter()
558 .filter_map(|m| m.tool_calls.as_ref())
559 .flat_map(|calls| calls.iter())
560 .filter_map(|c| c.id.clone())
561 .collect();
562 messages.retain(|m| {
563 m.role != MessageRole::Tool
564 || m.tool_call_id
565 .as_deref()
566 .is_some_and(|id| emitted.contains(id))
567 });
568}
569
570fn tail_start_index(messages: &[ChatMessage], policy: CompactionPolicy) -> Option<usize> {
571 let mut user_turns = 0usize;
572 let mut start = None;
573 for (idx, msg) in messages.iter().enumerate().rev() {
574 if msg.role == MessageRole::User {
575 user_turns += 1;
576 start = Some(idx);
577 if user_turns >= policy.tail_turns {
578 break;
579 }
580 }
581 }
582 let mut start = start?;
583 while estimate_messages_tokens(&messages[start..]) > policy.tail_token_budget {
584 let next_user = messages
585 .iter()
586 .enumerate()
587 .skip(start + 1)
588 .find(|(_, msg)| msg.role == MessageRole::User)
589 .map(|(idx, _)| idx);
590 match next_user {
591 Some(idx) => start = idx,
592 None => break,
593 }
594 }
595 Some(start)
596}
597
598fn format_history_excerpt(messages: &[ChatMessage], policy: CompactionPolicy) -> String {
599 let mut out = String::new();
600 for (idx, msg) in messages.iter().enumerate() {
601 let role = match msg.role {
602 MessageRole::User => "USER",
603 MessageRole::Assistant => "ASSISTANT",
604 MessageRole::System => "SYSTEM",
605 MessageRole::Tool => "TOOL",
606 };
607 out.push_str(&format!("\n\n--- MESSAGE {} [{}] ---\n", idx + 1, role));
608 if msg.kind != ChatMessageKind::Normal {
609 out.push_str(&format!("kind: {:?}\n", msg.kind));
610 }
611 if let Some(name) = &msg.tool_name {
612 out.push_str(&format!("tool_name: {}\n", name));
613 }
614 if let Some(id) = &msg.tool_call_id {
615 out.push_str(&format!("tool_call_id: {}\n", id));
616 }
617 if let Some(calls) = &msg.tool_calls {
618 let names: Vec<&str> = calls
619 .iter()
620 .map(|call| call.function.name.as_str())
621 .collect();
622 out.push_str(&format!("tool_calls: {}\n", names.join(", ")));
623 }
624 if let Some(images) = &msg.images
625 && !images.is_empty()
626 {
627 out.push_str(&format!("[{} image attachment(s) omitted]\n", images.len()));
628 }
629 for action in &msg.actions {
630 out.push_str(&format!(
631 "action: {}({}) duration={:?}\n",
632 action.action_type, action.target, action.duration_seconds
633 ));
634 if let Some(metadata) = &action.metadata {
635 out.push_str(&format!("action_metadata: {:?}\n", metadata));
636 }
637 }
638 let cap = if msg.role == MessageRole::Tool {
639 policy.tool_output_max_chars
640 } else {
641 policy.tool_output_max_chars.saturating_mul(4)
642 };
643 out.push_str(&truncate_middle(&msg.content, cap));
644 }
645 out
646}
647
648fn estimate_message_tokens(msg: &ChatMessage) -> usize {
649 let mut chars = msg.content.len();
650 chars = chars.saturating_add(format!("{:?}", msg.role).len());
651 chars = chars.saturating_add(msg.tool_name.as_deref().map(str::len).unwrap_or(0));
652 chars = chars.saturating_add(msg.tool_call_id.as_deref().map(str::len).unwrap_or(0));
653 if let Some(images) = &msg.images {
654 chars = chars.saturating_add(images.iter().map(String::len).sum::<usize>());
655 }
656 if let Some(tool_calls) = &msg.tool_calls {
661 for tc in tool_calls {
662 chars = chars.saturating_add(tc.function.name.len());
663 chars = chars.saturating_add(tc.function.arguments.to_string().len());
664 chars = chars.saturating_add(tc.id.as_deref().map(str::len).unwrap_or(0));
665 }
666 }
667 chars.div_ceil(4)
668}
669
670fn truncate_middle(text: &str, max_chars: usize) -> String {
671 if text.chars().count() <= max_chars {
672 return text.to_string();
673 }
674 if max_chars < 128 {
675 return text.chars().take(max_chars).collect();
676 }
677 let marker = "\n\n[... truncated during context compaction ...]\n\n";
678 let keep = max_chars.saturating_sub(marker.len());
679 let head = keep / 2;
680 let tail = keep.saturating_sub(head);
681 let start: String = text.chars().take(head).collect();
682 let end: String = text
683 .chars()
684 .rev()
685 .take(tail)
686 .collect::<Vec<_>>()
687 .into_iter()
688 .rev()
689 .collect();
690 format!("{start}{marker}{end}")
691}
692
693fn extract_tagged_summary(text: &str) -> Option<&str> {
694 let start_tag = "<summary>";
695 let end_tag = "</summary>";
696 let start = text.find(start_tag)? + start_tag.len();
697 let end = text[start..].find(end_tag)? + start;
698 Some(&text[start..end])
699}
700
701#[cfg(test)]
702mod tests {
703 use super::*;
704
705 fn request_with(messages: Vec<ChatMessage>) -> ChatRequest {
706 ChatRequest {
707 model_id: "ollama/test".to_string(),
708 messages,
709 system_prompt: "system".to_string(),
710 instructions: None,
711 reasoning: ReasoningLevel::Medium,
712 temperature: 0.7,
713 max_tokens: 4096,
714 tools: Vec::new(),
715 ollama_num_ctx: None,
716 ollama_allow_ram_offload: None,
717 }
718 }
719
720 #[test]
721 fn auto_compaction_triggers_by_percent() {
722 let snapshot = ContextUsageSnapshot::from_estimate(
723 super::super::state::PromptTokenBreakdown {
724 system_tokens: 0,
725 instructions_tokens: 0,
726 message_tokens: 86,
727 tool_schema_tokens: 0,
728 image_count: 0,
729 message_count: 2,
730 tool_count: 0,
731 },
732 Some(100),
733 );
734 let req = request_with(vec![ChatMessage::user("hello")]);
735 assert!(should_auto_compact(&snapshot, &req, CompactionPolicy::default()).is_ok());
736 }
737
738 #[test]
739 fn prepare_preserves_recent_two_user_turns() {
740 let messages = vec![
741 ChatMessage::user("one"),
742 ChatMessage::assistant("one answer"),
743 ChatMessage::user("two"),
744 ChatMessage::assistant("two answer"),
745 ChatMessage::user("three"),
746 ];
747 let request = CompactionRequest::manual(request_with(messages), None);
748 let prepared = prepare_compaction(&request, Some(100_000)).expect("prepared");
749 assert_eq!(prepared.archived_messages.len(), 2);
750 assert_eq!(prepared.preserved_messages.len(), 3);
751 assert_eq!(prepared.preserved_messages[0].content, "two");
752 }
753
754 fn tool_call(id: &str, name: &str) -> crate::models::tool_call::ToolCall {
755 crate::models::tool_call::ToolCall {
756 id: Some(id.to_string()),
757 function: crate::models::tool_call::FunctionCall {
758 name: name.to_string(),
759 arguments: serde_json::json!({}),
760 },
761 }
762 }
763
764 #[test]
765 fn prepare_strips_orphan_tool_call_from_preserved_tail() {
766 let mut orphan = ChatMessage::assistant("calling a tool");
769 orphan.tool_calls = Some(vec![tool_call("call_1", "do_thing")]);
770 let messages = vec![
771 ChatMessage::user("one"),
772 ChatMessage::assistant("one answer"),
773 ChatMessage::user("two"),
774 orphan,
775 ChatMessage::user("three"),
776 ];
777 let request = CompactionRequest::manual(request_with(messages), None);
778 let prepared = prepare_compaction(&request, Some(100_000)).expect("prepared");
779 let has_orphan = prepared
780 .preserved_messages
781 .iter()
782 .any(|m| m.tool_calls.as_ref().is_some_and(|c| !c.is_empty()));
783 assert!(
784 !has_orphan,
785 "orphan tool_use must be stripped from the tail"
786 );
787 assert!(
789 prepared
790 .preserved_messages
791 .iter()
792 .any(|m| m.content == "calling a tool")
793 );
794 }
795
796 #[test]
797 fn prepare_keeps_paired_tool_call_in_tail() {
798 let mut asst = ChatMessage::assistant("calling");
800 asst.tool_calls = Some(vec![tool_call("call_1", "do_thing")]);
801 let messages = vec![
802 ChatMessage::user("one"),
803 ChatMessage::assistant("one answer"),
804 ChatMessage::user("two"),
805 asst,
806 ChatMessage::tool("call_1", "do_thing", "ok"),
807 ChatMessage::user("three"),
808 ];
809 let request = CompactionRequest::manual(request_with(messages), None);
810 let prepared = prepare_compaction(&request, Some(100_000)).expect("prepared");
811 let kept = prepared
812 .preserved_messages
813 .iter()
814 .any(|m| m.tool_calls.as_ref().is_some_and(|c| !c.is_empty()));
815 assert!(kept, "a tool_call paired with its result must be preserved");
816 }
817
818 #[test]
819 fn prepare_drops_reverse_orphan_tool_result_from_tail() {
820 let mut asst = ChatMessage::assistant("calling");
825 asst.tool_calls = Some(vec![tool_call("call_1", "do_thing")]);
826 let messages = vec![
827 ChatMessage::user("one"),
828 asst,
829 ChatMessage::user("two"),
830 ChatMessage::tool("call_1", "do_thing", "result"),
831 ChatMessage::user("three"),
832 ];
833 let request = CompactionRequest::manual(request_with(messages), None);
836 let prepared = prepare_compaction(&request, Some(100_000)).expect("prepared");
837 assert!(
838 prepared
839 .preserved_messages
840 .iter()
841 .all(|m| m.role != MessageRole::Tool),
842 "an orphan tool_result whose tool_use was archived must be dropped"
843 );
844 }
845
846 #[test]
847 fn prepare_keeps_pending_trailing_tool_use_on_retry() {
848 let mut pending = ChatMessage::assistant("calling a tool");
852 pending.tool_calls = Some(vec![tool_call("call_9", "do_thing")]);
853 let messages = vec![
854 ChatMessage::user("one"),
855 ChatMessage::assistant("a1"),
856 ChatMessage::user("two"),
857 ChatMessage::assistant("a2"),
858 ChatMessage::user("three"),
859 pending,
860 ];
861 let request =
862 CompactionRequest::auto(request_with(messages), CompactionTrigger::ContextLimitRetry);
863 let prepared = prepare_compaction(&request, Some(100_000)).expect("prepared");
864 let last = prepared
865 .preserved_messages
866 .last()
867 .expect("non-empty preserved tail");
868 assert!(
869 last.tool_calls
870 .as_ref()
871 .is_some_and(|c| c.iter().any(|call| call.id.as_deref() == Some("call_9"))),
872 "a pending trailing tool_use must be preserved across a retry compaction"
873 );
874 }
875
876 #[test]
877 fn prepare_drops_trailing_tool_use_on_manual_compaction() {
878 let mut pending = ChatMessage::assistant("calling a tool");
882 pending.tool_calls = Some(vec![tool_call("call_9", "do_thing")]);
883 let messages = vec![
884 ChatMessage::user("one"),
885 ChatMessage::assistant("a1"),
886 ChatMessage::user("two"),
887 ChatMessage::assistant("a2"),
888 ChatMessage::user("three"),
889 pending,
890 ];
891 let request = CompactionRequest::manual(request_with(messages), None);
892 let prepared = prepare_compaction(&request, Some(100_000)).expect("prepared");
893 assert!(
894 !prepared
895 .preserved_messages
896 .iter()
897 .any(|m| m.tool_calls.as_ref().is_some_and(|c| !c.is_empty())),
898 "manual compaction must scrub the trailing orphan tool_use"
899 );
900 assert!(
901 prepared
902 .preserved_messages
903 .iter()
904 .any(|m| m.content == "calling a tool"),
905 "the assistant text is kept even though the orphan call is dropped"
906 );
907 }
908
909 #[test]
910 fn replacement_starts_with_checkpoint_and_ack() {
911 let prepared = PreparedCompaction {
912 archived_messages: vec![ChatMessage::user("old")],
913 preserved_messages: vec![ChatMessage::user("new")],
914 previous_summary: None,
915 history_excerpt: "old".to_string(),
916 };
917 let record = CompactionRecord {
918 id: "c1".to_string(),
919 trigger: CompactionTrigger::Manual,
920 created_at: Local::now(),
921 before_tokens: 100,
922 after_tokens: 25,
923 archived_message_count: 1,
924 preserved_message_count: 1,
925 summary_tokens: 10,
926 duration_secs: 1.0,
927 verified: true,
928 verification_error: None,
929 focus: None,
930 archive_path: None,
931 };
932 let messages = build_replacement_messages("## Goal\n- continue", &prepared, &record);
933 assert_eq!(messages[0].kind, ChatMessageKind::ContextCheckpoint);
934 assert!(messages[0].content.contains(CHECKPOINT_MARKER));
935 assert_eq!(messages[2].content, "new");
936 }
937
938 #[test]
939 fn replacement_metadata_records_verification_status() {
940 let prepared = PreparedCompaction {
941 archived_messages: vec![ChatMessage::user("old")],
942 preserved_messages: vec![ChatMessage::user("new")],
943 previous_summary: None,
944 history_excerpt: "old".to_string(),
945 };
946 let record = CompactionRecord {
947 id: "c1".to_string(),
948 trigger: CompactionTrigger::Manual,
949 created_at: Local::now(),
950 before_tokens: 100,
951 after_tokens: 25,
952 archived_message_count: 1,
953 preserved_message_count: 1,
954 summary_tokens: 10,
955 duration_secs: 1.0,
956 verified: false,
957 verification_error: Some("provider overloaded".to_string()),
958 focus: None,
959 archive_path: None,
960 };
961 let messages = build_replacement_messages("## Goal\n- continue", &prepared, &record);
962 let metadata = messages[0].metadata.as_ref().expect("metadata");
963 assert_eq!(
964 metadata.get("verified").and_then(|v| v.as_bool()),
965 Some(false)
966 );
967 assert_eq!(
968 metadata.get("verification_error").and_then(|v| v.as_str()),
969 Some("provider overloaded")
970 );
971 assert!(messages[1].content.contains("Used draft summary"));
972 }
973}