Skip to main content

vtcode_core/compaction/
two_pass.rs

1//! Pure builders for prefire two-pass compaction.
2//!
3//! Pass1 summarizes ~95% of history (by estimated-token weight) -> NOTE₁.
4//! Pass2 rewrites NOTE₁ + the ~5% tail into the successor-visible NOTE₂.
5//! Sampling lives in [`super`]; this module has no I/O.
6
7use crate::llm::provider::Message;
8
9/// Default history fraction covered by pass1; the remainder is the blocking
10/// pass2 tail, so keep it small (prod pass2 latency is dominated by tail prefill).
11pub const TWO_PASS_DEFAULT_SPLIT_FRACTION: f64 = 0.95;
12
13/// Minimum char length for a closed `<summary>` block to be preferred as NOTE₁
14/// over the full pass1 response.
15const TWO_PASS_MIN_SUMMARY_BLOCK_CHARS: usize = 1000;
16
17/// Cap on NOTE₁ text embedded in pass2 (once, in the carrier turn).
18const TWO_PASS_MAX_NOTE1_CHARS: usize = 12_000;
19
20/// Result of splitting a conversation for two-pass compaction.
21#[derive(Debug, Clone, Copy)]
22pub struct TwoPassSplit<'a> {
23    pub prefix: &'a [Message],
24    pub tail: &'a [Message],
25    pub split_idx: usize,
26}
27
28/// Choose a split index so prefix weight is at least `fraction` of total.
29#[allow(
30    clippy::cast_sign_loss,
31    reason = "Intentional compatibility, platform, or test-only suppression."
32)]
33fn split_index_by_token_fraction(weights: &[usize], fraction: f64) -> usize {
34    if weights.is_empty() {
35        return 0;
36    }
37    let frac = fraction.clamp(0.05, 0.95);
38    let total_w: usize = weights.iter().copied().sum::<usize>().max(1);
39    let target_w = (frac * total_w as f64).round() as u64 as usize;
40    let mut acc = 0usize;
41    let mut split_idx = weights.len().saturating_sub(1).max(1);
42    for (i, w) in weights.iter().enumerate() {
43        acc = acc.saturating_add(*w);
44        if acc >= target_w {
45            split_idx = (i + 1).max(1);
46            break;
47        }
48    }
49    if split_idx >= weights.len() && weights.len() > 1 {
50        split_idx = weights.len() - 1;
51    }
52    split_idx
53}
54
55/// Never separate an assistant `tool_calls` turn from its following `tool_call_id` results.
56fn snap_split_idx_to_tool_boundaries(conversation: &[Message], mut split_idx: usize) -> usize {
57    let n = conversation.len();
58    if n == 0 {
59        return 0;
60    }
61    split_idx = split_idx.min(n);
62
63    while split_idx < n {
64        let Some(msg) = conversation.get(split_idx) else {
65            break;
66        };
67        if msg.tool_call_id.is_some() {
68            split_idx += 1;
69        } else {
70            break;
71        }
72    }
73    if split_idx < n
74        && let Some(msg) = conversation.get(split_idx)
75        && msg.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty())
76    {
77        split_idx += 1;
78        while split_idx < n {
79            let Some(msg) = conversation.get(split_idx) else {
80                break;
81            };
82            if msg.tool_call_id.is_some() {
83                split_idx += 1;
84            } else {
85                break;
86            }
87        }
88    }
89    while split_idx > 0 && split_idx < n {
90        let Some(msg) = conversation.get(split_idx - 1) else {
91            break;
92        };
93        if msg.tool_calls.as_ref().is_some_and(|tool_calls| !tool_calls.is_empty()) {
94            if conversation.get(split_idx).map(|m| m.tool_call_id.is_some()) != Some(true) {
95                break;
96            }
97            while split_idx < n {
98                let Some(msg) = conversation.get(split_idx) else {
99                    break;
100                };
101                if msg.tool_call_id.is_some() {
102                    split_idx += 1;
103                } else {
104                    break;
105                }
106            }
107        } else {
108            break;
109        }
110    }
111
112    if split_idx >= n && n > 1 {
113        let mut candidate = n - 1;
114        while candidate > 1 {
115            if let Some(msg) = conversation.get(candidate) {
116                if msg.tool_call_id.is_none() {
117                    break;
118                }
119                candidate -= 1;
120            } else {
121                break;
122            }
123        }
124        if candidate > 0
125            && let Some(msg) = conversation.get(candidate)
126            && msg.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty())
127        {
128            if conversation.get(candidate + 1).map(|m| m.tool_call_id.is_some()) == Some(true) {
129                split_idx = candidate;
130            }
131        } else if candidate > 0 && conversation.get(candidate).map(|m| m.tool_call_id.is_some()) == Some(true) {
132            let mut i = candidate;
133            while i > 0 && conversation.get(i - 1).map(|m| m.tool_call_id.is_some()) == Some(true) {
134                i -= 1;
135            }
136            if conversation
137                .get(i)
138                .is_some_and(|m| m.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()))
139            {
140                split_idx = i;
141            }
142        }
143        if candidate >= 1 && candidate < n {
144            split_idx = candidate;
145        }
146    }
147
148    split_idx.min(n)
149}
150
151/// Split `conversation` into pass1 prefix / pass2 tail by estimated-token weight.
152pub fn split_conversation_for_two_pass(conversation: &[Message], split_fraction: f64) -> TwoPassSplit<'_> {
153    let weights: Vec<usize> = conversation.iter().map(|m| m.estimate_tokens()).collect();
154    let mut split_idx = split_index_by_token_fraction(&weights, split_fraction);
155    split_idx = snap_split_idx_to_tool_boundaries(conversation, split_idx);
156    let split_idx = split_idx.min(conversation.len());
157    TwoPassSplit {
158        prefix: &conversation[..split_idx],
159        tail: &conversation[split_idx..],
160        split_idx,
161    }
162}
163
164fn extract_summary_block(text: &str, min_chars: usize) -> Option<String> {
165    if text.is_empty() {
166        return None;
167    }
168    let open = "<summary>";
169    let close = "</summary>";
170    let lower = text.to_ascii_lowercase();
171    let mut blocks: Vec<String> = Vec::new();
172    let text_bytes = text.as_bytes();
173    let lower_bytes = lower.as_bytes();
174    let mut search_from = 0usize;
175    while search_from < lower_bytes.len() {
176        let Some(rel) = find_bytes(&lower_bytes[search_from..], open.as_bytes()) else {
177            break;
178        };
179        let start = search_from + rel + open.len();
180        let Some(rel_close) = find_bytes(&lower_bytes[start..], close.as_bytes()) else {
181            break;
182        };
183        let end = start + rel_close;
184        let inner = std::str::from_utf8(&text_bytes[start..end]).unwrap_or("");
185        blocks.push(inner.to_string());
186        search_from = end + close.len();
187    }
188    for block in blocks.into_iter().rev() {
189        let stripped = block.trim();
190        if stripped.chars().count() > min_chars {
191            return Some(stripped.to_string());
192        }
193    }
194    None
195}
196
197fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
198    haystack.windows(needle.len()).position(|w| w == needle)
199}
200
201/// Prefer a substantive `<summary>` inner for NOTE₁; otherwise the full pass1 response.
202pub fn note_for_two_pass_pass2(pass1_raw: &str) -> String {
203    let mut note = extract_summary_block(pass1_raw, TWO_PASS_MIN_SUMMARY_BLOCK_CHARS)
204        .unwrap_or_else(|| pass1_raw.trim().to_string());
205    let n = note.chars().count();
206    if n > TWO_PASS_MAX_NOTE1_CHARS {
207        note = note.chars().take(TWO_PASS_MAX_NOTE1_CHARS).collect();
208        note.push_str("\n\n[… NOTE₁ truncated for pass2 input budget …]");
209    }
210    note
211}
212
213fn format_two_pass_note1_carrier(note1: &str) -> String {
214    let note1 = note1.trim();
215    format!(
216        "Your conversation was summarized due to context constraints. \
217         Here is the summary of the conversation so far:\n\n\
218         <summary_content>\n{note1}\n</summary_content>\n\n\
219         Continue with the compaction task below."
220    )
221}
222
223/// Final pass2 instruction. NOTE₁ is already in the carrier turn at the start
224/// of the pass2 history, so this refers to it instead of repeating it.
225fn format_two_pass_special_pass2_user(compaction_prompt: &str) -> String {
226    let uq = if compaction_prompt.trim().is_empty() {
227        "Please summarize the conversation so far."
228    } else {
229        compaction_prompt
230    };
231    format!(
232        "This is a special compaction case (two-pass / hierarchical summarization).\n\
233         You are writing the final compaction note that a successor assistant will \
234         rely on as its only memory of the conversation; it will not see the prior summary.\n\n\
235         - Incorporate the whole prior summary (the <summary_content> block at the start \
236         of this conversation) into the final note. Do not omit its sections, defer to \
237         \"see prior compaction\", or drop early history because newer turns are in context.\n\
238         - Merge it with the more recent conversation turns above into one coherent, \
239         faithful, self-contained summary (the same structure/sections you normally use \
240         for compaction).\n\
241         - Preserve concrete values, file paths, errors/blockers, operational how-tos, \
242         key findings, and pending tasks from both the prior summary and the recent \
243         turns when they still matter.\n\n\
244         Compaction instruction:\n\
245         {uq}"
246    )
247}
248
249/// Pass1 sample history: `prefix` + compaction instruction user turn.
250pub fn build_two_pass_pass1_history(prefix: &[Message], compaction_prompt: &str) -> Vec<Message> {
251    let mut history = prefix.to_vec();
252    history.push(Message {
253        role: crate::llm::provider::MessageRole::User,
254        content: crate::llm::provider::MessageContent::Text(compaction_prompt.to_string()),
255        reasoning: None,
256        reasoning_details: None,
257        tool_calls: None,
258        tool_call_id: None,
259        phase: None,
260        origin_tool: None,
261        metadata: None,
262        clear_at: None,
263    });
264    history
265}
266
267/// Pass2 sample history: system (from prefix) + NOTE₁ carrier + tail + special turn.
268///
269/// Successor-visible artifact is the model output of *this* history only (NOTE₂).
270pub fn build_two_pass_pass2_history(
271    prefix: &[Message],
272    tail: &[Message],
273    note1: &str,
274    compaction_prompt: &str,
275) -> Vec<Message> {
276    let mut history: Vec<Message> = Vec::new();
277
278    for item in prefix {
279        if item.role == crate::llm::provider::MessageRole::System {
280            history.push(item.clone());
281        }
282    }
283    if history.is_empty() {
284        history.push(Message {
285            role: crate::llm::provider::MessageRole::System,
286            content: crate::llm::provider::MessageContent::Text(
287                "You are condensing a coding-agent session so work can continue without the full history.".to_string(),
288            ),
289            reasoning: None,
290            reasoning_details: None,
291            tool_calls: None,
292            tool_call_id: None,
293            phase: None,
294            origin_tool: None,
295            metadata: None,
296            clear_at: None,
297        });
298    }
299
300    history.push(Message {
301        role: crate::llm::provider::MessageRole::User,
302        content: crate::llm::provider::MessageContent::Text(format_two_pass_note1_carrier(note1)),
303        reasoning: None,
304        reasoning_details: None,
305        tool_calls: None,
306        tool_call_id: None,
307        phase: None,
308        origin_tool: None,
309        metadata: None,
310        clear_at: None,
311    });
312    history.extend(tail.iter().cloned());
313    history.push(Message {
314        role: crate::llm::provider::MessageRole::User,
315        content: crate::llm::provider::MessageContent::Text(format_two_pass_special_pass2_user(compaction_prompt)),
316        reasoning: None,
317        reasoning_details: None,
318        tool_calls: None,
319        tool_call_id: None,
320        phase: None,
321        origin_tool: None,
322        metadata: None,
323        clear_at: None,
324    });
325    history
326}
327
328use std::hash::{Hash, Hasher};
329
330/// Cheap fingerprint of a conversation prefix for prefire NOTE₁ validity.
331/// A mismatch means the prefix changed (edit / rewind / branch) since pass-1,
332/// so the cached NOTE₁ no longer summarizes the current prefix.
333pub fn fingerprint_prefix(prefix: &[Message]) -> u64 {
334    use std::collections::hash_map::DefaultHasher;
335    let mut h = DefaultHasher::new();
336    prefix.len().hash(&mut h);
337    for msg in prefix {
338        let tag: u8 = match msg.role {
339            crate::llm::provider::MessageRole::System => 0,
340            crate::llm::provider::MessageRole::User => 1,
341            crate::llm::provider::MessageRole::Assistant => 2,
342            crate::llm::provider::MessageRole::Tool => 3,
343        };
344        tag.hash(&mut h);
345        msg.content.as_text().hash(&mut h);
346    }
347    h.finish()
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353    use crate::llm::provider::{FunctionCall, MessageRole, ToolCall};
354
355    fn make_tool_call(id: &str, name: &str) -> Option<Vec<ToolCall>> {
356        Some(vec![ToolCall {
357            id: id.to_string(),
358            call_type: "function".to_string(),
359            function: Some(FunctionCall {
360                namespace: None,
361                name: name.to_string(),
362                arguments: "{}".to_string(),
363            }),
364            text: None,
365            thought_signature: None,
366        }])
367    }
368
369    #[test]
370    fn split_leaves_tail_when_possible() {
371        let weights = vec![10usize, 10, 10, 10, 10];
372        let idx = split_index_by_token_fraction(&weights, 0.9);
373        assert!(idx < weights.len());
374        assert!(idx >= 1);
375    }
376
377    #[test]
378    fn default_split_fraction_leaves_five_percent_tail() {
379        assert!((TWO_PASS_DEFAULT_SPLIT_FRACTION - 0.95).abs() < f64::EPSILON);
380        let weights = vec![10usize; 40];
381        let idx = split_index_by_token_fraction(&weights, TWO_PASS_DEFAULT_SPLIT_FRACTION);
382        assert_eq!(idx, 38);
383        assert!(idx < weights.len());
384    }
385
386    #[test]
387    fn split_does_not_sever_tool_pairs() {
388        let assistant = Message {
389            role: MessageRole::Assistant,
390            content: crate::llm::provider::MessageContent::Text("call".to_string()),
391            reasoning: None,
392            reasoning_details: None,
393            tool_calls: make_tool_call("tc1", "bash"),
394            tool_call_id: None,
395            phase: None,
396            origin_tool: None,
397            metadata: None,
398            clear_at: None,
399        };
400        let conv = vec![
401            Message {
402                role: MessageRole::User,
403                content: crate::llm::provider::MessageContent::Text("a".repeat(400)),
404                reasoning: None,
405                reasoning_details: None,
406                tool_calls: None,
407                tool_call_id: None,
408                phase: None,
409                origin_tool: None,
410                metadata: None,
411                clear_at: None,
412            },
413            Message {
414                role: MessageRole::Assistant,
415                content: crate::llm::provider::MessageContent::Text("b".repeat(400)),
416                reasoning: None,
417                reasoning_details: None,
418                tool_calls: None,
419                tool_call_id: None,
420                phase: None,
421                origin_tool: None,
422                metadata: None,
423                clear_at: None,
424            },
425            assistant.clone(),
426            Message {
427                role: MessageRole::Tool,
428                content: crate::llm::provider::MessageContent::Text("ok".to_string()),
429                reasoning: None,
430                reasoning_details: None,
431                tool_calls: None,
432                tool_call_id: Some("tc1".to_string()),
433                phase: None,
434                origin_tool: None,
435                metadata: None,
436                clear_at: None,
437            },
438            Message {
439                role: MessageRole::User,
440                content: crate::llm::provider::MessageContent::Text("tail".to_string()),
441                reasoning: None,
442                reasoning_details: None,
443                tool_calls: None,
444                tool_call_id: None,
445                phase: None,
446                origin_tool: None,
447                metadata: None,
448                clear_at: None,
449            },
450        ];
451        let split = split_conversation_for_two_pass(&conv, 0.9);
452        if let Some(msg) = split.prefix.last() {
453            if msg.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()) {
454                assert!(!matches!(split.tail.first(), Some(m) if m.tool_call_id.is_some()));
455            }
456        }
457        assert!(!matches!(split.tail.first(), Some(m) if m.tool_call_id.is_some()));
458        let prefix_has_call = split.prefix.iter().any(|m| {
459            m.tool_calls
460                .as_ref()
461                .map(|tc| tc.iter().any(|t| t.id == "tc1"))
462                .unwrap_or(false)
463        });
464        let prefix_has_result = split.prefix.iter().any(|m| m.tool_call_id.as_deref() == Some("tc1"));
465        let tail_has_call = split.tail.iter().any(|m| {
466            m.tool_calls
467                .as_ref()
468                .map(|tc| tc.iter().any(|t| t.id == "tc1"))
469                .unwrap_or(false)
470        });
471        let tail_has_result = split.tail.iter().any(|m| m.tool_call_id.as_deref() == Some("tc1"));
472        assert_eq!(prefix_has_call, prefix_has_result);
473        assert_eq!(tail_has_call, tail_has_result);
474    }
475
476    #[test]
477    fn note_prefers_summary_block_and_caps_huge_raw() {
478        let text = format!("<summary>short</summary>\n<summary>\n{}\n</summary>", "x".repeat(1001));
479        assert_eq!(note_for_two_pass_pass2(&text), "x".repeat(1001));
480        assert_eq!(note_for_two_pass_pass2("no tags"), "no tags");
481        let huge = "n".repeat(TWO_PASS_MAX_NOTE1_CHARS + 500);
482        let note = note_for_two_pass_pass2(&huge);
483        assert!(note.chars().count() <= TWO_PASS_MAX_NOTE1_CHARS + 80);
484        assert!(note.contains("truncated"));
485    }
486
487    #[test]
488    fn pass_histories_shape() {
489        let conv = vec![
490            Message {
491                role: MessageRole::System,
492                content: crate::llm::provider::MessageContent::Text("You are VT Code.".to_string()),
493                reasoning: None,
494                reasoning_details: None,
495                tool_calls: None,
496                tool_call_id: None,
497                phase: None,
498                origin_tool: None,
499                metadata: None,
500                clear_at: None,
501            },
502            Message {
503                role: MessageRole::User,
504                content: crate::llm::provider::MessageContent::Text("early".to_string()),
505                reasoning: None,
506                reasoning_details: None,
507                tool_calls: None,
508                tool_call_id: None,
509                phase: None,
510                origin_tool: None,
511                metadata: None,
512                clear_at: None,
513            },
514            Message {
515                role: MessageRole::Assistant,
516                content: crate::llm::provider::MessageContent::Text("early-a".to_string()),
517                reasoning: None,
518                reasoning_details: None,
519                tool_calls: None,
520                tool_call_id: None,
521                phase: None,
522                origin_tool: None,
523                metadata: None,
524                clear_at: None,
525            },
526            Message {
527                role: MessageRole::User,
528                content: crate::llm::provider::MessageContent::Text("late".to_string()),
529                reasoning: None,
530                reasoning_details: None,
531                tool_calls: None,
532                tool_call_id: None,
533                phase: None,
534                origin_tool: None,
535                metadata: None,
536                clear_at: None,
537            },
538            Message {
539                role: MessageRole::Assistant,
540                content: crate::llm::provider::MessageContent::Text("late-a".to_string()),
541                reasoning: None,
542                reasoning_details: None,
543                tool_calls: None,
544                tool_call_id: None,
545                phase: None,
546                origin_tool: None,
547                metadata: None,
548                clear_at: None,
549            },
550        ];
551        let split = split_conversation_for_two_pass(&conv, 0.5);
552        let prompt = "1. Primary Request and Intent: x\n5. Optional Next Step: y\n";
553        let pass1 = build_two_pass_pass1_history(split.prefix, prompt);
554        assert!(matches!(pass1.last(), Some(m) if m.role == MessageRole::User));
555
556        let note1 = "x".repeat(1001);
557        let pass2 = build_two_pass_pass2_history(split.prefix, split.tail, &note1, prompt);
558        assert!(pass2.iter().any(|m| m.role == MessageRole::System));
559        let texts: Vec<String> = pass2
560            .iter()
561            .map(|m| match &m.content {
562                crate::llm::provider::MessageContent::Text(t) => t.clone(),
563                crate::llm::provider::MessageContent::Parts(parts) => parts
564                    .iter()
565                    .filter_map(|p| match p {
566                        crate::llm::provider::ContentPart::Text { text } => Some(text.clone()),
567                        _ => None,
568                    })
569                    .collect(),
570            })
571            .collect();
572        assert!(texts.iter().any(|t| t.contains("<summary_content>")));
573        assert!(texts.last().is_some_and(|t| t.contains("special compaction case")));
574        // NOTE₁ is embedded once (carrier turn); the final turn refers back to it.
575        let note1_copies: usize = texts.iter().map(|t| t.matches(note1.as_str()).count()).sum();
576        assert_eq!(note1_copies, 1);
577        assert!(texts.last().is_some_and(|t| !t.contains(note1.as_str())));
578    }
579
580    #[test]
581    fn pass2_fallback_system_prompt_states_compaction_role() {
582        let conv = vec![
583            Message {
584                role: MessageRole::User,
585                content: crate::llm::provider::MessageContent::Text("early".to_string()),
586                reasoning: None,
587                reasoning_details: None,
588                tool_calls: None,
589                tool_call_id: None,
590                phase: None,
591                origin_tool: None,
592                metadata: None,
593                clear_at: None,
594            },
595            Message {
596                role: MessageRole::User,
597                content: crate::llm::provider::MessageContent::Text("late".to_string()),
598                reasoning: None,
599                reasoning_details: None,
600                tool_calls: None,
601                tool_call_id: None,
602                phase: None,
603                origin_tool: None,
604                metadata: None,
605                clear_at: None,
606            },
607        ];
608        let split = split_conversation_for_two_pass(&conv, 0.5);
609        let pass2 = build_two_pass_pass2_history(split.prefix, split.tail, "note", "summarize");
610        let system = pass2.first().expect("system message");
611        assert_eq!(system.role, MessageRole::System);
612        assert!(system.content.as_text().contains("condensing a coding-agent session"));
613    }
614}